forked from tangger/lerobot
Address Michel's comments
This commit is contained in:
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user