From b10960140d27da691b5ed887d7feefb95bb856f7 Mon Sep 17 00:00:00 2001
From: Simon Alibert <alibert.sim@gmail.com>
Date: Wed, 20 Mar 2024 15:52:34 +0100
Subject: [PATCH] WIP

---
 examples/pretrained.ipynb | 21 +++++++++++++++------
 lerobot/scripts/eval.py   | 27 +++++++++++++++++----------
 2 files changed, 32 insertions(+), 16 deletions(-)

diff --git a/examples/pretrained.ipynb b/examples/pretrained.ipynb
index 990f1d0a..64f98946 100644
--- a/examples/pretrained.ipynb
+++ b/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)"
    ]
   }
  ],
diff --git a/lerobot/scripts/eval.py b/lerobot/scripts/eval.py
index 1ec34480..699d6d1f 100644
--- a/lerobot/scripts/eval.py
+++ b/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)
 
+    # WIP
+    policy = make_policy(cfg)
+    policy = TensorDictModule(
+        policy,
+        in_keys=["observation", "step_count"],
+        out_keys=["action"],
+    )
     # 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
+    # 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,