hangX/lerobot
1
0
Fork 0
forked from tangger/lerobot
Code Pull Requests Activity

fix unit test

Browse Source
This commit is contained in:
Remi Cadene
2024-09-26 11:48:22 +02:00
parent 500d505bf6
commit f2b1842d69
2 changed files with 4 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
lerobot/scripts/control_robot.py
View File

@@ -561,9 +561,9 @@ def record(
else: else:
imgs_dir = videos_dir / f"{key}_episode_{episode_index:06d}" imgs_dir = videos_dir / f"{key}_episode_{episode_index:06d}"
ep_dict[key] = [] ep_dict[key] = []
for frame_index in range(num_frames): for i in range(num_frames):
img_path = imgs_dir / f"frame_{frame_index:06d}.png" img_path = imgs_dir / f"frame_{i:06d}.png"
ep_dict[key].append(img_path) ep_dict[key].append({"path": str(img_path)})
for key in not_image_keys: for key in not_image_keys:
ep_dict[key] = torch.stack(ep_dict[key]) ep_dict[key] = torch.stack(ep_dict[key])

1
tests/test_control_robot.py
View File

@@ -114,6 +114,7 @@ def test_record_and_replay_and_policy(tmpdir, request, robot_type, mock):
cfg, cfg,
warmup_time_s=1, warmup_time_s=1,
episode_time_s=1, episode_time_s=1,
num_episodes=2,
run_compute_stats=False, run_compute_stats=False,
push_to_hub=False, push_to_hub=False,
video=False, video=False,