From bb1d5f6c708a0877bc0146b8642be95eaae56e00 Mon Sep 17 00:00:00 2001 From: Steven Palma Date: Tue, 6 Jan 2026 17:12:24 +0100 Subject: [PATCH 1/3] feat(visualization): allow remote viewer + compress rerun images --- src/lerobot/scripts/lerobot_dataset_viz.py | 14 ++++++++++-- src/lerobot/scripts/lerobot_record.py | 15 +++++++++++-- src/lerobot/scripts/lerobot_teleoperate.py | 13 ++++++++++- src/lerobot/utils/visualization_utils.py | 25 +++++++++++++++++----- 4 files changed, 57 insertions(+), 10 deletions(-) diff --git a/src/lerobot/scripts/lerobot_dataset_viz.py b/src/lerobot/scripts/lerobot_dataset_viz.py index 974762b0b21..2cd48eab8b4 100644 --- a/src/lerobot/scripts/lerobot_dataset_viz.py +++ b/src/lerobot/scripts/lerobot_dataset_viz.py @@ -96,6 +96,7 @@ def visualize_dataset( ws_port: int = 9087, save: bool = False, output_dir: Path | None = None, + display_compressed_images: bool = False, ) -> Path | None: if save: assert output_dir is not None, ( @@ -137,8 +138,9 @@ def visualize_dataset( # display each camera image for key in dataset.meta.camera_keys: - # TODO(rcadene): add `.compress()`? is it lossless? - rr.log(key, rr.Image(to_hwc_uint8_numpy(batch[key][i]))) + img = to_hwc_uint8_numpy(batch[key][i]) + img_entity = rr.Image(img).compress() if display_compressed_images else rr.Image(img) + rr.log(key, entity=img_entity) # display each dimension of action space (e.g. actuators command) if ACTION in batch: @@ -261,6 +263,14 @@ def main(): ), ) + parser.add_argument( + "--display-compressed-images", + type=bool, + required=True, + default=False, + help="If set, display compressed images in Rerun instead of uncompressed ones.", + ) + args = parser.parse_args() kwargs = vars(args) repo_id = kwargs.pop("repo_id") diff --git a/src/lerobot/scripts/lerobot_record.py b/src/lerobot/scripts/lerobot_record.py index eafd32acefc..20c58785167 100644 --- a/src/lerobot/scripts/lerobot_record.py +++ b/src/lerobot/scripts/lerobot_record.py @@ -185,6 +185,10 @@ class RecordConfig: policy: PreTrainedConfig | None = None # Display all cameras on screen display_data: bool = False + # Display data on a remote Rerun server + display_url: str | None = None + # Port of the remote Rerun server + display_port: int | None = None # Use vocal synthesis to read events. play_sounds: bool = True # Resume recording on an existing dataset. @@ -261,6 +265,7 @@ def record_loop( control_time_s: int | None = None, single_task: str | None = None, display_data: bool = False, + display_compressed_images: bool = False, ): if dataset is not None and dataset.fps != fps: raise ValueError(f"The dataset fps should be equal to requested fps ({dataset.fps} != {fps}).") @@ -371,7 +376,9 @@ def record_loop( dataset.add_frame(frame) if display_data: - log_rerun_data(observation=obs_processed, action=action_values) + log_rerun_data( + observation=obs_processed, action=action_values, compress_images=display_compressed_images + ) dt_s = time.perf_counter() - start_loop_t precise_sleep(1 / fps - dt_s) @@ -384,7 +391,10 @@ def record(cfg: RecordConfig) -> LeRobotDataset: init_logging() logging.info(pformat(asdict(cfg))) if cfg.display_data: - init_rerun(session_name="recording") + init_rerun(session_name="recording", url=cfg.display_url, port=cfg.display_port) + display_compressed_images = ( + cfg.display_data and cfg.display_url is not None and cfg.display_port is not None + ) robot = make_robot_from_config(cfg.robot) teleop = make_teleoperator_from_config(cfg.teleop) if cfg.teleop is not None else None @@ -478,6 +488,7 @@ def record(cfg: RecordConfig) -> LeRobotDataset: control_time_s=cfg.dataset.episode_time_s, single_task=cfg.dataset.single_task, display_data=cfg.display_data, + display_compressed_images=display_compressed_images, ) # Execute a few seconds without recording to give time to manually reset the environment diff --git a/src/lerobot/scripts/lerobot_teleoperate.py b/src/lerobot/scripts/lerobot_teleoperate.py index bf722d6f15b..9138b36615f 100644 --- a/src/lerobot/scripts/lerobot_teleoperate.py +++ b/src/lerobot/scripts/lerobot_teleoperate.py @@ -108,6 +108,10 @@ class TeleoperateConfig: teleop_time_s: float | None = None # Display all cameras on screen display_data: bool = False + # Display data on a remote Rerun server + display_url: str | None = None + # Port of the remote Rerun server + display_port: int | None = None def teleop_loop( @@ -119,6 +123,7 @@ def teleop_loop( robot_observation_processor: RobotProcessorPipeline[RobotObservation, RobotObservation], display_data: bool = False, duration: float | None = None, + display_compressed_images: bool = False, ): """ This function continuously reads actions from a teleoperation device, processes them through optional @@ -130,6 +135,7 @@ def teleop_loop( robot: The robot instance being controlled. fps: The target frequency for the control loop in frames per second. display_data: If True, fetches robot observations and displays them in the console and Rerun. + display_compressed_images: If True, compresses images before sending them to Rerun for display. duration: The maximum duration of the teleoperation loop in seconds. If None, the loop runs indefinitely. teleop_action_processor: An optional pipeline to process raw actions from the teleoperator. robot_action_processor: An optional pipeline to process actions before they are sent to the robot. @@ -167,6 +173,7 @@ def teleop_loop( log_rerun_data( observation=obs_transition, action=teleop_action, + compress_images=display_compressed_images, ) print("\n" + "-" * (display_len + 10)) @@ -191,7 +198,10 @@ def teleoperate(cfg: TeleoperateConfig): init_logging() logging.info(pformat(asdict(cfg))) if cfg.display_data: - init_rerun(session_name="teleoperation") + init_rerun(session_name="teleoperation", url=cfg.display_url, port=cfg.display_port) + display_compressed_images = ( + cfg.display_data and cfg.display_url is not None and cfg.display_port is not None + ) teleop = make_teleoperator_from_config(cfg.teleop) robot = make_robot_from_config(cfg.robot) @@ -210,6 +220,7 @@ def teleoperate(cfg: TeleoperateConfig): teleop_action_processor=teleop_action_processor, robot_action_processor=robot_action_processor, robot_observation_processor=robot_observation_processor, + display_compressed_images=display_compressed_images, ) except KeyboardInterrupt: pass diff --git a/src/lerobot/utils/visualization_utils.py b/src/lerobot/utils/visualization_utils.py index 991b10247fa..a48ef35bb96 100644 --- a/src/lerobot/utils/visualization_utils.py +++ b/src/lerobot/utils/visualization_utils.py @@ -22,13 +22,25 @@ from .constants import OBS_PREFIX, OBS_STR -def init_rerun(session_name: str = "lerobot_control_loop") -> None: - """Initializes the Rerun SDK for visualizing the control loop.""" +def init_rerun( + session_name: str = "lerobot_control_loop", url: str | None = None, port: int | None = None +) -> None: + """ + Initializes the Rerun SDK for visualizing the control loop. + + Args: + session_name: Name of the Rerun session. + url: Optional URL for connecting to a Rerun server. + port: Optional port for connecting to a Rerun server. + """ batch_size = os.getenv("RERUN_FLUSH_NUM_BYTES", "8000") os.environ["RERUN_FLUSH_NUM_BYTES"] = batch_size rr.init(session_name) memory_limit = os.getenv("LEROBOT_RERUN_MEMORY_LIMIT", "10%") - rr.spawn(memory_limit=memory_limit) + if url and port: + rr.connect_grpc(url=f"rerun+http://{url}:{port}/proxy") + else: + rr.spawn(memory_limit=memory_limit) def _is_scalar(x): @@ -40,6 +52,7 @@ def _is_scalar(x): def log_rerun_data( observation: dict[str, Any] | None = None, action: dict[str, Any] | None = None, + compress_images: bool = False, ) -> None: """ Logs observation and action data to Rerun for real-time visualization. @@ -48,7 +61,7 @@ def log_rerun_data( to the Rerun viewer. It handles different data types appropriately: - Scalars values (floats, ints) are logged as `rr.Scalars`. - 3D NumPy arrays that resemble images (e.g., with 1, 3, or 4 channels first) are transposed - from CHW to HWC format and logged as `rr.Image`. + from CHW to HWC format, (optionally) compressed to JPEG and logged as `rr.Image` or `rr.EncodedImage`. - 1D NumPy arrays are logged as a series of individual scalars, with each element indexed. - Other multi-dimensional arrays are flattened and logged as individual scalars. @@ -57,6 +70,7 @@ def log_rerun_data( Args: observation: An optional dictionary containing observation data to log. action: An optional dictionary containing action data to log. + compress_images: Whether to compress images before logging to save bandwidth & memory in exchange for cpu and quality. """ if observation: for k, v in observation.items(): @@ -75,7 +89,8 @@ def log_rerun_data( for i, vi in enumerate(arr): rr.log(f"{key}_{i}", rr.Scalars(float(vi))) else: - rr.log(key, rr.Image(arr), static=True) + img_entity = rr.Image(arr).compress() if compress_images else rr.Image(arr) + rr.log(key, entity=img_entity, static=True) if action: for k, v in action.items(): From d963999cc8921944160455465ba39bde781d8969 Mon Sep 17 00:00:00 2001 From: Steven Palma Date: Tue, 6 Jan 2026 17:25:42 +0100 Subject: [PATCH 2/3] fix(tests): allow named argument in mocked rerun --- tests/utils/test_visualization_utils.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/utils/test_visualization_utils.py b/tests/utils/test_visualization_utils.py index f573de1668b..408f636cb5a 100644 --- a/tests/utils/test_visualization_utils.py +++ b/tests/utils/test_visualization_utils.py @@ -41,7 +41,10 @@ class DummyImage: def __init__(self, arr): self.arr = arr - def dummy_log(key, obj, **kwargs): + def dummy_log(key, obj=None, **kwargs): + # Accept either positional `obj` or keyword `entity` and record remaining kwargs. + if obj is None and "entity" in kwargs: + obj = kwargs.pop("entity") calls.append((key, obj, kwargs)) dummy_rr = SimpleNamespace( From 2025de6d0b73a80a3ca6ec6b2fdc3d383a542e8a Mon Sep 17 00:00:00 2001 From: Steven Palma Date: Wed, 7 Jan 2026 12:02:59 +0100 Subject: [PATCH 3/3] feat(visualization): ip instead or url & cli arg for compressing images --- src/lerobot/scripts/lerobot_record.py | 10 +++++++--- src/lerobot/scripts/lerobot_teleoperate.py | 10 +++++++--- src/lerobot/utils/visualization_utils.py | 8 ++++---- 3 files changed, 18 insertions(+), 10 deletions(-) diff --git a/src/lerobot/scripts/lerobot_record.py b/src/lerobot/scripts/lerobot_record.py index 20c58785167..5325f441a3a 100644 --- a/src/lerobot/scripts/lerobot_record.py +++ b/src/lerobot/scripts/lerobot_record.py @@ -186,9 +186,11 @@ class RecordConfig: # Display all cameras on screen display_data: bool = False # Display data on a remote Rerun server - display_url: str | None = None + display_ip: str | None = None # Port of the remote Rerun server display_port: int | None = None + # Whether to display compressed images in Rerun + display_compressed_images: bool = False # Use vocal synthesis to read events. play_sounds: bool = True # Resume recording on an existing dataset. @@ -391,9 +393,11 @@ def record(cfg: RecordConfig) -> LeRobotDataset: init_logging() logging.info(pformat(asdict(cfg))) if cfg.display_data: - init_rerun(session_name="recording", url=cfg.display_url, port=cfg.display_port) + init_rerun(session_name="recording", ip=cfg.display_ip, port=cfg.display_port) display_compressed_images = ( - cfg.display_data and cfg.display_url is not None and cfg.display_port is not None + True + if (cfg.display_data and cfg.display_ip is not None and cfg.display_port is not None) + else cfg.display_compressed_images ) robot = make_robot_from_config(cfg.robot) diff --git a/src/lerobot/scripts/lerobot_teleoperate.py b/src/lerobot/scripts/lerobot_teleoperate.py index 9138b36615f..821abe0e24b 100644 --- a/src/lerobot/scripts/lerobot_teleoperate.py +++ b/src/lerobot/scripts/lerobot_teleoperate.py @@ -109,9 +109,11 @@ class TeleoperateConfig: # Display all cameras on screen display_data: bool = False # Display data on a remote Rerun server - display_url: str | None = None + display_ip: str | None = None # Port of the remote Rerun server display_port: int | None = None + # Whether to display compressed images in Rerun + display_compressed_images: bool = False def teleop_loop( @@ -198,9 +200,11 @@ def teleoperate(cfg: TeleoperateConfig): init_logging() logging.info(pformat(asdict(cfg))) if cfg.display_data: - init_rerun(session_name="teleoperation", url=cfg.display_url, port=cfg.display_port) + init_rerun(session_name="teleoperation", ip=cfg.display_ip, port=cfg.display_port) display_compressed_images = ( - cfg.display_data and cfg.display_url is not None and cfg.display_port is not None + True + if (cfg.display_data and cfg.display_ip is not None and cfg.display_port is not None) + else cfg.display_compressed_images ) teleop = make_teleoperator_from_config(cfg.teleop) diff --git a/src/lerobot/utils/visualization_utils.py b/src/lerobot/utils/visualization_utils.py index a48ef35bb96..9143d0f660f 100644 --- a/src/lerobot/utils/visualization_utils.py +++ b/src/lerobot/utils/visualization_utils.py @@ -23,22 +23,22 @@ def init_rerun( - session_name: str = "lerobot_control_loop", url: str | None = None, port: int | None = None + session_name: str = "lerobot_control_loop", ip: str | None = None, port: int | None = None ) -> None: """ Initializes the Rerun SDK for visualizing the control loop. Args: session_name: Name of the Rerun session. - url: Optional URL for connecting to a Rerun server. + ip: Optional IP for connecting to a Rerun server. port: Optional port for connecting to a Rerun server. """ batch_size = os.getenv("RERUN_FLUSH_NUM_BYTES", "8000") os.environ["RERUN_FLUSH_NUM_BYTES"] = batch_size rr.init(session_name) memory_limit = os.getenv("LEROBOT_RERUN_MEMORY_LIMIT", "10%") - if url and port: - rr.connect_grpc(url=f"rerun+http://{url}:{port}/proxy") + if ip and port: + rr.connect_grpc(url=f"rerun+http://{ip}:{port}/proxy") else: rr.spawn(memory_limit=memory_limit)