feat(rerun unit): adding correct depth unit display for rerun (foxglove does not support units yet)

src/lerobot/scripts/lerobot_dataset_viz.py

@@ -84,6 +84,7 @@ import torch
 import torch.utils.data
 import tqdm
 
+from lerobot.configs import DEFAULT_DEPTH_UNIT, DEPTH_MILLIMETER_UNIT
 from lerobot.datasets import LeRobotDataset
 from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD, SUCCESS
 from lerobot.utils.utils import init_logging
@@ -228,6 +229,9 @@ def visualize_dataset(
 
     logging.info("Logging to Rerun")
 
+    # In this script, depth frames and stats are dequantized to the default depth_output_unit on load.
+    depth_meter = 1000.0 if DEFAULT_DEPTH_UNIT == DEPTH_MILLIMETER_UNIT else 1.0
+
     # Use the dataset's q01/q99 depth statistics for robust depth range bounds
     depth_ranges = {}
     for key in dataset.meta.depth_keys:
@@ -254,6 +258,7 @@ def visualize_dataset(
                     depth = to_hwc_float32_numpy(batch[key][i])
                     depth_entity = rr.DepthImage(
                         depth,
+                        meter=depth_meter,
                         colormap=rr.components.Colormap.Viridis,
                         depth_range=depth_ranges.get(key),
                     )

src/lerobot/utils/rerun_visualization.py

@@ -24,6 +24,7 @@ import os
 
 import numpy as np
 
+from lerobot.configs import DEPTH_MILLIMETER_UNIT, infer_depth_unit
 from lerobot.types import RobotAction, RobotObservation
 
 from .constants import ACTION, ACTION_PREFIX, OBS_PREFIX, OBS_STR
@@ -161,7 +162,13 @@ def log_rerun_data(
                     observation_paths.add(key)
                 else:
                     if arr.shape[-1] == 1:
-                        img_entity = rr.DepthImage(arr, colormap=rr.components.Colormap.Viridis)
+                        # At record time, the depth unit is inferred from the frame type.
+                        depth_unit = infer_depth_unit(arr.dtype)
+                        img_entity = rr.DepthImage(
+                            arr,
+                            meter=1000.0 if depth_unit == DEPTH_MILLIMETER_UNIT else 1.0,
+                            colormap=rr.components.Colormap.Viridis,
+                        )
                     else:
                         img_entity = rr.Image(arr).compress() if compress_images else rr.Image(arr)
                     rr.log(key, entity=img_entity, static=True)