Merge pull request #53 from alexander-soare/finish_examples

Add examples 2 and 3

README.md

@@ -146,11 +146,7 @@ hydra.run.dir=outputs/visualize_dataset/example
 
 ### Evaluate a pretrained policy
 
-You can import our environment class, download pretrained policies from the HuggingFace hub, and use our rollout utilities with rendering:
+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.
-```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
@@ -160,7 +156,7 @@ eval_episodes=10 \
 hydra.run.dir=outputs/eval/example_hub
 ```
 
-After launching training of your own policy, you can also re-evaluate the checkpoints with:
+After training your own policy, you can also re-evaluate the checkpoints with:
 ```bash
 python lerobot/scripts/eval.py \
 --config PATH/TO/FOLDER/config.yaml \
@@ -173,19 +169,9 @@ 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):
+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!
-```python
-""" Copy pasted from `examples/3_train_policy.py`
-# TODO
-```
 
-Or you can achieve the same result by executing our script from the command line:
+In general, you can use our training script to easily train any policy on any environment:
-```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 \

examples/2_evaluate_pretrained_policy.py

@@ -1 +1,39 @@
-# TODO
+"""
+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,
+)

examples/3_train_policy.py

@@ -1 +1,55 @@
-# TODO
+"""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(output_directory / "model.pt")
+OmegaConf.save(cfg, output_directory / "config.yaml")
+torch.save(offline_buffer.transform[-1].stats, output_directory / "stats.pth")

lerobot/common/utils.py

@@ -1,9 +1,13 @@
 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:
@@ -63,3 +67,31 @@ 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/scripts/eval.py

@@ -30,14 +30,12 @@ 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 datetime import datetime as dt
 from pathlib import Path
 
 import einops
-import hydra
 import imageio
 import numpy as np
 import torch
@@ -52,7 +50,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_logging, set_global_seed
+from lerobot.common.utils import get_safe_torch_device, init_hydra_config, init_logging, set_global_seed
 
 
 def write_video(video_path, stacked_frames, fps):
@@ -195,6 +193,7 @@ 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.")
@@ -229,19 +228,6 @@ 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
@@ -259,19 +245,14 @@ if __name__ == "__main__":
 
     if args.config is not None:
         # Note: For the config_path, Hydra wants a path relative to this script file.
-        hydra.initialize(
+        cfg = init_hydra_config(args.config, args.overrides)
-            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 = hydra.initialize(config_path=str(_relative_path_between(folder, Path(__file__).parent)))
+        cfg = init_hydra_config(
-        cfg = hydra.compose("config", args.overrides)
+            folder / "config.yaml", [f"policy.pretrained_model_path={folder / 'model.pt'}", *args.overrides]
-        cfg.policy.pretrained_model_path = folder / "model.pt"
+        )
         stats_path = folder / "stats.pth"
 
     eval(

tests/test_datasets.py

@@ -2,8 +2,9 @@ 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, init_config
+from .utils import DEVICE, DEFAULT_CONFIG_PATH
 
 
 @pytest.mark.parametrize(
@@ -18,7 +19,10 @@ from .utils import DEVICE, init_config
     ],
 )
 def test_factory(env_name, dataset_id):
-    cfg = init_config(overrides=[f"env={env_name}", f"env.task={dataset_id}", f"device={DEVICE}"])
+    cfg = init_hydra_config(
+        DEFAULT_CONFIG_PATH,
+        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,12 +4,13 @@ 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.envs.aloha.env import AlohaEnv
+from lerobot.common.utils import init_hydra_config
 
-from .utils import DEVICE, init_config
+from .utils import DEVICE, DEFAULT_CONFIG_PATH
 
 
 def print_spec_rollout(env):
@@ -110,7 +111,10 @@ def test_pusht(from_pixels, pixels_only):
     ],
 )
 def test_factory(env_name):
-    cfg = init_config(overrides=[f"env={env_name}", f"device={DEVICE}"])
+    cfg = init_hydra_config(
+        DEFAULT_CONFIG_PATH,
+        overrides=[f"env={env_name}", f"device={DEVICE}"],
+    )
 
     offline_buffer = make_offline_buffer(cfg)
 

tests/test_examples.py

@@ -1,19 +1,70 @@
-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):
 
-    with open(path, 'r') as file:
+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:
         file_contents = file.read()
     exec(file_contents)
 
-    if path == "examples/1_visualize_dataset.py":
+    assert Path("outputs/visualize_dataset/example/episode_0.mp4").exists()
-        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()

tests/test_policies.py

@@ -1,4 +1,3 @@
-from omegaconf import open_dict
 import pytest
 from tensordict import TensorDict
 from tensordict.nn import TensorDictModule
@@ -10,8 +9,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, init_config
+from .utils import DEVICE, DEFAULT_CONFIG_PATH
 
 @pytest.mark.parametrize(
     "env_name,policy_name,extra_overrides",
@@ -34,7 +33,8 @@ def test_concrete_policy(env_name, policy_name, extra_overrides):
         - Updating the policy.
         - Using the policy to select actions at inference time.
     """
-    cfg = init_config(
+    cfg = init_hydra_config(
+        DEFAULT_CONFIG_PATH,
         overrides=[
             f"env={env_name}",
             f"policy={policy_name}",

tests/utils.py

@@ -1,13 +1,6 @@
 import os
-import hydra
-from hydra import compose, initialize
 
-CONFIG_PATH = "../lerobot/configs"
+# Pass this as the first argument to init_hydra_config.
+DEFAULT_CONFIG_PATH = "lerobot/configs/default.yaml"
 
 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