test(visualization): fixing visualization tests

src/lerobot/utils/visualization_utils.py

@@ -115,7 +115,7 @@ def log_rerun_data(
                         )
                     else:
                         img_entity = rr.Image(arr).compress() if compress_images else rr.Image(arr)
-                    rr.log(key, entity=img_entity)
+                    rr.log(key, entity=img_entity, static=True)
 
     if action:
         for k, v in action.items():

tests/utils/test_visualization_utils.py

@@ -43,6 +43,11 @@ def mock_rerun(monkeypatch):
         def __init__(self, arr):
             self.arr = arr
 
+    class DummyDepthImage:
+        def __init__(self, arr, colormap=None):
+            self.arr = arr
+            self.colormap = colormap
+
     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:
@@ -55,6 +60,8 @@ def mock_rerun(monkeypatch):
         __spec__=SimpleNamespace(name="rerun", submodule_search_locations=None),
         Scalars=DummyScalar,
         Image=DummyImage,
+        DepthImage=DummyDepthImage,
+        components=SimpleNamespace(Colormap=SimpleNamespace(Viridis="viridis")),
         log=dummy_log,
         init=lambda *a, **k: None,
         spawn=lambda *a, **k: None,
@@ -225,7 +232,7 @@ def test_log_rerun_data_kwargs_only(mock_rerun):
     assert temp.value == pytest.approx(10.0)
 
     img = _obj_for(calls, "observation.gray")
-    assert type(img).__name__ == "DummyImage"
+    assert type(img).__name__ == "DummyDepthImage"  # single-channel -> DepthImage
     assert img.arr.shape == (8, 8, 1)  # remains HWC
     assert _kwargs_for(calls, "observation.gray").get("static", False) is True