forked from tangger/lerobot
102 lines
3.0 KiB
Plaintext
102 lines
3.0 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"%load_ext autoreload\n",
|
|
"%autoreload 2"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from pathlib import Path\n",
|
|
"from pprint import pprint\n",
|
|
"\n",
|
|
"from huggingface_hub import snapshot_download\n",
|
|
"from hydra import compose, initialize\n",
|
|
"from hydra.core.global_hydra import GlobalHydra\n",
|
|
"from IPython.display import Video\n",
|
|
"from omegaconf import OmegaConf\n",
|
|
"from omegaconf.dictconfig import DictConfig\n",
|
|
"\n",
|
|
"from examples.notebook_utils import config_notebook\n",
|
|
"from examples.pretrained_script import download_eval_pretrained\n",
|
|
"from lerobot.scripts.eval import eval\n",
|
|
"\n",
|
|
"# Select policy and env\n",
|
|
"POLICY = \"diffusion\" # \"tdmpc\" | \"diffusion\"\n",
|
|
"ENV = \"pusht\" # \"pusht\" | \"simxarm\"\n",
|
|
"\n",
|
|
"# Select device\n",
|
|
"DEVICE = \"mps\" # \"cuda\" | \"mps\"\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",
|
|
"PRETRAINED_REPO = \"lerobot/diffusion_policy_pusht_image\"\n",
|
|
"pretrained_folder = Path(snapshot_download(repo_id=PRETRAINED_REPO, repo_type=\"model\", revision=\"v1.0\"))\n",
|
|
"pretrained_model_path = pretrained_folder / \"model.pt\"\n",
|
|
"\n",
|
|
"cfg_path = pretrained_folder / \"config.yaml\"\n",
|
|
"GlobalHydra.instance().clear()\n",
|
|
"\n",
|
|
"print(pretrained_folder)\n",
|
|
"\n",
|
|
"initialize(config_path=\"../../../.cache/huggingface/hub/models--lerobot--diffusion_policy_pusht_image/snapshots/163d168f5c193c356b82e3bf6bbf5b4eeaa780d7\")\n",
|
|
"overrides = [\n",
|
|
" f\"env={ENV}\",\n",
|
|
" f\"policy={POLICY}\",\n",
|
|
" f\"device={DEVICE}\",\n",
|
|
" f\"+policy.pretrained_model_path={pretrained_model_path}\",\n",
|
|
" f\"eval_episodes=1\",\n",
|
|
" f\"+env.episode_length=200\",\n",
|
|
"]\n",
|
|
"cfg = compose(config_name=\"config\", overrides=overrides)\n",
|
|
"pprint(OmegaConf.to_container(cfg))\n",
|
|
"# Setup config\n",
|
|
"#cfg = config_notebook(cfg_path, policy=POLICY, env=ENV, device=DEVICE, print_config=False, pretrained_model_path=pretrained_model_path)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# eval(cfg, out_dir=OUT_DIR)\n",
|
|
"download_eval_pretrained(OUT_DIR, cfg)\n",
|
|
"Video(OUT_EXAMPLE, embed=True)"
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"kernelspec": {
|
|
"display_name": "lerobot",
|
|
"language": "python",
|
|
"name": "python3"
|
|
},
|
|
"language_info": {
|
|
"codemirror_mode": {
|
|
"name": "ipython",
|
|
"version": 3
|
|
},
|
|
"file_extension": ".py",
|
|
"mimetype": "text/x-python",
|
|
"name": "python",
|
|
"nbconvert_exporter": "python",
|
|
"pygments_lexer": "ipython3",
|
|
"version": "3.10.14"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 2
|
|
}
|