diff --git a/lerobot/common/robot_devices/robots/factory.py b/lerobot/common/robot_devices/robots/factory.py index 94466115..593603e5 100644 --- a/lerobot/common/robot_devices/robots/factory.py +++ b/lerobot/common/robot_devices/robots/factory.py @@ -34,8 +34,8 @@ def make_robot(name): ), }, cameras={ - "macbookpro": OpenCVCamera(1, fps=30, width=640, height=480), - "iphone": OpenCVCamera(2, fps=30, width=640, height=480), + "laptop": OpenCVCamera(1, fps=30, width=640, height=480), + "phone": OpenCVCamera(2, fps=30, width=640, height=480), }, ) else: diff --git a/lerobot/scripts/control_robot.py b/lerobot/scripts/control_robot.py index 396c7834..cccd6ef0 100644 --- a/lerobot/scripts/control_robot.py +++ b/lerobot/scripts/control_robot.py @@ -209,33 +209,35 @@ def record_dataset( # Save images using threads to reach high fps (30 and more) # Using `with` to exist smoothly if an execption is raised. # Using only 4 worker threads to avoid blocking the main thread. - with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: - # Execute a few seconds without recording data, to give times - # to the robot devices to connect and start synchronizing. - timestamp = 0 - start_time = time.perf_counter() - is_warmup_print = False - while timestamp < warmup_time_s: - if not is_warmup_print: - logging.info("Warming up (no data recording)") - os.system('say "Warmup" &') - is_warmup_print = True + futures = [] - now = time.perf_counter() - observation, action = robot.teleop_step(record_data=True) + # Execute a few seconds without recording data, to give times + # to the robot devices to connect and start synchronizing. + timestamp = 0 + start_time = time.perf_counter() + is_warmup_print = False + while timestamp < warmup_time_s: + if not is_warmup_print: + logging.info("Warming up (no data recording)") + os.system('say "Warmup" &') + is_warmup_print = True - dt_s = time.perf_counter() - now - busy_wait(1 / fps - dt_s) + now = time.perf_counter() + observation, action = robot.teleop_step(record_data=True) - dt_s = time.perf_counter() - now - log_control_info(robot, dt_s) + dt_s = time.perf_counter() - now + busy_wait(1 / fps - dt_s) - timestamp = time.perf_counter() - start_time + dt_s = time.perf_counter() - now + log_control_info(robot, dt_s) - # Start recording all episodes - ep_dicts = [] - for episode_index in range(num_episodes): + timestamp = time.perf_counter() - start_time + + # Start recording all episodes + ep_dicts = [] + for episode_index in range(num_episodes): + with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: ep_dict = {} frame_index = 0 timestamp = 0 @@ -254,7 +256,8 @@ def record_dataset( not_image_keys = [key for key in observation if "image" not in key] for key in image_keys: - executor.submit(save_image, observation[key], key, frame_index, episode_index, videos_dir) + future = executor.submit(save_image, observation[key], key, frame_index, episode_index, videos_dir) + futures.append(future) for key in not_image_keys: if key not in ep_dict: @@ -276,40 +279,40 @@ def record_dataset( timestamp = time.perf_counter() - start_time - logging.info("Encoding images to videos") + logging.info("Encoding images to videos") - num_frames = frame_index + num_frames = frame_index - for key in image_keys: - tmp_imgs_dir = videos_dir / f"{key}_episode_{episode_index:06d}" - fname = f"{key}_episode_{episode_index:06d}.mp4" - video_path = local_dir / "videos" / fname - encode_video_frames(tmp_imgs_dir, video_path, fps) + for key in image_keys: + tmp_imgs_dir = videos_dir / f"{key}_episode_{episode_index:06d}" + fname = f"{key}_episode_{episode_index:06d}.mp4" + video_path = local_dir / "videos" / fname + encode_video_frames(tmp_imgs_dir, video_path, fps) - # TODO(rcadene): uncomment? - # clean temporary images directory - # shutil.rmtree(tmp_imgs_dir) + # TODO(rcadene): uncomment? + # clean temporary images directory + # shutil.rmtree(tmp_imgs_dir) - # store the reference to the video frame - ep_dict[key] = [] - for i in range(num_frames): - ep_dict[key].append({"path": f"videos/{fname}", "timestamp": i / fps}) + # store the reference to the video frame + ep_dict[key] = [] + for i in range(num_frames): + ep_dict[key].append({"path": f"videos/{fname}", "timestamp": i / fps}) - for key in not_image_keys: - ep_dict[key] = torch.stack(ep_dict[key]) + for key in not_image_keys: + ep_dict[key] = torch.stack(ep_dict[key]) - for key in action: - ep_dict[key] = torch.stack(ep_dict[key]) + for key in action: + ep_dict[key] = torch.stack(ep_dict[key]) - ep_dict["episode_index"] = torch.tensor([episode_index] * num_frames) - ep_dict["frame_index"] = torch.arange(0, num_frames, 1) - ep_dict["timestamp"] = torch.arange(0, num_frames, 1) / fps + ep_dict["episode_index"] = torch.tensor([episode_index] * num_frames) + ep_dict["frame_index"] = torch.arange(0, num_frames, 1) + ep_dict["timestamp"] = torch.arange(0, num_frames, 1) / fps - done = torch.zeros(num_frames, dtype=torch.bool) - done[-1] = True - ep_dict["next.done"] = done + done = torch.zeros(num_frames, dtype=torch.bool) + done[-1] = True + ep_dict["next.done"] = done - ep_dicts.append(ep_dict) + ep_dicts.append(ep_dict) data_dict = concatenate_episodes(ep_dicts) @@ -338,6 +341,7 @@ def record_dataset( videos_dir=videos_dir, ) stats = compute_stats(lerobot_dataset) if run_compute_stats else {} + lerobot_dataset.stats = stats hf_dataset = hf_dataset.with_format(None) # to remove transforms that cant be saved hf_dataset.save_to_disk(str(local_dir / "train")) diff --git a/poetry.lock b/poetry.lock index f39c13db..cc917d70 100644 --- a/poetry.lock +++ b/poetry.lock @@ -2375,9 +2375,8 @@ description = "Nvidia JIT LTO Library" optional = false python-versions = ">=3" files = [ - {file = "nvidia_nvjitlink_cu12-12.5.40-py3-none-manylinux2014_aarch64.whl", hash = "sha256:004186d5ea6a57758fd6d57052a123c73a4815adf365eb8dd6a85c9eaa7535ff"}, - {file = "nvidia_nvjitlink_cu12-12.5.40-py3-none-manylinux2014_x86_64.whl", hash = "sha256:d9714f27c1d0f0895cd8915c07a87a1d0029a0aa36acaf9156952ec2a8a12189"}, - {file = "nvidia_nvjitlink_cu12-12.5.40-py3-none-win_amd64.whl", hash = "sha256:c3401dc8543b52d3a8158007a0c1ab4e9c768fcbd24153a48c86972102197ddd"}, + {file = "nvidia_nvjitlink_cu12-12.5.82-py3-none-manylinux2014_x86_64.whl", hash = "sha256:f9b37bc5c8cf7509665cb6ada5aaa0ce65618f2332b7d3e78e9790511f111212"}, + {file = "nvidia_nvjitlink_cu12-12.5.82-py3-none-win_amd64.whl", hash = "sha256:e782564d705ff0bf61ac3e1bf730166da66dd2fe9012f111ede5fc49b64ae697"}, ] [[package]] @@ -4364,4 +4363,4 @@ xarm = ["gym-xarm"] [metadata] lock-version = "2.0" python-versions = ">=3.10,<3.13" -content-hash = "81dc830d3d36c67e2fe2aea6cc30829eb2977edbf49a037df21a5f329a01aee5" +content-hash = "223a6496a630da8181f119634f96bed3e0de3aaca714f1f1abd7edd562e3f1c6" diff --git a/tests/test_control_robot.py b/tests/test_control_robot.py index f52c1146..e7dd4966 100644 --- a/tests/test_control_robot.py +++ b/tests/test_control_robot.py @@ -19,14 +19,14 @@ def test_teleoperate(): def test_record_dataset_and_replay_episode_and_run_policy(tmpdir): robot_name = "koch" env_name = "koch_real" - policy_name = "act_real" + policy_name = "act_koch_real" #root = Path(tmpdir) root = Path("tmp/data") repo_id = "lerobot/debug" robot = make_robot(robot_name) - dataset = record_dataset(robot, fps=30, root=root, repo_id=repo_id, warmup_time_s=2, episode_time_s=2, num_episodes=2) + dataset = record_dataset(robot, fps=30, root=root, repo_id=repo_id, warmup_time_s=1, episode_time_s=1, num_episodes=2) replay_episode(robot, episode=0, fps=30, root=root, repo_id=repo_id)