fix(visualization): init & log also for lekiwi remote robot

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

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()))

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