Address Michel's comments

This commit is contained in:
Simon Alibert
2025-06-04 12:47:24 +02:00
parent f5b1ef0045
commit 52028f5201
2 changed files with 14 additions and 14 deletions

View File

@@ -92,7 +92,7 @@ class DatasetRecordConfig:
single_task: str
# Root directory where the dataset will be stored (e.g. 'dataset/path').
root: str | Path | None = None
# Limit the frames per second. By default, uses the policy fps.
# Limit the frames per second.
fps: int = 30
# Number of seconds for data recording for each episode.
episode_time_s: int | float = 60
@@ -159,15 +159,15 @@ class RecordConfig:
def record_loop(
robot: Robot,
events: dict,
fps: int,
dataset: LeRobotDataset | None = None,
teleop: Teleoperator | None = None,
policy: PreTrainedPolicy | None = None,
control_time_s: int | None = None,
fps: int | None = None,
single_task: str | None = None,
display_data: bool = False,
):
if dataset is not None and fps is not None and dataset.fps != fps:
if dataset is not None and dataset.fps != fps:
raise ValueError(f"The dataset fps should be equal to requested fps ({dataset.fps} != {fps}).")
# if policy is given it needs cleaning up
@@ -216,12 +216,8 @@ def record_loop(
if isinstance(val, float):
rr.log(f"action.{act}", rr.Scalar(val))
if fps is not None:
dt_s = time.perf_counter() - start_loop_t
busy_wait(1 / fps - dt_s)
dt_s = time.perf_counter() - start_loop_t
# log_control_info(robot, dt_s, fps=fps)
busy_wait(1 / fps - dt_s)
timestamp = time.perf_counter() - start_episode_t
if events["exit_early"]:
@@ -283,11 +279,11 @@ def record(cfg: RecordConfig) -> LeRobotDataset:
record_loop(
robot=robot,
events=events,
fps=cfg.dataset.fps,
teleop=teleop,
policy=policy,
dataset=dataset,
control_time_s=cfg.dataset.episode_time_s,
fps=cfg.dataset.fps,
single_task=cfg.dataset.single_task,
display_data=cfg.display_data,
)
@@ -301,9 +297,9 @@ def record(cfg: RecordConfig) -> LeRobotDataset:
record_loop(
robot=robot,
events=events,
fps=cfg.dataset.fps,
teleop=teleop,
control_time_s=cfg.dataset.reset_time_s,
fps=cfg.dataset.fps,
single_task=cfg.dataset.single_task,
display_data=cfg.display_data,
)

View File

@@ -54,6 +54,7 @@ from lerobot.common.teleoperators import (
TeleoperatorConfig,
make_teleoperator_from_config,
)
from lerobot.common.utils.robot_utils import busy_wait
from lerobot.common.utils.utils import init_logging, move_cursor_up
from lerobot.common.utils.visualization_utils import _init_rerun
@@ -64,15 +65,15 @@ from .common.teleoperators import koch_leader, so100_leader, so101_leader # noq
class TeleoperateConfig:
teleop: TeleoperatorConfig
robot: RobotConfig
# Limit the maximum frames per second. By default, no limit.
fps: int | None = None
# Limit the maximum frames per second.
fps: int = 60
teleop_time_s: float | None = None
# Display all cameras on screen
display_data: bool = False
def teleop_loop(
teleop: Teleoperator, robot: Robot, display_data: bool = False, duration: float | None = None
teleop: Teleoperator, robot: Robot, fps: int, display_data: bool = False, duration: float | None = None
):
display_len = max(len(key) for key in robot.action_features)
start = time.perf_counter()
@@ -91,6 +92,9 @@ def teleop_loop(
rr.log(f"action_{act}", rr.Scalar(val))
robot.send_action(action)
dt_s = time.perf_counter() - loop_start
busy_wait(1 / fps - dt_s)
loop_s = time.perf_counter() - loop_start
print("\n" + "-" * (display_len + 10))
@@ -119,7 +123,7 @@ def teleoperate(cfg: TeleoperateConfig):
robot.connect()
try:
teleop_loop(teleop, robot, display_data=cfg.display_data, duration=cfg.teleop_time_s)
teleop_loop(teleop, robot, cfg.fps, display_data=cfg.display_data, duration=cfg.teleop_time_s)
except KeyboardInterrupt:
pass
finally: