WIP

examples/pretrained.ipynb

@@ -13,12 +13,21 @@
     "from examples.notebook_utils import config_notebook\n",
     "from lerobot.scripts.eval import eval\n",
     "\n",
-    "OUT_DIR = Path(\"./outputs\")\n",
-    "POLICY = \"act\"\n",
-    "ENV = \"aloha\"\n",
+    "# Select policy and env\n",
+    "POLICY = \"act\"  # \"tdmpc\" | \"diffusion\"\n",
+    "ENV = \"aloha\"  # \"pusht\" | \"simxarm\"\n",
     "\n",
-    "# setup config\n",
-    "cfg = config_notebook(policy=POLICY, env=ENV, device=\"cpu\", print_config=True)"
+    "# Select device\n",
+    "DEVICE = \"mps\"  # \"cpu\" | \"cuda\"\n",
+    "\n",
+    "# Generated videos will be written here\n",
+    "OUT_DIR = Path(\"./outputs\")\n",
+    "OUT_EXAMPLE = OUT_DIR / \"eval\" / \"eval_episode_0.mp4\"\n",
+    "\n",
+    "# Setup config\n",
+    "cfg = config_notebook(policy=POLICY, env=ENV, device=DEVICE, print_config=False)\n",
+    "print(cfg.env.episode_length)\n",
+    "print(cfg.n_action_steps)"
    ]
   },
   {
@@ -28,7 +37,7 @@
    "outputs": [],
    "source": [
     "eval(cfg, out_dir=OUT_DIR)\n",
-    "Video(OUT_DIR / \"eval\" / \"eval_episode_0.mp4\", embed=True)"
+    "Video(OUT_EXAMPLE, embed=True)"
    ]
   }
  ],

lerobot/scripts/eval.py

@@ -145,17 +145,24 @@ def eval(cfg: dict, out_dir=None):
     logging.info("make_env")
     env = make_env(cfg, transform=offline_buffer.transform)
 
-    # TODO(aliberts, Cadene): fetch pretrained model from HF hub
-    if cfg.policy.pretrained_model_path:
+    # WIP
     policy = make_policy(cfg)
     policy = TensorDictModule(
         policy,
         in_keys=["observation", "step_count"],
         out_keys=["action"],
     )
-    else:
-        # when policy is None, rollout a random policy
-        policy = None
+    # TODO(aliberts, Cadene): fetch pretrained model from HF hub
+    # if cfg.policy.pretrained_model_path:
+    #     policy = make_policy(cfg)
+    #     policy = TensorDictModule(
+    #         policy,
+    #         in_keys=["observation", "step_count"],
+    #         out_keys=["action"],
+    #     )
+    # else:
+    #     # when policy is None, rollout a random policy
+    #     policy = None
 
     metrics = eval_policy(
         env,