diff --git a/src/lerobot/robots/so_follower/so_follower.py b/src/lerobot/robots/so_follower/so_follower.py index 0651f566c..759bc5618 100644 --- a/src/lerobot/robots/so_follower/so_follower.py +++ b/src/lerobot/robots/so_follower/so_follower.py @@ -68,9 +68,16 @@ class SOFollower(Robot): @property def _cameras_ft(self) -> dict[str, tuple]: - return { - cam: (self.config.cameras[cam].height, self.config.cameras[cam].width, 3) for cam in self.cameras - } + features: dict[str, tuple] = {} + for cam in self.cameras: + cam_cfg = self.config.cameras[cam] + features[cam] = (cam_cfg.height, cam_cfg.width, 3) + # Cameras with a depth stream (e.g. RealSense with use_depth=True) also + # emit a 2D depth feature; hw_to_dataset_features routes 2D shapes to + # ``observation.depth.`` with the depth-map marker. + if getattr(cam_cfg, "use_depth", False): + features[f"{cam}_depth"] = (cam_cfg.height, cam_cfg.width) + return features @cached_property def observation_features(self) -> dict[str, type | tuple]: @@ -190,6 +197,14 @@ class SOFollower(Robot): dt_ms = (time.perf_counter() - start) * 1e3 logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + # Cameras with a depth stream populate a sibling ``_depth`` key + # (consumed by hw_to_dataset_features / build_dataset_frame). + if getattr(self.config.cameras[cam_key], "use_depth", False): + start = time.perf_counter() + obs_dict[f"{cam_key}_depth"] = cam.read_latest_depth() + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read {cam_key} depth: {dt_ms:.1f}ms") + return obs_dict @check_if_not_connected