WIP

examples/pretrained.py

@@ -0,0 +1,67 @@
+import logging
+from pathlib import Path
+
+import torch
+from huggingface_hub import snapshot_download
+from omegaconf import OmegaConf
+from tensordict.nn import TensorDictModule
+
+from lerobot.common.datasets.factory import make_offline_buffer
+from lerobot.common.envs.factory import make_env
+from lerobot.common.logger import log_output_dir
+from lerobot.common.policies.factory import make_policy
+from lerobot.common.utils import get_safe_torch_device, init_logging, set_seed
+from lerobot.scripts.eval import eval_policy
+
+folder = Path(snapshot_download("lerobot/diffusion_policy_pusht_image", revision="v1.0"))
+cfg = OmegaConf.load(folder / "config.yaml")
+cfg.policy.pretrained_model_path = folder / "model.pt"
+cfg.eval_episodes = 1
+cfg.episode_length = 50
+# cfg.device = "cpu"
+
+out_dir = "tmp/"
+
+if out_dir is None:
+    raise NotImplementedError()
+
+init_logging()
+
+# Check device is available
+get_safe_torch_device(cfg.device, log=True)
+
+torch.backends.cudnn.benchmark = True
+torch.backends.cuda.matmul.allow_tf32 = True
+set_seed(cfg.seed)
+
+log_output_dir(out_dir)
+
+logging.info("make_offline_buffer")
+offline_buffer = make_offline_buffer(cfg)
+
+logging.info("make_env")
+env = make_env(cfg, transform=offline_buffer.transform)
+
+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,
+    policy=policy,
+    save_video=True,
+    video_dir=Path(out_dir) / "eval",
+    fps=cfg.env.fps,
+    max_steps=cfg.env.episode_length,
+    num_episodes=cfg.eval_episodes,
+)
+print(metrics)
+
+logging.info("End of eval")