wangzixiang
/
lerobot
mirror of https://github.com/huggingface/lerobot.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Almost done

This commit is contained in:
Remi Cadene 2024-07-10 00:07:40 +02:00
parent 798373e7bf
commit 52e760a88e
4 changed files with 58 additions and 55 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
lerobot/common/robot_devices/robots/factory.py
View File

@ -34,8 +34,8 @@ def make_robot(name):
), ),
}, },
cameras={ cameras={
"macbookpro": OpenCVCamera(1, fps=30, width=640, height=480), "laptop": OpenCVCamera(1, fps=30, width=640, height=480),
"iphone": OpenCVCamera(2, fps=30, width=640, height=480), "phone": OpenCVCamera(2, fps=30, width=640, height=480),
}, },
) )
else: else:

98
lerobot/scripts/control_robot.py
View File

@ -209,33 +209,35 @@ def record_dataset(
# Save images using threads to reach high fps (30 and more) # Save images using threads to reach high fps (30 and more)
# Using `with` to exist smoothly if an execption is raised. # Using `with` to exist smoothly if an execption is raised.
# Using only 4 worker threads to avoid blocking the main thread. # 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 futures = []
# 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
now = time.perf_counter() # Execute a few seconds without recording data, to give times
observation, action = robot.teleop_step(record_data=True) # 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 now = time.perf_counter()
busy_wait(1 / fps - dt_s) observation, action = robot.teleop_step(record_data=True)
dt_s = time.perf_counter() - now dt_s = time.perf_counter() - now
log_control_info(robot, dt_s) 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 timestamp = time.perf_counter() - start_time
ep_dicts = []
for episode_index in range(num_episodes): # Start recording all episodes
ep_dicts = []
for episode_index in range(num_episodes):
with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
ep_dict = {} ep_dict = {}
frame_index = 0 frame_index = 0
timestamp = 0 timestamp = 0
@ -254,7 +256,8 @@ def record_dataset(
not_image_keys = [key for key in observation if "image" not in key] not_image_keys = [key for key in observation if "image" not in key]
for key in image_keys: 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: for key in not_image_keys:
if key not in ep_dict: if key not in ep_dict:
@ -276,40 +279,40 @@ def record_dataset(
timestamp = time.perf_counter() - start_time 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: for key in image_keys:
tmp_imgs_dir = videos_dir / f"{key}_episode_{episode_index:06d}" tmp_imgs_dir = videos_dir / f"{key}_episode_{episode_index:06d}"
fname = f"{key}_episode_{episode_index:06d}.mp4" fname = f"{key}_episode_{episode_index:06d}.mp4"
video_path = local_dir / "videos" / fname video_path = local_dir / "videos" / fname
encode_video_frames(tmp_imgs_dir, video_path, fps) encode_video_frames(tmp_imgs_dir, video_path, fps)
# TODO(rcadene): uncomment? # TODO(rcadene): uncomment?
# clean temporary images directory # clean temporary images directory
# shutil.rmtree(tmp_imgs_dir) # shutil.rmtree(tmp_imgs_dir)
# store the reference to the video frame # store the reference to the video frame
ep_dict[key] = [] ep_dict[key] = []
for i in range(num_frames): for i in range(num_frames):
ep_dict[key].append({"path": f"videos/{fname}", "timestamp": i / fps}) ep_dict[key].append({"path": f"videos/{fname}", "timestamp": i / fps})
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])
for key in action: for key in action:
ep_dict[key] = torch.stack(ep_dict[key]) ep_dict[key] = torch.stack(ep_dict[key])
ep_dict["episode_index"] = torch.tensor([episode_index] * num_frames) ep_dict["episode_index"] = torch.tensor([episode_index] * num_frames)
ep_dict["frame_index"] = torch.arange(0, num_frames, 1) ep_dict["frame_index"] = torch.arange(0, num_frames, 1)
ep_dict["timestamp"] = torch.arange(0, num_frames, 1) / fps ep_dict["timestamp"] = torch.arange(0, num_frames, 1) / fps
done = torch.zeros(num_frames, dtype=torch.bool) done = torch.zeros(num_frames, dtype=torch.bool)
done[-1] = True done[-1] = True
ep_dict["next.done"] = done ep_dict["next.done"] = done
ep_dicts.append(ep_dict) ep_dicts.append(ep_dict)
data_dict = concatenate_episodes(ep_dicts) data_dict = concatenate_episodes(ep_dicts)
@ -338,6 +341,7 @@ def record_dataset(
videos_dir=videos_dir, videos_dir=videos_dir,
) )
stats = compute_stats(lerobot_dataset) if run_compute_stats else {} 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 = hf_dataset.with_format(None) # to remove transforms that cant be saved
hf_dataset.save_to_disk(str(local_dir / "train")) hf_dataset.save_to_disk(str(local_dir / "train"))

7
poetry.lock generated
View File

@ -2375,9 +2375,8 @@ description = "Nvidia JIT LTO Library"
optional = false optional = false
python-versions = ">=3" python-versions = ">=3"
files = [ files = [
{file = "nvidia_nvjitlink_cu12-12.5.40-py3-none-manylinux2014_aarch64.whl", hash = "sha256:004186d5ea6a57758fd6d57052a123c73a4815adf365eb8dd6a85c9eaa7535ff"}, {file = "nvidia_nvjitlink_cu12-12.5.82-py3-none-manylinux2014_x86_64.whl", hash = "sha256:f9b37bc5c8cf7509665cb6ada5aaa0ce65618f2332b7d3e78e9790511f111212"},
{file = "nvidia_nvjitlink_cu12-12.5.40-py3-none-manylinux2014_x86_64.whl", hash = "sha256:d9714f27c1d0f0895cd8915c07a87a1d0029a0aa36acaf9156952ec2a8a12189"}, {file = "nvidia_nvjitlink_cu12-12.5.82-py3-none-win_amd64.whl", hash = "sha256:e782564d705ff0bf61ac3e1bf730166da66dd2fe9012f111ede5fc49b64ae697"},
{file = "nvidia_nvjitlink_cu12-12.5.40-py3-none-win_amd64.whl", hash = "sha256:c3401dc8543b52d3a8158007a0c1ab4e9c768fcbd24153a48c86972102197ddd"},
] ]
[[package]] [[package]]
@ -4364,4 +4363,4 @@ xarm = ["gym-xarm"]
[metadata] [metadata]
lock-version = "2.0" lock-version = "2.0"
python-versions = ">=3.10,<3.13" python-versions = ">=3.10,<3.13"
content-hash = "81dc830d3d36c67e2fe2aea6cc30829eb2977edbf49a037df21a5f329a01aee5" content-hash = "223a6496a630da8181f119634f96bed3e0de3aaca714f1f1abd7edd562e3f1c6"

4
tests/test_control_robot.py
View File

@ -19,14 +19,14 @@ def test_teleoperate():
def test_record_dataset_and_replay_episode_and_run_policy(tmpdir): def test_record_dataset_and_replay_episode_and_run_policy(tmpdir):
robot_name = "koch" robot_name = "koch"
env_name = "koch_real" env_name = "koch_real"
policy_name = "act_real" policy_name = "act_koch_real"
#root = Path(tmpdir) #root = Path(tmpdir)
root = Path("tmp/data") root = Path("tmp/data")
repo_id = "lerobot/debug" repo_id = "lerobot/debug"
robot = make_robot(robot_name) 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) replay_episode(robot, episode=0, fps=30, root=root, repo_id=repo_id)