finish examples 2 and 3

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")