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

Fix slow camera fps with Aloha (#433)

This commit is contained in:
Remi 2024-09-12 14:20:24 +02:00 committed by GitHub
parent a88dd602d9
commit beaa427504
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 18 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

31
lerobot/scripts/control_robot.py
View File

@ -179,13 +179,18 @@ def none_or_int(value):
def log_control_info(robot, dt_s, episode_index=None, frame_index=None, fps=None): def log_control_info(robot, dt_s, episode_index=None, frame_index=None, fps=None):
log_items = [] log_items = []
if episode_index is not None: if episode_index is not None:
log_items += [f"ep:{episode_index}"] log_items.append(f"ep:{episode_index}")
if frame_index is not None: if frame_index is not None:
log_items += [f"frame:{frame_index}"] log_items.append(f"frame:{frame_index}")
def log_dt(shortname, dt_val_s): def log_dt(shortname, dt_val_s):
nonlocal log_items nonlocal log_items, fps
log_items += [f"{shortname}:{dt_val_s * 1000:5.2f} ({1/ dt_val_s:3.1f}hz)"] info_str = f"{shortname}:{dt_val_s * 1000:5.2f} ({1/ dt_val_s:3.1f}hz)"
if fps is not None:
actual_fps = 1 / dt_val_s
if actual_fps < fps - 1:
info_str = colored(info_str, "yellow")
log_items.append(info_str)
# total step time displayed in milliseconds and its frequency # total step time displayed in milliseconds and its frequency
log_dt("dt", dt_s) log_dt("dt", dt_s)
@ -210,10 +215,6 @@ def log_control_info(robot, dt_s, episode_index=None, frame_index=None, fps=None
log_dt(f"dtR{name}", robot.logs[key]) log_dt(f"dtR{name}", robot.logs[key])
info_str = " ".join(log_items) info_str = " ".join(log_items)
if fps is not None:
actual_fps = 1 / dt_s
if actual_fps < fps - 1:
info_str = colored(info_str, "yellow")
logging.info(info_str) logging.info(info_str)
@ -320,7 +321,7 @@ def record(
run_compute_stats=True, run_compute_stats=True,
push_to_hub=True, push_to_hub=True,
tags=None, tags=None,
num_image_writers=8, num_image_writers_per_camera=4,
force_override=False, force_override=False,
): ):
# TODO(rcadene): Add option to record logs # TODO(rcadene): Add option to record logs
@ -442,8 +443,8 @@ def record(
# 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.
futures = [] futures = []
num_image_writers = num_image_writers_per_camera * len(robot.cameras)
with concurrent.futures.ThreadPoolExecutor(max_workers=num_image_writers) as executor: with concurrent.futures.ThreadPoolExecutor(max_workers=num_image_writers) as executor:
# Start recording all episodes # Start recording all episodes
while episode_index < num_episodes: while episode_index < num_episodes:
@ -803,10 +804,14 @@ if __name__ == "__main__":
help="Add tags to your dataset on the hub.", help="Add tags to your dataset on the hub.",
) )
parser_record.add_argument( parser_record.add_argument(
"--num-image-writers", "--num-image-writers-per-camera",
type=int, type=int,
default=8, default=4,
help="Number of threads writing the frames as png images on disk. Don't set too much as you might get unstable fps due to main thread being blocked.", help=(
"Number of threads writing the frames as png images on disk, per camera. "
"Too much threads might cause unstable teleoperation fps due to main thread being blocked. "
"Not enough threads might cause low camera fps."
),
) )
parser_record.add_argument( parser_record.add_argument(
"--force-override", "--force-override",