diff --git a/lerobot/scripts/control_robot.py b/lerobot/scripts/control_robot.py index 995b19c9e..c07602a19 100644 --- a/lerobot/scripts/control_robot.py +++ b/lerobot/scripts/control_robot.py @@ -561,9 +561,9 @@ def record( else: imgs_dir = videos_dir / f"{key}_episode_{episode_index:06d}" ep_dict[key] = [] - for frame_index in range(num_frames): - img_path = imgs_dir / f"frame_{frame_index:06d}.png" - ep_dict[key].append(img_path) + for i in range(num_frames): + img_path = imgs_dir / f"frame_{i:06d}.png" + ep_dict[key].append({"path": str(img_path)}) for key in not_image_keys: ep_dict[key] = torch.stack(ep_dict[key]) diff --git a/tests/test_control_robot.py b/tests/test_control_robot.py index 857d809bb..2b35e7984 100644 --- a/tests/test_control_robot.py +++ b/tests/test_control_robot.py @@ -114,6 +114,7 @@ def test_record_and_replay_and_policy(tmpdir, request, robot_type, mock): cfg, warmup_time_s=1, episode_time_s=1, + num_episodes=2, run_compute_stats=False, push_to_hub=False, video=False,