From b960524d93248152bca64e17828de8f00bb93f59 Mon Sep 17 00:00:00 2001
From: CarolinePascal <caroline8.pascal@gmail.com>
Date: Tue, 19 May 2026 16:11:21 +0200
Subject: [PATCH] feat(depth): persist depth metadata

---
 src/lerobot/datasets/dataset_metadata.py | 13 ++++++++++
 tests/datasets/test_dataset_metadata.py  | 31 ++++++++++++++++++++++++
 2 files changed, 44 insertions(+)

diff --git a/src/lerobot/datasets/dataset_metadata.py b/src/lerobot/datasets/dataset_metadata.py
index 3c58774c3..feaa6d543 100644
--- a/src/lerobot/datasets/dataset_metadata.py
+++ b/src/lerobot/datasets/dataset_metadata.py
@@ -338,6 +338,19 @@ class LeRobotDatasetMetadata:
         """Keys to access visual modalities stored as videos."""
         return [key for key, ft in self.features.items() if ft["dtype"] == "video"]
 
+    @property
+    def depth_keys(self) -> list[str]:
+        """Keys to access depth-map modalities stored as videos or images.
+
+        A depth key is a feature whose ``info`` dict carries ``"<dtype>.is_depth_map": True``.
+        """
+        return [
+            key
+            for key, ft in self.features.items()
+            # TODO(CarolinePascal): Make sure the legacy video_info works here as well.
+            if (ft.get("info") or {}).get(ft["dtype"] + ".is_depth_map", False)
+        ]
+
     @property
     def camera_keys(self) -> list[str]:
         """Keys to access visual modalities (regardless of their storage method)."""
diff --git a/tests/datasets/test_dataset_metadata.py b/tests/datasets/test_dataset_metadata.py
index 6c784c90b..9f417af5e 100644
--- a/tests/datasets/test_dataset_metadata.py
+++ b/tests/datasets/test_dataset_metadata.py
@@ -53,6 +53,16 @@ IMAGE_FEATURES = {
     },
 }
 
+DEPTH_FEATURES = {
+    **SIMPLE_FEATURES,
+    "observation.depth.laptop": {
+        "dtype": "video",
+        "shape": (64, 96),
+        "names": ["height", "width"],
+        "info": {"video.is_depth_map": True},
+    },
+}
+
 
 def _make_dummy_stats(features: dict) -> dict:
     """Create minimal episode stats matching the given features."""
@@ -142,6 +152,27 @@ def test_create_without_videos_has_no_video_path(tmp_path):
     assert meta.video_keys == []
 
 
+def test_depth_keys_property_filters_by_marker(tmp_path):
+    """``depth_keys`` selects only video features carrying ``video.is_depth_map=True``."""
+    features = {
+        **VIDEO_FEATURES,
+        **DEPTH_FEATURES,
+    }
+    meta = LeRobotDatasetMetadata.create(
+        repo_id="test/depth_keys", fps=DEFAULT_FPS, features=features, root=tmp_path / "depth_keys"
+    )
+
+    assert set(meta.video_keys) == {"observation.images.laptop", "observation.depth.laptop"}
+    assert meta.depth_keys == ["observation.depth.laptop"]
+
+
+def test_depth_keys_empty_when_no_marker(tmp_path):
+    meta = LeRobotDatasetMetadata.create(
+        repo_id="test/no_depth", fps=DEFAULT_FPS, features=VIDEO_FEATURES, root=tmp_path / "no_depth"
+    )
+    assert meta.depth_keys == []
+
+
 def test_create_raises_on_existing_directory(tmp_path):
     """create() raises if root directory already exists."""
     root = tmp_path / "existing"