This commit is contained in:
Simon Alibert
2024-03-20 15:52:34 +01:00
parent eb56a96e67
commit b10960140d
2 changed files with 32 additions and 16 deletions

View File

@@ -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)"
]
}
],