Improve readme format

.github/poetry/cpu/poetry.lock

@@ -327,35 +327,6 @@ files = [
     {file = "cloudpickle-3.0.0.tar.gz", hash = "sha256:996d9a482c6fb4f33c1a35335cf8afd065d2a56e973270364840712d9131a882"},
 ]
 
-[[package]]
-name = "cmake"
-version = "3.29.0.1"
-description = "CMake is an open-source, cross-platform family of tools designed to build, test and package software"
-optional = false
-python-versions = ">=3.7"
-files = [
-    {file = "cmake-3.29.0.1-py3-none-macosx_10_10_universal2.macosx_10_10_x86_64.macosx_11_0_arm64.macosx_11_0_universal2.whl", hash = "sha256:ec8b39fdeb75c48fd5a2894658a1ca75f94fb49b421c1f753d86d3e5d5e9f196"},
-    {file = "cmake-3.29.0.1-py3-none-manylinux_2_12_i686.manylinux2010_i686.whl", hash = "sha256:ead7dc5176a6c6347b3fc19532c25ec328f9279b6213902ac930242334e7b621"},
-    {file = "cmake-3.29.0.1-py3-none-manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:f7c7dabf3dd40cb830d2eded43d51c5a3737625bbc5ab6916041c04e352b74f8"},
-    {file = "cmake-3.29.0.1-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:de8f55198f4a820daf2c57645a4bb8cd1064dc92d950ad95be14c5ffabc15bd4"},
-    {file = "cmake-3.29.0.1-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:7bb8aaa3419eafd466931e4dcc161d3e5e6a82730ab508c75946ff4fc883b3f6"},
-    {file = "cmake-3.29.0.1-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:e956e4b6c2d8d6bbee399bf3c77e5b901d916fe8f35d6b2f58444d5892c4602b"},
-    {file = "cmake-3.29.0.1-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:2a8f7c8b07e6ab0dd444c5b74e658d5013ca0da456041029f734d751090bb7ec"},
-    {file = "cmake-3.29.0.1-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:22493049b6383ea2baa7237a326c2914ab4a7b3e1642f4233245e3a34aae39f6"},
-    {file = "cmake-3.29.0.1-py3-none-musllinux_1_1_aarch64.whl", hash = "sha256:d09573411901fc8a3ede7433713c000ac7b81cda0d771874e9182770acf29eb4"},
-    {file = "cmake-3.29.0.1-py3-none-musllinux_1_1_i686.whl", hash = "sha256:c85e35ec572e54152154637f24d6bde316fbcf94dfea644bd8f22b1855a09abe"},
-    {file = "cmake-3.29.0.1-py3-none-musllinux_1_1_ppc64le.whl", hash = "sha256:a36f3b19a5caa6c63aa70bfa0b262bae8d296e68c6e6d8e918edc5c51d952bf8"},
-    {file = "cmake-3.29.0.1-py3-none-musllinux_1_1_s390x.whl", hash = "sha256:85a0e28eaaee311d50fcee60f730e5a44a65b3cefc556a1163bbabd7328acd60"},
-    {file = "cmake-3.29.0.1-py3-none-musllinux_1_1_x86_64.whl", hash = "sha256:58879ef15dd8344e1583a36cead794fb0fee13f78c590a56749283ac2e27d30a"},
-    {file = "cmake-3.29.0.1-py3-none-win32.whl", hash = "sha256:068a3e7461dd9e487f5f3f720ae8072c2eeb37239ebe4642c4ae29058d83347f"},
-    {file = "cmake-3.29.0.1-py3-none-win_amd64.whl", hash = "sha256:c097892b3653e2d2d41d055c80a9af029fdd24f9eff2ecb66576e4da8b85b1c7"},
-    {file = "cmake-3.29.0.1-py3-none-win_arm64.whl", hash = "sha256:1808071047cb49ed0fe2359e4c310b49880c2805cad4ea9f03c959f51b881ac7"},
-    {file = "cmake-3.29.0.1.tar.gz", hash = "sha256:ec49a7a4480959c229d9d2aa77f7885859c17a45fc66981aaf4551ceffb4d030"},
-]
-
-[package.extras]
-test = ["coverage (>=4.2)", "pytest (>=3.0.3)", "pytest-cov (>=2.4.0)"]
-
 [[package]]
 name = "colorama"
 version = "0.4.6"
@@ -3330,4 +3301,4 @@ testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "p
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.10"
-content-hash = "8800bb8b24312d17b765cd2ce2799f49436171dd5fbf1bec3b07f853cfa9befd"
+content-hash = "9f3f9a0c7ec2aed90c4f51737b71a07e0819c7cb895fe4e84fa4ada1f352b063"

.github/poetry/cpu/pyproject.toml

@@ -52,7 +52,6 @@ robomimic = "0.2.0"
 huggingface-hub = "^0.21.4"
 gymnasium-robotics = "^1.2.4"
 gymnasium = "^0.29.1"
-cmake = "^3.29.0.1"
 
 
 [tool.poetry.group.dev.dependencies]

.github/workflows/test.yml

@@ -154,16 +154,15 @@ jobs:
             device=cpu \
             policy.pretrained_model_path=tests/outputs/act/models/2.pt
 
