From f2b1842d6929fa6ca981a3a0d3a3ad3136011aa3 Mon Sep 17 00:00:00 2001 From: Remi Cadene Date: Thu, 26 Sep 2024 11:48:22 +0200 Subject: [PATCH] fix unit test --- lerobot/scripts/control_robot.py | 6 +++--- tests/test_control_robot.py | 1 + 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/lerobot/scripts/control_robot.py b/lerobot/scripts/control_robot.py index 995b19c9..c07602a1 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 857d809b..2b35e798 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,