From 28d8fe8858402dec93af0aac534178169ad72899 Mon Sep 17 00:00:00 2001 From: CarolinePascal Date: Wed, 24 Jun 2026 21:29:02 +0200 Subject: [PATCH] feat(depth): add depth support for follower robots --- src/lerobot/robots/hope_jr/hope_jr_arm.py | 26 ++++++++++++++----- src/lerobot/robots/hope_jr/hope_jr_hand.py | 26 ++++++++++++++----- .../robots/koch_follower/koch_follower.py | 26 ++++++++++++++----- .../robots/omx_follower/omx_follower.py | 26 ++++++++++++++----- .../openarm_follower/openarm_follower.py | 26 ++++++++++++++----- .../rebot_b601_follower.py | 26 ++++++++++++++----- src/lerobot/robots/unitree_g1/unitree_g1.py | 16 +++++++++--- 7 files changed, 126 insertions(+), 46 deletions(-) diff --git a/src/lerobot/robots/hope_jr/hope_jr_arm.py b/src/lerobot/robots/hope_jr/hope_jr_arm.py index 4918bcae3..b606a4fe7 100644 --- a/src/lerobot/robots/hope_jr/hope_jr_arm.py +++ b/src/lerobot/robots/hope_jr/hope_jr_arm.py @@ -66,9 +66,14 @@ class HopeJrArm(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: + cfg = self.config.cameras[cam] + if getattr(cfg, "use_rgb", True): + features[cam] = (cfg.height, cfg.width, 3) + if getattr(cfg, "use_depth", False): + features[f"{cam}_depth"] = (cfg.height, cfg.width, 1) + return features @cached_property def observation_features(self) -> dict[str, type | tuple]: @@ -139,10 +144,17 @@ class HopeJrArm(Robot): # Capture images from cameras for cam_key, cam in self.cameras.items(): - start = time.perf_counter() - obs_dict[cam_key] = cam.read_latest() - dt_ms = (time.perf_counter() - start) * 1e3 - logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + if getattr(cam, "use_rgb", True): + start = time.perf_counter() + obs_dict[cam_key] = cam.read_latest() + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + + if getattr(cam, "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 diff --git a/src/lerobot/robots/hope_jr/hope_jr_hand.py b/src/lerobot/robots/hope_jr/hope_jr_hand.py index 566628724..ce70e7e13 100644 --- a/src/lerobot/robots/hope_jr/hope_jr_hand.py +++ b/src/lerobot/robots/hope_jr/hope_jr_hand.py @@ -102,9 +102,14 @@ class HopeJrHand(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: + cfg = self.config.cameras[cam] + if getattr(cfg, "use_rgb", True): + features[cam] = (cfg.height, cfg.width, 3) + if getattr(cfg, "use_depth", False): + features[f"{cam}_depth"] = (cfg.height, cfg.width, 1) + return features @cached_property def observation_features(self) -> dict[str, type | tuple]: @@ -170,10 +175,17 @@ class HopeJrHand(Robot): # Capture images from cameras for cam_key, cam in self.cameras.items(): - start = time.perf_counter() - obs_dict[cam_key] = cam.read_latest() - dt_ms = (time.perf_counter() - start) * 1e3 - logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + if getattr(cam, "use_rgb", True): + start = time.perf_counter() + obs_dict[cam_key] = cam.read_latest() + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + + if getattr(cam, "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 diff --git a/src/lerobot/robots/koch_follower/koch_follower.py b/src/lerobot/robots/koch_follower/koch_follower.py index 3f40ac738..de6f9c4a3 100644 --- a/src/lerobot/robots/koch_follower/koch_follower.py +++ b/src/lerobot/robots/koch_follower/koch_follower.py @@ -68,9 +68,14 @@ class KochFollower(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: + cfg = self.config.cameras[cam] + if getattr(cfg, "use_rgb", True): + features[cam] = (cfg.height, cfg.width, 3) + if getattr(cfg, "use_depth", False): + features[f"{cam}_depth"] = (cfg.height, cfg.width, 1) + return features @cached_property def observation_features(self) -> dict[str, type | tuple]: @@ -192,10 +197,17 @@ class KochFollower(Robot): # Capture images from cameras for cam_key, cam in self.cameras.items(): - start = time.perf_counter() - obs_dict[cam_key] = cam.read_latest() - dt_ms = (time.perf_counter() - start) * 1e3 - logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + if getattr(cam, "use_rgb", True): + start = time.perf_counter() + obs_dict[cam_key] = cam.read_latest() + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + + if getattr(cam, "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 diff --git a/src/lerobot/robots/omx_follower/omx_follower.py b/src/lerobot/robots/omx_follower/omx_follower.py index c30eec97a..b2cfb52e9 100644 --- a/src/lerobot/robots/omx_follower/omx_follower.py +++ b/src/lerobot/robots/omx_follower/omx_follower.py @@ -68,9 +68,14 @@ class OmxFollower(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: + cfg = self.config.cameras[cam] + if getattr(cfg, "use_rgb", True): + features[cam] = (cfg.height, cfg.width, 3) + if getattr(cfg, "use_depth", False): + features[f"{cam}_depth"] = (cfg.height, cfg.width, 1) + return features @cached_property def observation_features(self) -> dict[str, type | tuple]: @@ -175,10 +180,17 @@ class OmxFollower(Robot): # Capture images from cameras for cam_key, cam in self.cameras.items(): - start = time.perf_counter() - obs_dict[cam_key] = cam.read_latest() - dt_ms = (time.perf_counter() - start) * 1e3 - logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + if getattr(cam, "use_rgb", True): + start = time.perf_counter() + obs_dict[cam_key] = cam.read_latest() + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + + if getattr(cam, "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 diff --git a/src/lerobot/robots/openarm_follower/openarm_follower.py b/src/lerobot/robots/openarm_follower/openarm_follower.py index 020f24052..e2c7c8cf5 100644 --- a/src/lerobot/robots/openarm_follower/openarm_follower.py +++ b/src/lerobot/robots/openarm_follower/openarm_follower.py @@ -101,9 +101,14 @@ class OpenArmFollower(Robot): @property def _cameras_ft(self) -> dict[str, tuple]: """Camera features for observation space.""" - 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: + cfg = self.config.cameras[cam] + if getattr(cfg, "use_rgb", True): + features[cam] = (cfg.height, cfg.width, 3) + if getattr(cfg, "use_depth", False): + features[f"{cam}_depth"] = (cfg.height, cfg.width, 1) + return features @cached_property def observation_features(self) -> dict[str, type | tuple]: @@ -242,10 +247,17 @@ class OpenArmFollower(Robot): # Capture images from cameras for cam_key, cam in self.cameras.items(): - start = time.perf_counter() - obs_dict[cam_key] = cam.read_latest() - dt_ms = (time.perf_counter() - start) * 1e3 - logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + if getattr(cam, "use_rgb", True): + start = time.perf_counter() + obs_dict[cam_key] = cam.read_latest() + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + + if getattr(cam, "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") dt_ms = (time.perf_counter() - start) * 1e3 logger.debug(f"{self} get_observation took: {dt_ms:.1f}ms") diff --git a/src/lerobot/robots/rebot_b601_follower/rebot_b601_follower.py b/src/lerobot/robots/rebot_b601_follower/rebot_b601_follower.py index ec00f4aa9..bf989702b 100644 --- a/src/lerobot/robots/rebot_b601_follower/rebot_b601_follower.py +++ b/src/lerobot/robots/rebot_b601_follower/rebot_b601_follower.py @@ -80,9 +80,14 @@ class RebotB601Follower(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: + cfg = self.config.cameras[cam] + if getattr(cfg, "use_rgb", True): + features[cam] = (cfg.height, cfg.width, 3) + if getattr(cfg, "use_depth", False): + features[f"{cam}_depth"] = (cfg.height, cfg.width, 1) + return features @cached_property def observation_features(self) -> dict[str, type | tuple]: @@ -213,10 +218,17 @@ class RebotB601Follower(Robot): logger.debug(f"{self} read state: {dt_ms:.1f}ms") for cam_key, cam in self.cameras.items(): - start = time.perf_counter() - obs_dict[cam_key] = cam.read_latest() - dt_ms = (time.perf_counter() - start) * 1e3 - logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + if getattr(cam, "use_rgb", True): + start = time.perf_counter() + obs_dict[cam_key] = cam.read_latest() + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + + if getattr(cam, "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 diff --git a/src/lerobot/robots/unitree_g1/unitree_g1.py b/src/lerobot/robots/unitree_g1/unitree_g1.py index 25ec32716..5b8be0941 100644 --- a/src/lerobot/robots/unitree_g1/unitree_g1.py +++ b/src/lerobot/robots/unitree_g1/unitree_g1.py @@ -222,9 +222,14 @@ class UnitreeG1(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: + cfg = self.config.cameras[cam] + if getattr(cfg, "use_rgb", True): + features[cam] = (cfg.height, cfg.width, 3) + if getattr(cfg, "use_depth", False): + features[f"{cam}_depth"] = (cfg.height, cfg.width, 1) + return features @cached_property def observation_features(self) -> dict[str, type | tuple]: @@ -458,7 +463,10 @@ class UnitreeG1(Robot): # Cameras - read images from ZMQ cameras for cam_name, cam in self._cameras.items(): - obs[cam_name] = cam.read_latest() + if getattr(cam, "use_rgb", True): + obs[cam_name] = cam.read_latest() + if getattr(cam, "use_depth", False): + obs[f"{cam_name}_depth"] = cam.read_latest_depth() return obs