-      # TODO(aliberts): This takes ~2mn to run, needs to be improved
-      # - name: Test eval ACT on ALOHA end-to-end (policy is None)
-      #   run: |
-      #     source .venv/bin/activate
-      #     python lerobot/scripts/eval.py \
-      #       --config lerobot/configs/default.yaml \
-      #       policy=act \
-      #       env=aloha \
-      #       eval_episodes=1 \
-      #       device=cpu
+      - name: Test eval ACT on ALOHA end-to-end (policy is None)
+        run: |
+          source .venv/bin/activate
+          python lerobot/scripts/eval.py \
+            --config lerobot/configs/default.yaml \
+            policy=act \
+            env=aloha \
+            eval_episodes=1 \
+            device=cpu
 
       - name: Test train Diffusion on PushT end-to-end
         run: |

README.md

@@ -21,12 +21,9 @@
 
 </div>
 
-<h3 align="center">
-    <p>State-of-the-art Machine Learning for real-world robotics</p>
-</h3>
-
 ---
 
+**State-of-the-art machine learning for real-world robotics**
 
 🤗 LeRobot aims to provide models, datasets, and tools for real-world robotics in PyTorch. The goal is to lower the barrier for entry to robotics so that everyone can contribute and benefit from sharing datasets and pretrained models.
 
@@ -60,21 +57,26 @@
 
 ## Installation
 
