From 9a22c9834da2f4a250b0c377bdb961b040ed24a1 Mon Sep 17 00:00:00 2001 From: Steven Palma Date: Thu, 27 Mar 2025 14:15:49 +0100 Subject: [PATCH] fix(visualization): init & log also for lekiwi remote robot --- lerobot/common/robot_devices/control_configs.py | 2 ++ lerobot/common/robot_devices/control_utils.py | 3 ++- lerobot/scripts/control_robot.py | 4 +++- 3 files changed, 7 insertions(+), 2 deletions(-) diff --git a/lerobot/common/robot_devices/control_configs.py b/lerobot/common/robot_devices/control_configs.py index fd7b5425..f61907bd 100644 --- a/lerobot/common/robot_devices/control_configs.py +++ b/lerobot/common/robot_devices/control_configs.py @@ -116,6 +116,8 @@ class ReplayControlConfig(ControlConfig): @dataclass class RemoteRobotConfig(ControlConfig): log_interval: int = 100 + # Display all cameras on screen + display_data: bool = False @dataclass diff --git a/lerobot/common/robot_devices/control_utils.py b/lerobot/common/robot_devices/control_utils.py index 63a3e7bb..c24e8778 100644 --- a/lerobot/common/robot_devices/control_utils.py +++ b/lerobot/common/robot_devices/control_utils.py @@ -264,7 +264,8 @@ def control_loop( frame = {**observation, **action, "task": single_task} dataset.add_frame(frame) - if display_data and not is_headless(): + # TODO(Steven): This should be more general (for RemoteRobot instead of checking the name, but anyways it will change soon) + if (display_data and not is_headless()) or (display_data and robot.robot_type.startswith("lekiwi")): for k, v in action.items(): for i, vv in enumerate(v): rr.log(f"sent_{k}_{i}", rr.Scalar(vv.numpy())) diff --git a/lerobot/scripts/control_robot.py b/lerobot/scripts/control_robot.py index f6020447..337d052b 100644 --- a/lerobot/scripts/control_robot.py +++ b/lerobot/scripts/control_robot.py @@ -378,7 +378,9 @@ def _init_rerun(control_config: ControlConfig, session_name: str = "lerobot_cont Raises: ValueError: If viewer IP is missing for non-remote configurations with display enabled. """ - if control_config.display_data and not is_headless(): + if (control_config.display_data and not is_headless()) or ( + control_config.display_data and isinstance(control_config, RemoteRobotConfig) + ): # Configure Rerun flush batch size batch_size = os.getenv("RERUN_FLUSH_NUM_BYTES", "8000") os.environ["RERUN_FLUSH_NUM_BYTES"] = batch_size