wangzixiang
/
lerobot
mirror of https://github.com/huggingface/lerobot.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

WIP

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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
examples/pretrained.ipynb
View File

@ -13,12 +13,21 @@
"from examples.notebook_utils import config_notebook\n", "from examples.notebook_utils import config_notebook\n",
"from lerobot.scripts.eval import eval\n", "from lerobot.scripts.eval import eval\n",
"\n", "\n",
"OUT_DIR = Path(\"./outputs\")\n", "# Select policy and env\n",
"POLICY = \"act\"\n", "POLICY = \"act\" # \"tdmpc\" | \"diffusion\"\n",
"ENV = \"aloha\"\n", "ENV = \"aloha\" # \"pusht\" | \"simxarm\"\n",
"\n", "\n",
"# setup config\n", "# Select device\n",
"cfg = config_notebook(policy=POLICY, env=ENV, device=\"cpu\", print_config=True)" "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": [], "outputs": [],
"source": [ "source": [
"eval(cfg, out_dir=OUT_DIR)\n", "eval(cfg, out_dir=OUT_DIR)\n",
"Video(OUT_DIR / \"eval\" / \"eval_episode_0.mp4\", embed=True)" "Video(OUT_EXAMPLE, embed=True)"
] ]
} }
], ],

17
lerobot/scripts/eval.py
View File

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