-Download our source code:
-```bash
-git clone https://github.com/huggingface/lerobot.git
-cd lerobot
-```
-
-Create a virtual environment with Python 3.10 and activate it, e.g. with [`miniconda`](https://docs.anaconda.com/free/miniconda/index.html):
+Create a virtual environment with Python 3.10, e.g. using `conda`:
 ```bash
 conda create -y -n lerobot python=3.10
 conda activate lerobot
 ```
 
-Then, install 🤗 LeRobot:
+[Install `poetry`](https://python-poetry.org/docs/#installation) (if you don't have it already)
 ```bash
-python -m pip install .
+curl -sSL https://install.python-poetry.org | python -
+```
+
+Install dependencies
+```bash
+poetry install
+```
+
+If you encounter a disk space error, try to change your tmp dir to a location where you have enough disk space, e.g.
+```bash
+mkdir ~/tmp
+export TMPDIR='~/tmp'
 ```
 
 To use [Weights and Biases](https://docs.wandb.ai/quickstart) for experiments tracking, log in with
@@ -146,17 +148,22 @@ hydra.run.dir=outputs/visualize_dataset/example
 
 ### Evaluate a pretrained policy
 
-Check out [example 2](./examples/2_evaluate_pretrained_policy.py) to see how you can load a pretrained policy from HuggingFace hub, load up the corresponding environment and model, and run an evaluation.
+You can import our environment class, download pretrained policies from the HuggingFace hub, and use our rollout utilities with rendering:
+```python
+""" Copy pasted from `examples/2_evaluate_pretrained_policy.py`
+# TODO
+```
 
 Or you can achieve the same result by executing our script from the command line:
 ```bash
 python lerobot/scripts/eval.py \
 --hub-id lerobot/diffusion_policy_pusht_image \
+--revision v1.0 \
 eval_episodes=10 \
 hydra.run.dir=outputs/eval/example_hub
 ```
 
-After training your own policy, you can also re-evaluate the checkpoints with:
+After launching training of your own policy, you can also re-evaluate the checkpoints with:
 ```bash
 python lerobot/scripts/eval.py \
 --config PATH/TO/FOLDER/config.yaml \
@@ -169,9 +176,19 @@ See `python lerobot/scripts/eval.py --help` for more instructions.
 
 ### Train your own policy
 
-You can import our dataset, environment, policy classes, and use our training utilities (if some data is missing, it will be automatically downloaded from HuggingFace hub): check out [example 3](./examples/3_train_policy.py). After you run this, you may want to revisit [example 2](./examples/2_evaluate_pretrained_policy.py) to evaluate your training output!
+You can import our dataset, environment, policy classes, and use our training utilities (if some data is missing, it will be automatically downloaded from HuggingFace hub):
+```python
+""" Copy pasted from `examples/3_train_policy.py`
+# TODO
+```
 
-In general, you can use our training script to easily train any policy on any environment:
+Or you can achieve the same result by executing our script from the command line:
+```bash
+python lerobot/scripts/train.py \
+hydra.run.dir=outputs/train/example
+```
+
+You can easily train any policy on any environment:
 ```bash
 python lerobot/scripts/train.py \
 env=aloha \
@@ -185,11 +202,11 @@ hydra.run.dir=outputs/train/aloha_act
 
 Feel free to open issues and PRs, and to coordinate your efforts with the community on our [Discord Channel](https://discord.gg/VjFz58wn3R). For specific inquiries, reach out to [Remi Cadene](remi.cadene@huggingface.co).
 
-### TODO
+**TODO**
 
 If you are not sure how to contribute or want to know the next features we working on, look on this project page: [LeRobot TODO](https://github.com/orgs/huggingface/projects/46)
 
-### Follow our style
+**Follow our style**
 
 ```bash
 # install if needed
@@ -198,33 +215,25 @@ pre-commit install
 pre-commit
 ```
 
-### Add dependencies
+**Add dependencies**
 
-Instead of using `pip` directly, we use `poetry` for development purposes to easily track our dependencies.
-If you don't have it already, follow the [instructions](https://python-poetry.org/docs/#installation) to install it.
-
-Install the project with:
-```bash
-poetry install
-```
-
-Then, the equivalent of `pip install some-package`, would just be:
+Instead of `pip install some-package`, we use `poetry` to track the versions of our dependencies:
 ```bash
 poetry add some-package
 ```
 
-**NOTE:** Currently, to ensure the CI works properly, any new package must also be added in the CPU-only environment dedicated to the CI. To do this, you should create a separate environment and add the new package there as well. For example:
+**NOTE:** Currently, to ensure the CI works properly, any new package must also be added in the CPU-only environment dedicated CI. To do this, you should create a separate environment and add the new package there as well. For example:
 ```bash
-# Add the new package to your main poetry env
+# add the new package to your main poetry env
 poetry add some-package
-# Add the same package to the CPU-only env dedicated to CI
+# add the same package to the CPU-only env dedicated to CI
 conda create -y -n lerobot-ci python=3.10
 conda activate lerobot-ci
 cd .github/poetry/cpu
 poetry add some-package
 ```
 
-### Run tests locally
+**Run tests locally**
 
 Install [git lfs](https://git-lfs.com/) to retrieve test artifacts (if you don't have it already).
 
@@ -255,7 +264,7 @@ Run tests
 DATA_DIR="tests/data" pytest -sx tests
 ```
 
-### Add a new dataset
+**Add a new dataset**
 
 To add a dataset to the hub, first login and use a token generated from [huggingface settings](https://huggingface.co/settings/tokens) with write access:
 ```bash
@@ -316,7 +325,7 @@ Finally, you might want to mock the dataset if you need to update the unit tests
 python tests/scripts/mock_dataset.py --in-data-dir data/$DATASET --out-data-dir tests/data/$DATASET
 ```
 
-### Add a pretrained policy
+**Add a pretrained policy**
 
 Once you have trained a policy you may upload it to the HuggingFace hub.
 
@@ -352,7 +361,7 @@ huggingface-cli upload $HUB_ID to_upload
 See `eval.py` for an example of how a user may use your policy.
 
 
-### Improve your code with profiling
+**Improve your code with profiling**
 
 An example of a code snippet to profile the evaluation of a policy:
 ```python

examples/2_evaluate_pretrained_policy.py

@@ -1,39 +1 @@
-"""
-This scripts demonstrates how to evaluate a pretrained policy from the HuggingFace Hub or from your local
-training outputs directory. In the latter case, you might want to run examples/3_train_policy.py first.
-"""
-
-from pathlib import Path
-
-from huggingface_hub import snapshot_download
-
-from lerobot.common.utils import init_hydra_config
-from lerobot.scripts.eval import eval
-
-# Get a pretrained policy from the hub.
-hub_id = "lerobot/diffusion_policy_pusht_image"
-folder = Path(snapshot_download(hub_id))
-# OR uncomment the following to evaluate a policy from the local outputs/train folder.
-# folder = Path("outputs/train/example_pusht_diffusion")
-
-config_path = folder / "config.yaml"
-weights_path = folder / "model.pt"
-stats_path = folder / "stats.pth"  # normalization stats
-
-# Override some config parameters to do with evaluation.
-overrides = [
-    f"policy.pretrained_model_path={weights_path}",
-    "eval_episodes=10",
-    "rollout_batch_size=10",
-    "device=cuda",
-]
-
-# Create a Hydra config.
-cfg = init_hydra_config(config_path, overrides)
-
-# Evaluate the policy and save the outputs including metrics and videos.
-eval(
-    cfg,
-    out_dir=f"outputs/eval/example_{cfg.env.name}_{cfg.policy.name}",
-    stats_path=stats_path,
-)
+# TODO

examples/3_train_policy.py

@@ -1,55 +1 @@
-"""This scripts demonstrates how to train Diffusion Policy on the PushT environment.
-
-Once you have trained a model with this script, you can try to evaluate it on
-examples/2_evaluate_pretrained_policy.py
-"""
-
-import os
-from pathlib import Path
-
-import torch
-from omegaconf import OmegaConf
-from tqdm import trange
-
-from lerobot.common.datasets.factory import make_offline_buffer
-from lerobot.common.policies.diffusion.policy import DiffusionPolicy
-from lerobot.common.utils import init_hydra_config
-
-output_directory = Path("outputs/train/example_pusht_diffusion")
-os.makedirs(output_directory, exist_ok=True)
-
-overrides = [
-    "env=pusht",
-    "policy=diffusion",
-    # Adjust as you prefer. 5000 steps are needed to get something worth evaluating.
-    "offline_steps=5000",
-    "log_freq=250",
-    "device=cuda",
-]
-
-cfg = init_hydra_config("lerobot/configs/default.yaml", overrides)
-
-policy = DiffusionPolicy(
-    cfg=cfg.policy,
-    cfg_device=cfg.device,
-    cfg_noise_scheduler=cfg.noise_scheduler,
-    cfg_rgb_model=cfg.rgb_model,
-    cfg_obs_encoder=cfg.obs_encoder,
-    cfg_optimizer=cfg.optimizer,
-    cfg_ema=cfg.ema,
-    n_action_steps=cfg.n_action_steps + cfg.n_latency_steps,
-    **cfg.policy,
-)
-policy.train()
-
-offline_buffer = make_offline_buffer(cfg)
-
-for offline_step in trange(cfg.offline_steps):
-    train_info = policy.update(offline_buffer, offline_step)
-    if offline_step % cfg.log_freq == 0:
-        print(train_info)
-
-# Save the policy, configuration, and normalization stats for later use.
-policy.save_pretrained(output_directory / "model.pt")
-OmegaConf.save(cfg, output_directory / "config.yaml")
-torch.save(offline_buffer.transform[-1].stats, output_directory / "stats.pth")
+# TODO

lerobot/common/datasets/factory.py

@@ -81,8 +81,13 @@ def make_offline_buffer(
     else:
         raise ValueError(cfg.env.name)
 
+    # TODO(rcadene): backward compatiblity to load pretrained pusht policy
+    dataset_id = cfg.get("dataset_id")
+    if dataset_id is None and cfg.env.name == "pusht":
+        dataset_id = "pusht"
+
     offline_buffer = clsfunc(
-        dataset_id=cfg.dataset_id,
+        dataset_id=dataset_id,
         sampler=sampler,
         batch_size=batch_size,
         root=DATA_DIR,

lerobot/common/logger.py

@@ -5,7 +5,6 @@ from pathlib import Path
 from omegaconf import OmegaConf
 from termcolor import colored
 
-from lerobot.common.policies.abstract import AbstractPolicy
 
 def log_output_dir(out_dir):
     logging.info(colored("Output dir:", "yellow", attrs=["bold"]) + f" {out_dir}")
@@ -68,11 +67,11 @@ class Logger:
             logging.info(f"Track this run --> {colored(wandb.run.get_url(), 'yellow', attrs=['bold'])}")
             self._wandb = wandb
 
-    def save_model(self, policy: AbstractPolicy, identifier):
+    def save_model(self, policy, identifier):
         if self._save_model:
             self._model_dir.mkdir(parents=True, exist_ok=True)
             fp = self._model_dir / f"{str(identifier)}.pt"
-            policy.save_pretrained(fp)
+            policy.save(fp)
             if self._wandb and not self._disable_wandb_artifact:
                 # note wandb artifact does not accept ":" in its name
                 artifact = self._wandb.Artifact(

lerobot/common/policies/abstract.py

@@ -2,25 +2,14 @@ from collections import deque
 
 import torch
 from torch import Tensor, nn
-from huggingface_hub import PyTorchModelHubMixin
 
 
-class AbstractPolicy(nn.Module, PyTorchModelHubMixin):
+class AbstractPolicy(nn.Module):
     """Base policy which all policies should be derived from.
 
     The forward method should generally not be overriden as it plays the role of handling multi-step policies. See its
     documentation for more information.
 
-    The policy is a PyTorchModelHubMixin, which means that it can be saved and loaded from the Hugging Face Hub and/or to a local directory.
-    # Save policy weights to local directory
-    >>> policy.save_pretrained("my-awesome-policy")
-
-    # Push policy weights to the Hub
-    >>> policy.push_to_hub("my-awesome-policy")
-
-    # Download and initialize policy from the Hub
-    >>> policy = MyPolicy.from_pretrained("username/my-awesome-policy")
-
     Note:
         When implementing a concrete class (e.g. `AlohaDataset`, `PushtEnv`, `DiffusionPolicy`), you need to:
             1. set the required class attributes:
@@ -33,7 +22,7 @@ class AbstractPolicy(nn.Module, PyTorchModelHubMixin):
 
     name: str | None = None  # same name should be used to instantiate the policy in factory.py
 
-    def __init__(self, n_action_steps: int | None = None):
+    def __init__(self, n_action_steps: int | None):
         """
         n_action_steps: Sets the cache size for storing action trajectories. If None, it is assumed that a single
             action is returned by `select_actions` and that doesn't have a horizon dimension. The `forward` method then
@@ -48,10 +37,10 @@ class AbstractPolicy(nn.Module, PyTorchModelHubMixin):
         """One step of the policy's learning algorithm."""
         raise NotImplementedError("Abstract method")
 
-    def save(self, fp):  # TODO: remove this method since we are using PyTorchModelHubMixin
+    def save(self, fp):
         torch.save(self.state_dict(), fp)
 
-    def load(self, fp):  # TODO: remove this method since we are using PyTorchModelHubMixin
+    def load(self, fp):
         d = torch.load(fp)
         self.load_state_dict(d)
 

lerobot/common/policies/act/policy.py

@@ -136,8 +136,8 @@ class ActionChunkingTransformerPolicy(AbstractPolicy):
     def save(self, fp):
         torch.save(self.state_dict(), fp)
 
-    def load(self, fp, device=None):
-        d = torch.load(fp, map_location=device)
+    def load(self, fp):
+        d = torch.load(fp)
         self.load_state_dict(d)
 
     def compute_loss(self, batch):

lerobot/common/policies/diffusion/diffusion_unet_image_policy.py

@@ -32,7 +32,7 @@ assert len(unexpected_keys) == 0
 Then in that same runtime you can also save the weights with the new aligned state_dict:
 
 ```
-policy.save_pretrained("my-policy")
+policy.save("weights.pt")
 ```
 
 Now you can remove the breakpoint and extra code and load in the weights just like with any other lerobot checkpoint.

lerobot/common/policies/diffusion/policy.py

@@ -203,8 +203,8 @@ class DiffusionPolicy(AbstractPolicy):
     def save(self, fp):
         torch.save(self.state_dict(), fp)
 
-    def load(self, fp, device=None):
-        d = torch.load(fp, map_location=device)
+    def load(self, fp):
+        d = torch.load(fp)
         missing_keys, unexpected_keys = self.load_state_dict(d, strict=False)
         if len(missing_keys) > 0:
             assert all(k.startswith("ema_diffusion.") for k in missing_keys)

lerobot/common/policies/factory.py

@@ -1,53 +1,35 @@
-""" Factory for policies
-"""
-
-from lerobot.common.policies.abstract import AbstractPolicy
-
-
-def make_policy(cfg: dict) -> AbstractPolicy:
-    """ Instantiate a policy from the configuration.
-        Currently supports TD-MPC, Diffusion, and ACT: select the policy with cfg.policy.name: tdmpc, diffusion, act.
-    
-    Args:
-        cfg: The configuration (DictConfig)
-
-    """
-    policy_kwargs = {}
+def make_policy(cfg):
     if cfg.policy.name != "diffusion" and cfg.rollout_batch_size > 1:
         raise NotImplementedError("Only diffusion policy supports rollout_batch_size > 1 for the time being.")
 
     if cfg.policy.name == "tdmpc":
         from lerobot.common.policies.tdmpc.policy import TDMPCPolicy
 
-        policy_cls = TDMPCPolicy
-        policy_kwargs = {"cfg": cfg.policy, "device": cfg.device}
+        policy = TDMPCPolicy(cfg.policy, cfg.device)
     elif cfg.policy.name == "diffusion":
         from lerobot.common.policies.diffusion.policy import DiffusionPolicy
 
-        policy_cls = DiffusionPolicy
-        policy_kwargs = {
-            "cfg": cfg.policy,
-            "cfg_device": cfg.device,
-            "cfg_noise_scheduler": cfg.noise_scheduler,
-            "cfg_rgb_model": cfg.rgb_model,
-            "cfg_obs_encoder": cfg.obs_encoder,
-            "cfg_optimizer": cfg.optimizer,
-            "cfg_ema": cfg.ema,
-            "n_action_steps": cfg.n_action_steps + cfg.n_latency_steps,
+        policy = DiffusionPolicy(
+            cfg=cfg.policy,
+            cfg_device=cfg.device,
+            cfg_noise_scheduler=cfg.noise_scheduler,
+            cfg_rgb_model=cfg.rgb_model,
+            cfg_obs_encoder=cfg.obs_encoder,
+            cfg_optimizer=cfg.optimizer,
+            cfg_ema=cfg.ema,
+            n_action_steps=cfg.n_action_steps + cfg.n_latency_steps,
             **cfg.policy,
-        }
+        )
     elif cfg.policy.name == "act":
         from lerobot.common.policies.act.policy import ActionChunkingTransformerPolicy
 
-        policy_cls = ActionChunkingTransformerPolicy
-        policy_kwargs = {"cfg": cfg.policy, "device": cfg.device, "n_action_steps": cfg.n_action_steps + cfg.n_latency_steps}
+        policy = ActionChunkingTransformerPolicy(
+            cfg.policy, cfg.device, n_action_steps=cfg.n_action_steps + cfg.n_latency_steps
+        )
     else:
         raise ValueError(cfg.policy.name)
 
     if cfg.policy.pretrained_model_path:
-        # policy.load(cfg.policy.pretrained_model_path, device=cfg.device)
-        policy = policy_cls.from_pretrained(cfg.policy.pretrained_model_path, map_location=cfg.device, **policy_kwargs)
-    
         # TODO(rcadene): hack for old pretrained models from fowm
         if cfg.policy.name == "tdmpc" and "fowm" in cfg.policy.pretrained_model_path:
             if "offline" in cfg.pretrained_model_path:
@@ -56,5 +38,6 @@ def make_policy(cfg: dict) -> AbstractPolicy:
                 policy.step[0] = 100000
             else:
                 raise NotImplementedError()
+        policy.load(cfg.policy.pretrained_model_path)
 
     return policy

lerobot/common/policies/tdmpc/policy.py

@@ -122,9 +122,9 @@ class TDMPCPolicy(AbstractPolicy):
         """Save state dict of TOLD model to filepath."""
         torch.save(self.state_dict(), fp)
 
-    def load(self, fp, device=None):
+    def load(self, fp):
         """Load a saved state dict from filepath into current agent."""
-        d = torch.load(fp, map_location=device)
+        d = torch.load(fp)
         self.model.load_state_dict(d["model"])
         self.model_target.load_state_dict(d["model_target"])
 

lerobot/common/utils.py

@@ -1,13 +1,9 @@
 import logging
-import os.path as osp
 import random
 from datetime import datetime
-from pathlib import Path
 
-import hydra
 import numpy as np
 import torch
-from omegaconf import DictConfig
 
 
 def get_safe_torch_device(cfg_device: str, log: bool = False) -> torch.device:
@@ -67,31 +63,3 @@ def format_big_number(num):
         num /= divisor
 
     return num
-
-
-def _relative_path_between(path1: Path, path2: Path) -> Path:
-    """Returns path1 relative to path2."""
-    path1 = path1.absolute()
-    path2 = path2.absolute()
-    try:
-        return path1.relative_to(path2)
-    except ValueError:  # most likely because path1 is not a subpath of path2
-        common_parts = Path(osp.commonpath([path1, path2])).parts
-        return Path(
-            "/".join([".."] * (len(path2.parts) - len(common_parts)) + list(path1.parts[len(common_parts) :]))
-        )
-
-
-def init_hydra_config(config_path: str, overrides: list[str] | None = None) -> DictConfig:
-    """Initialize a Hydra config given only the path to the relevant config file.
-
-    For config resolution, it is assumed that the config file's parent is the Hydra config dir.
-    """
-    # TODO(alexander-soare): Resolve configs without Hydra initialization.
-    hydra.core.global_hydra.GlobalHydra.instance().clear()
-    # Hydra needs a path relative to this file.
-    hydra.initialize(
-        str(_relative_path_between(Path(config_path).absolute().parent, Path(__file__).absolute().parent))
-    )
-    cfg = hydra.compose(Path(config_path).stem, overrides)
-    return cfg

lerobot/configs/policy/diffusion.yaml

@@ -103,3 +103,29 @@ optimizer:
   betas: [0.95, 0.999]
   eps: 1.0e-8
   weight_decay: 1.0e-6
+
+training:
+  device: "cuda:0"
+  seed: 42
+  debug: False
+  resume: True
+  # optimization
+  # lr_scheduler: cosine
+  # lr_warmup_steps: 500
+  num_epochs: 8000
+  # gradient_accumulate_every: 1
+  # EMA destroys performance when used with BatchNorm
+  # replace BatchNorm with GroupNorm.
+  # use_ema: True
+  freeze_encoder: False
+  # training loop control
+  # in epochs
+  rollout_every: 50
+  checkpoint_every: 50
+  val_every: 1
+  sample_every: 5
+  # steps per epoch
+  max_train_steps: null
+  max_val_steps: null
+  # misc
+  tqdm_interval_sec: 1.0

lerobot/scripts/eval.py

@@ -30,13 +30,14 @@ python lerobot/scripts/eval.py --hub-id HUB/ID --revision v1.0 eval_episodes=10
 import argparse
 import json
 import logging
+import os.path as osp
 import threading
 import time
-from typing import Tuple, Union
 from datetime import datetime as dt
 from pathlib import Path
 
 import einops
+import hydra
 import imageio
 import numpy as np
 import torch
@@ -51,7 +52,7 @@ from lerobot.common.envs.factory import make_env
 from lerobot.common.logger import log_output_dir
 from lerobot.common.policies.abstract import AbstractPolicy
 from lerobot.common.policies.factory import make_policy
-from lerobot.common.utils import get_safe_torch_device, init_hydra_config, init_logging, set_global_seed
+from lerobot.common.utils import get_safe_torch_device, init_logging, set_global_seed
 
 
 def write_video(video_path, stacked_frames, fps):
@@ -67,19 +68,7 @@ def eval_policy(
     video_dir: Path = None,
     fps: int = 15,
     return_first_video: bool = False,
-) -> Union[dict, Tuple[dict, torch.Tensor]]:
-    """ Evaluate a policy on an environment by running rollouts and computing metrics.
-
-    Args:
-        env: The environment to evaluate.
-        policy: The policy to evaluate.
-        num_episodes: The number of episodes to evaluate.
-        max_steps: The maximum number of steps per episode.
-        save_video: Whether to save videos of the evaluation episodes.
-        video_dir: The directory to save the videos.
-        fps: The frames per second for the videos.
-        return_first_video: Whether to return the first video as a tensor.
-    """
+):
     if policy is not None:
         policy.eval()
     start = time.time()
@@ -158,7 +147,7 @@ def eval_policy(
     for thread in threads:
         thread.join()
 
-    info = {  # TODO: change to dataclass
+    info = {
         "per_episode": [
             {
                 "episode_ix": i,
@@ -191,13 +180,6 @@ def eval_policy(
 
 
 def eval(cfg: dict, out_dir=None, stats_path=None):
-    """ Evaluate a policy.
-
-    Args:
-        cfg: The configuration (DictConfig).
-        out_dir: The directory to save the evaluation results (JSON file and videos)
-        stats_path: The path to the stats file.
-    """
     if out_dir is None:
         raise NotImplementedError()
 
@@ -213,7 +195,6 @@ def eval(cfg: dict, out_dir=None, stats_path=None):
     log_output_dir(out_dir)
 
     logging.info("Making transforms.")
-    # TODO(alexander-soare): Completely decouple datasets from evaluation.
     offline_buffer = make_offline_buffer(cfg, stats_path=stats_path)
 
     logging.info("Making environment.")
@@ -248,6 +229,19 @@ def eval(cfg: dict, out_dir=None, stats_path=None):
     logging.info("End of eval")
 
 
+def _relative_path_between(path1: Path, path2: Path) -> Path:
+    """Returns path1 relative to path2."""
+    path1 = path1.absolute()
+    path2 = path2.absolute()
+    try:
+        return path1.relative_to(path2)
+    except ValueError:  # most likely because path1 is not a subpath of path2
+        common_parts = Path(osp.commonpath([path1, path2])).parts
+        return Path(
+            "/".join([".."] * (len(path2.parts) - len(common_parts)) + list(path1.parts[len(common_parts) :]))
+        )
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(
         description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
@@ -265,15 +259,19 @@ if __name__ == "__main__":
 
     if args.config is not None:
         # Note: For the config_path, Hydra wants a path relative to this script file.
-        cfg = init_hydra_config(args.config, args.overrides)
+        hydra.initialize(
+            config_path=str(
+                _relative_path_between(Path(args.config).absolute().parent, Path(__file__).parent)
+            )
+        )
+        cfg = hydra.compose(Path(args.config).stem, args.overrides)
         # TODO(alexander-soare): Save and load stats in trained model directory.
         stats_path = None
     elif args.hub_id is not None:
-        folder = Path(snapshot_download(args.hub_id, revision=args.revision))
-        cfg = init_hydra_config(
-            folder / "config.yaml", [*args.overrides]
-            # folder / "config.yaml"  # , [f"policy.pretrained_model_path={folder / 'model.pt'}", *args.overrides]
-        )
+        folder = Path(snapshot_download(args.hub_id, revision="v1.0"))
+        cfg = hydra.initialize(config_path=str(_relative_path_between(folder, Path(__file__).parent)))
+        cfg = hydra.compose("config", args.overrides)
+        cfg.policy.pretrained_model_path = folder / "model.pt"
         stats_path = folder / "stats.pth"
 
     eval(

poetry.lock

@@ -327,35 +327,6 @@ files = [
     {file = "cloudpickle-3.0.0.tar.gz", hash = "sha256:996d9a482c6fb4f33c1a35335cf8afd065d2a56e973270364840712d9131a882"},
 ]
 
-[[package]]
-name = "cmake"
-version = "3.29.0.1"
-description = "CMake is an open-source, cross-platform family of tools designed to build, test and package software"
-optional = false
-python-versions = ">=3.7"
-files = [
-    {file = "cmake-3.29.0.1-py3-none-macosx_10_10_universal2.macosx_10_10_x86_64.macosx_11_0_arm64.macosx_11_0_universal2.whl", hash = "sha256:ec8b39fdeb75c48fd5a2894658a1ca75f94fb49b421c1f753d86d3e5d5e9f196"},
-    {file = "cmake-3.29.0.1-py3-none-manylinux_2_12_i686.manylinux2010_i686.whl", hash = "sha256:ead7dc5176a6c6347b3fc19532c25ec328f9279b6213902ac930242334e7b621"},
-    {file = "cmake-3.29.0.1-py3-none-manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:f7c7dabf3dd40cb830d2eded43d51c5a3737625bbc5ab6916041c04e352b74f8"},
-    {file = "cmake-3.29.0.1-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:de8f55198f4a820daf2c57645a4bb8cd1064dc92d950ad95be14c5ffabc15bd4"},
-    {file = "cmake-3.29.0.1-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:7bb8aaa3419eafd466931e4dcc161d3e5e6a82730ab508c75946ff4fc883b3f6"},
-    {file = "cmake-3.29.0.1-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:e956e4b6c2d8d6bbee399bf3c77e5b901d916fe8f35d6b2f58444d5892c4602b"},
-    {file = "cmake-3.29.0.1-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:2a8f7c8b07e6ab0dd444c5b74e658d5013ca0da456041029f734d751090bb7ec"},
-    {file = "cmake-3.29.0.1-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:22493049b6383ea2baa7237a326c2914ab4a7b3e1642f4233245e3a34aae39f6"},
-    {file = "cmake-3.29.0.1-py3-none-musllinux_1_1_aarch64.whl", hash = "sha256:d09573411901fc8a3ede7433713c000ac7b81cda0d771874e9182770acf29eb4"},
-    {file = "cmake-3.29.0.1-py3-none-musllinux_1_1_i686.whl", hash = "sha256:c85e35ec572e54152154637f24d6bde316fbcf94dfea644bd8f22b1855a09abe"},
-    {file = "cmake-3.29.0.1-py3-none-musllinux_1_1_ppc64le.whl", hash = "sha256:a36f3b19a5caa6c63aa70bfa0b262bae8d296e68c6e6d8e918edc5c51d952bf8"},
-    {file = "cmake-3.29.0.1-py3-none-musllinux_1_1_s390x.whl", hash = "sha256:85a0e28eaaee311d50fcee60f730e5a44a65b3cefc556a1163bbabd7328acd60"},
-    {file = "cmake-3.29.0.1-py3-none-musllinux_1_1_x86_64.whl", hash = "sha256:58879ef15dd8344e1583a36cead794fb0fee13f78c590a56749283ac2e27d30a"},
-    {file = "cmake-3.29.0.1-py3-none-win32.whl", hash = "sha256:068a3e7461dd9e487f5f3f720ae8072c2eeb37239ebe4642c4ae29058d83347f"},
-    {file = "cmake-3.29.0.1-py3-none-win_amd64.whl", hash = "sha256:c097892b3653e2d2d41d055c80a9af029fdd24f9eff2ecb66576e4da8b85b1c7"},
-    {file = "cmake-3.29.0.1-py3-none-win_arm64.whl", hash = "sha256:1808071047cb49ed0fe2359e4c310b49880c2805cad4ea9f03c959f51b881ac7"},
-    {file = "cmake-3.29.0.1.tar.gz", hash = "sha256:ec49a7a4480959c229d9d2aa77f7885859c17a45fc66981aaf4551ceffb4d030"},
-]
-
-[package.extras]
-test = ["coverage (>=4.2)", "pytest (>=3.0.3)", "pytest-cov (>=2.4.0)"]
-
 [[package]]
 name = "colorama"
 version = "0.4.6"
@@ -3589,4 +3560,4 @@ testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "p
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.10"
-content-hash = "174c7d42f8039eedd2c447a4e6cae5169782cbd94346b5606572a0010194ca05"
+content-hash = "e5c1e40c1cab71b9f2620e7b160c2f03f99e3e18804268a1667ad69699c34bb3"

pyproject.toml

@@ -51,7 +51,6 @@ huggingface-hub = {extras = ["hf-transfer"], version = "^0.21.4"}
 robomimic = "0.2.0"
 gymnasium-robotics = "^1.2.4"
 gymnasium = "^0.29.1"
-cmake = "^3.29.0.1"
 
 
 [tool.poetry.group.dev.dependencies]

tests/test_datasets.py

@@ -2,9 +2,8 @@ import pytest
 import torch
 
 from lerobot.common.datasets.factory import make_offline_buffer
-from lerobot.common.utils import init_hydra_config
 
-from .utils import DEVICE, DEFAULT_CONFIG_PATH
+from .utils import DEVICE, init_config
 
 
 @pytest.mark.parametrize(
@@ -19,10 +18,7 @@ from .utils import DEVICE, DEFAULT_CONFIG_PATH
     ],
 )
 def test_factory(env_name, dataset_id):
-    cfg = init_hydra_config(
-        DEFAULT_CONFIG_PATH,
-        overrides=[f"env={env_name}", f"env.task={dataset_id}", f"device={DEVICE}"]
-    )
+    cfg = init_config(overrides=[f"env={env_name}", f"env.task={dataset_id}", f"device={DEVICE}"])
     offline_buffer = make_offline_buffer(cfg)
     for key in offline_buffer.image_keys:
         img = offline_buffer[0].get(key)

tests/test_envs.py

@@ -4,13 +4,12 @@ import torch
 from torchrl.envs.utils import check_env_specs, step_mdp
 from lerobot.common.datasets.factory import make_offline_buffer
 
-from lerobot.common.envs.aloha.env import AlohaEnv
 from lerobot.common.envs.factory import make_env
 from lerobot.common.envs.pusht.env import PushtEnv
 from lerobot.common.envs.simxarm.env import SimxarmEnv
-from lerobot.common.utils import init_hydra_config
+from lerobot.common.envs.aloha.env import AlohaEnv
 
-from .utils import DEVICE, DEFAULT_CONFIG_PATH
+from .utils import DEVICE, init_config
 
 
 def print_spec_rollout(env):
@@ -111,10 +110,7 @@ def test_pusht(from_pixels, pixels_only):
     ],
 )
 def test_factory(env_name):
-    cfg = init_hydra_config(
-        DEFAULT_CONFIG_PATH,
-        overrides=[f"env={env_name}", f"device={DEVICE}"],
-    )
+    cfg = init_config(overrides=[f"env={env_name}", f"device={DEVICE}"])
 
     offline_buffer = make_offline_buffer(cfg)
 

tests/test_examples.py

@@ -1,70 +1,19 @@
+import pytest
 from pathlib import Path
 
+@pytest.mark.parametrize(
+    "path",
+    [
+        "examples/1_visualize_dataset.py",
+        "examples/2_evaluate_pretrained_policy.py",
+        "examples/3_train_policy.py",
+    ],
+)
+def test_example(path):
 
-def _find_and_replace(text: str, finds: list[str], replaces: list[str]) -> str:
-    for f, r in zip(finds, replaces):
-        assert f in text
-        text = text.replace(f, r)
-    return text
-
-
-def test_example_1():
-    path = "examples/1_visualize_dataset.py"
-
-    with open(path, "r") as file:
+    with open(path, 'r') as file:
         file_contents = file.read()
     exec(file_contents)
 
-    assert Path("outputs/visualize_dataset/example/episode_0.mp4").exists()
-
-
-def test_examples_3_and_2():
-    """
-    Train a model with example 3, check the outputs.
-    Evaluate the trained model with example 2, check the outputs.
-    """
-
-    path = "examples/3_train_policy.py"
-
-    with open(path, "r") as file:
-        file_contents = file.read()
-
-    # Do less steps and use CPU.
-    file_contents = _find_and_replace(
-        file_contents,
-        ['"offline_steps=5000"', '"device=cuda"'],
-        ['"offline_steps=1"', '"device=cpu"'],
-    )
-
-    exec(file_contents)
-
-    for file_name in ["model.pt", "stats.pth", "config.yaml"]:
-        assert Path(f"outputs/train/example_pusht_diffusion/{file_name}").exists()
-
-    path = "examples/2_evaluate_pretrained_policy.py"
-
-    with open(path, "r") as file:
-        file_contents = file.read()
-
-    # Do less evals, use CPU, and use the local model.
-    file_contents = _find_and_replace(
-        file_contents,
-        [
-            '"eval_episodes=10"',
-            '"rollout_batch_size=10"',
-            '"device=cuda"',
-            '# folder = Path("outputs/train/example_pusht_diffusion")',
-            'hub_id = "lerobot/diffusion_policy_pusht_image"',
-            "folder = Path(snapshot_download(hub_id)",
-        ],
-        [
-            '"eval_episodes=1"',
-            '"rollout_batch_size=1"',
-            '"device=cpu"',
-            'folder = Path("outputs/train/example_pusht_diffusion")',
-            "",
-            "",
-        ],
-    )
-
-    assert Path(f"outputs/train/example_pusht_diffusion").exists()
+    if path == "examples/1_visualize_dataset.py":
+        assert Path("outputs/visualize_dataset/example/episode_0.mp4").exists()

tests/test_policies.py

@@ -1,3 +1,4 @@
+from omegaconf import open_dict
 import pytest
 from tensordict import TensorDict
 from tensordict.nn import TensorDictModule
@@ -9,8 +10,8 @@ from lerobot.common.policies.factory import make_policy
 from lerobot.common.envs.factory import make_env
 from lerobot.common.datasets.factory import make_offline_buffer
 from lerobot.common.policies.abstract import AbstractPolicy
-from lerobot.common.utils import init_hydra_config
-from .utils import DEVICE, DEFAULT_CONFIG_PATH
+
+from .utils import DEVICE, init_config
 
 @pytest.mark.parametrize(
     "env_name,policy_name,extra_overrides",
@@ -33,8 +34,7 @@ def test_concrete_policy(env_name, policy_name, extra_overrides):
         - Updating the policy.
         - Using the policy to select actions at inference time.
     """
-    cfg = init_hydra_config(
-        DEFAULT_CONFIG_PATH,
+    cfg = init_config(
         overrides=[
             f"env={env_name}",
             f"policy={policy_name}",

tests/utils.py

@@ -1,6 +1,13 @@
 import os
+import hydra
+from hydra import compose, initialize
 
-# Pass this as the first argument to init_hydra_config.
-DEFAULT_CONFIG_PATH = "lerobot/configs/default.yaml"
+CONFIG_PATH = "../lerobot/configs"
 
 DEVICE = os.environ.get('LEROBOT_TESTS_DEVICE', "cuda")
+
+def init_config(config_name="default", overrides=None):
+    hydra.core.global_hydra.GlobalHydra.instance().clear()
+    initialize(config_path=CONFIG_PATH)
+    cfg = compose(config_name=config_name, overrides=overrides)
+    return cfg