Eval reproduced! Train running (but not reproduced)

lerobot/scripts/visualize_dataset.py

@@ -0,0 +1,59 @@
+import pickle
+from pathlib import Path
+
+import imageio
+import simxarm
+import torch
+from torchrl.data.replay_buffers import (
+    SamplerWithoutReplacement,
+    SliceSampler,
+    SliceSamplerWithoutReplacement,
+)
+
+from lerobot.common.datasets.simxarm import SimxarmExperienceReplay
+
+
+def visualize_simxarm_dataset(dataset_id="xarm_lift_medium"):
+    sampler = SliceSamplerWithoutReplacement(
+        num_slices=1,
+        strict_length=False,
+        shuffle=False,
+    )
+
+    dataset = SimxarmExperienceReplay(
+        dataset_id,
+        # download="force",
+        download=True,
+        streaming=False,
+        root="data",
+        sampler=sampler,
+    )
+
+    NUM_EPISODES_TO_RENDER = 10
+    MAX_NUM_STEPS = 50
+    FIRST_FRAME = 0
+    for _ in range(NUM_EPISODES_TO_RENDER):
+        episode = dataset.sample(MAX_NUM_STEPS)
+
+        ep_idx = episode["episode"][FIRST_FRAME].item()
+        ep_frames = torch.cat(
+            [
+                episode["observation"]["image"][FIRST_FRAME][None, ...],
+                episode["next", "observation"]["image"],
+            ],
+            dim=0,
+        )
+
+        video_dir = Path("tmp/2024_02_03_xarm_lift_medium")
+        video_dir.mkdir(parents=True, exist_ok=True)
+        # TODO(rcadene): make fps configurable
+        video_path = video_dir / f"eval_episode_{ep_idx}.mp4"
+        imageio.mimsave(video_path, ep_frames.numpy().transpose(0, 2, 3, 1), fps=15)
+
+        # ran out of episodes
+        if dataset._sampler._sample_list.numel() == 0:
+            break
+
+
+if __name__ == "__main__":
+    visualize_simxarm_dataset()