From 6b8d4c75a6e692a50991fd497a3b021f75499830 Mon Sep 17 00:00:00 2001 From: Martino Russi <77496684+nepyope@users.noreply.github.com> Date: Mon, 12 Jan 2026 17:31:39 +0100 Subject: [PATCH 1/5] Feat/g1 improvements record sim (#2765) This PR extends the integration of Unitree g1 with the LeRobot codebase. By converting robot state to a flat dict we can now record and replay episodes (example groot/holosoma scripts need to be adjusted as well). We also improve the simulation integration by calling .step @ _subscribe_motor_state instead of it running in a separate thread. We also add ZMQ camera to lerobot, streaming base64 images over json --- docs/source/unitree_g1.mdx | 6 +- examples/unitree_g1/gr00t_locomotion.py | 47 ++- examples/unitree_g1/holosoma_locomotion.py | 28 +- src/lerobot/cameras/utils.py | 5 + src/lerobot/cameras/zmq/__init__.py | 20 ++ src/lerobot/cameras/zmq/camera_zmq.py | 235 +++++++++++++ src/lerobot/cameras/zmq/configuration_zmq.py | 46 +++ src/lerobot/cameras/zmq/image_server.py | 114 +++++++ .../robots/unitree_g1/config_unitree_g1.py | 5 + src/lerobot/robots/unitree_g1/unitree_g1.py | 314 ++++++++++++------ src/lerobot/scripts/lerobot_record.py | 7 + src/lerobot/scripts/lerobot_replay.py | 1 + 12 files changed, 679 insertions(+), 149 deletions(-) create mode 100644 src/lerobot/cameras/zmq/__init__.py create mode 100644 src/lerobot/cameras/zmq/camera_zmq.py create mode 100644 src/lerobot/cameras/zmq/configuration_zmq.py create mode 100644 src/lerobot/cameras/zmq/image_server.py diff --git a/docs/source/unitree_g1.mdx b/docs/source/unitree_g1.mdx index bdc7eb33d..e6bffdf1b 100644 --- a/docs/source/unitree_g1.mdx +++ b/docs/source/unitree_g1.mdx @@ -7,7 +7,7 @@ This guide covers the complete setup process for the Unitree G1 humanoid, from i We support both 29 and 23 DOF G1 EDU version. We introduce: - **`unitree g1` robot class, handling low level read/write from/to the humanoid** -- **ZMQ socket bridge** for remote communication over wlan, allowing for remote policy deployment as well as over eth or directly on the Orin +- **ZMQ socket bridge** for remote communication and camera streaming, allowing for remote policy deployment over wlan, eth or directly on the robot - **Locomotion policies** from NVIDIA gr00t and Amazon FAR Holosoma - **Simulation mode** for testing policies without the physical robot in mujoco @@ -110,7 +110,7 @@ ssh unitree@ # Password: 123 ``` -Replace `` with your robot's actual WiFi IP address (e.g., `172.18.129.215`). +Replace `` with your robot's actual WiFi IP address. --- @@ -188,7 +188,7 @@ Press `Ctrl+C` to stop the policy. ## Running in Simulation Mode (MuJoCo) -You can now test and develop policies without a physical robot using MuJoCo. To do so simply set `is_simulation=True` in config. +You can now test policies before unleashing them on the physical robot using MuJoCo. To do so simply set `is_simulation=True` in config. ## Additional Resources diff --git a/examples/unitree_g1/gr00t_locomotion.py b/examples/unitree_g1/gr00t_locomotion.py index 30e5a27e6..0123b5206 100644 --- a/examples/unitree_g1/gr00t_locomotion.py +++ b/examples/unitree_g1/gr00t_locomotion.py @@ -111,34 +111,29 @@ class GrootLocomotionController: def run_step(self): # Get current observation - robot_state = self.robot.get_observation() + obs = self.robot.get_observation() - if robot_state is None: + if not obs: return # Get command from remote controller - if robot_state.wireless_remote is not None: - self.robot.remote_controller.set(robot_state.wireless_remote) - if self.robot.remote_controller.button[0]: # R1 - raise waist - self.groot_height_cmd += 0.001 - self.groot_height_cmd = np.clip(self.groot_height_cmd, 0.50, 1.00) - if self.robot.remote_controller.button[4]: # R2 - lower waist - self.groot_height_cmd -= 0.001 - self.groot_height_cmd = np.clip(self.groot_height_cmd, 0.50, 1.00) - else: - self.robot.remote_controller.lx = 0.0 - self.robot.remote_controller.ly = 0.0 - self.robot.remote_controller.rx = 0.0 - self.robot.remote_controller.ry = 0.0 + if obs["remote.buttons"][0]: # R1 - raise waist + self.groot_height_cmd += 0.001 + self.groot_height_cmd = np.clip(self.groot_height_cmd, 0.50, 1.00) + if obs["remote.buttons"][4]: # R2 - lower waist + self.groot_height_cmd -= 0.001 + self.groot_height_cmd = np.clip(self.groot_height_cmd, 0.50, 1.00) - self.cmd[0] = self.robot.remote_controller.ly # Forward/backward - self.cmd[1] = self.robot.remote_controller.lx * -1 # Left/right - self.cmd[2] = self.robot.remote_controller.rx * -1 # Rotation rate + self.cmd[0] = obs["remote.ly"] # Forward/backward + self.cmd[1] = obs["remote.lx"] * -1 # Left/right + self.cmd[2] = obs["remote.rx"] * -1 # Rotation rate - # Get joint positions and velocities - for i in range(29): - self.groot_qj_all[i] = robot_state.motor_state[i].q - self.groot_dqj_all[i] = robot_state.motor_state[i].dq + # Get joint positions and velocities from flat dict + for motor in G1_29_JointIndex: + name = motor.name + idx = motor.value + self.groot_qj_all[idx] = obs[f"{name}.q"] + self.groot_dqj_all[idx] = obs[f"{name}.dq"] # Adapt observation for g1_23dof for idx in MISSING_JOINTS: @@ -150,8 +145,8 @@ class GrootLocomotionController: dqj_obs = self.groot_dqj_all.copy() # Express IMU data in gravity frame of reference - quat = robot_state.imu_state.quaternion - ang_vel = np.array(robot_state.imu_state.gyroscope, dtype=np.float32) + quat = [obs["imu.quat.w"], obs["imu.quat.x"], obs["imu.quat.y"], obs["imu.quat.z"]] + ang_vel = np.array([obs["imu.gyro.x"], obs["imu.gyro.y"], obs["imu.gyro.z"]], dtype=np.float32) gravity_orientation = self.robot.get_gravity_orientation(quat) # Scale joint positions and velocities before policy inference @@ -219,6 +214,8 @@ def run(repo_id: str = DEFAULT_GROOT_REPO_ID) -> None: config = UnitreeG1Config() robot = UnitreeG1(config) + robot.connect() + # Initialize gr00T locomotion controller groot_controller = GrootLocomotionController( policy_balance=policy_balance, @@ -234,7 +231,7 @@ def run(repo_id: str = DEFAULT_GROOT_REPO_ID) -> None: logger.info("Press Ctrl+C to stop") # Run step - while True: + while not robot._shutdown_event.is_set(): start_time = time.time() groot_controller.run_step() elapsed = time.time() - start_time diff --git a/examples/unitree_g1/holosoma_locomotion.py b/examples/unitree_g1/holosoma_locomotion.py index 017f7835a..3a07023de 100644 --- a/examples/unitree_g1/holosoma_locomotion.py +++ b/examples/unitree_g1/holosoma_locomotion.py @@ -126,24 +126,23 @@ class HolosomaLocomotionController: def run_step(self): # Get current observation - robot_state = self.robot.get_observation() + obs = self.robot.get_observation() - if robot_state is None: + if not obs: return # Get command from remote controller - if robot_state.wireless_remote is not None: - self.robot.remote_controller.set(robot_state.wireless_remote) - - ly = self.robot.remote_controller.ly if abs(self.robot.remote_controller.ly) > 0.1 else 0.0 - lx = self.robot.remote_controller.lx if abs(self.robot.remote_controller.lx) > 0.1 else 0.0 - rx = self.robot.remote_controller.rx if abs(self.robot.remote_controller.rx) > 0.1 else 0.0 + ly = obs["remote.ly"] if abs(obs["remote.ly"]) > 0.1 else 0.0 + lx = obs["remote.lx"] if abs(obs["remote.lx"]) > 0.1 else 0.0 + rx = obs["remote.rx"] if abs(obs["remote.rx"]) > 0.1 else 0.0 self.cmd[:] = [ly, -lx, -rx] # Get joint positions and velocities - for i in range(29): - self.qj[i] = robot_state.motor_state[i].q - self.dqj[i] = robot_state.motor_state[i].dq + for motor in G1_29_JointIndex: + name = motor.name + idx = motor.value + self.qj[idx] = obs[f"{name}.q"] + self.dqj[idx] = obs[f"{name}.dq"] # Adapt observation for g1_23dof for idx in MISSING_JOINTS: @@ -151,8 +150,8 @@ class HolosomaLocomotionController: self.dqj[idx] = 0.0 # Express IMU data in gravity frame of reference - quat = robot_state.imu_state.quaternion - ang_vel = np.array(robot_state.imu_state.gyroscope, dtype=np.float32) + quat = [obs["imu.quat.w"], obs["imu.quat.x"], obs["imu.quat.y"], obs["imu.quat.z"]] + ang_vel = np.array([obs["imu.gyro.x"], obs["imu.gyro.y"], obs["imu.gyro.z"]], dtype=np.float32) gravity = self.robot.get_gravity_orientation(quat) # Scale joint positions and velocities before policy inference @@ -220,6 +219,7 @@ def run(repo_id: str = DEFAULT_HOLOSOMA_REPO_ID, policy_type: str = "fastsac") - # Initialize robot config = UnitreeG1Config() robot = UnitreeG1(config) + robot.connect() holosoma_controller = HolosomaLocomotionController(policy, robot, kp, kd) @@ -230,7 +230,7 @@ def run(repo_id: str = DEFAULT_HOLOSOMA_REPO_ID, policy_type: str = "fastsac") - logger.info("Press Ctrl+C to stop") # Run step - while True: + while not robot._shutdown_event.is_set(): start_time = time.time() holosoma_controller.run_step() elapsed = time.time() - start_time diff --git a/src/lerobot/cameras/utils.py b/src/lerobot/cameras/utils.py index 1b2d386d6..c0e7b6284 100644 --- a/src/lerobot/cameras/utils.py +++ b/src/lerobot/cameras/utils.py @@ -43,6 +43,11 @@ def make_cameras_from_configs(camera_configs: dict[str, CameraConfig]) -> dict[s cameras[key] = Reachy2Camera(cfg) + elif cfg.type == "zmq": + from .zmq.camera_zmq import ZMQCamera + + cameras[key] = ZMQCamera(cfg) + else: try: cameras[key] = cast(Camera, make_device_from_device_class(cfg)) diff --git a/src/lerobot/cameras/zmq/__init__.py b/src/lerobot/cameras/zmq/__init__.py new file mode 100644 index 000000000..d760c5325 --- /dev/null +++ b/src/lerobot/cameras/zmq/__init__.py @@ -0,0 +1,20 @@ +#!/usr/bin/env python + +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .camera_zmq import ZMQCamera +from .configuration_zmq import ZMQCameraConfig + +__all__ = ["ZMQCamera", "ZMQCameraConfig"] diff --git a/src/lerobot/cameras/zmq/camera_zmq.py b/src/lerobot/cameras/zmq/camera_zmq.py new file mode 100644 index 000000000..1a4155f4b --- /dev/null +++ b/src/lerobot/cameras/zmq/camera_zmq.py @@ -0,0 +1,235 @@ +#!/usr/bin/env python + +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +ZMQCamera - Captures frames from remote cameras via ZeroMQ using JSON protocol in the +following format: + { + "timestamps": {"camera_name": float}, + "images": {"camera_name": ""} + } +""" + +import base64 +import json +import logging +import time +from threading import Event, Lock, Thread +from typing import Any + +import cv2 +import numpy as np +from numpy.typing import NDArray + +from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError + +from ..camera import Camera +from ..configs import ColorMode +from .configuration_zmq import ZMQCameraConfig + +logger = logging.getLogger(__name__) + + +class ZMQCamera(Camera): + """ + Example usage: + ```python + from lerobot.cameras.zmq import ZMQCamera, ZMQCameraConfig + + config = ZMQCameraConfig(server_address="192.168.123.164", port=5555, camera_name="head_camera") + camera = ZMQCamera(config) + camera.connect() + frame = camera.read() + camera.disconnect() + ``` + """ + + def __init__(self, config: ZMQCameraConfig): + super().__init__(config) + import zmq + + self.config = config + self.server_address = config.server_address + self.port = config.port + self.camera_name = config.camera_name + self.color_mode = config.color_mode + self.timeout_ms = config.timeout_ms + + self.context: zmq.Context | None = None + self.socket: zmq.Socket | None = None + self._connected = False + + self.thread: Thread | None = None + self.stop_event: Event | None = None + self.frame_lock: Lock = Lock() + self.latest_frame: NDArray[Any] | None = None + self.new_frame_event: Event = Event() + + def __str__(self) -> str: + return f"ZMQCamera({self.camera_name}@{self.server_address}:{self.port})" + + @property + def is_connected(self) -> bool: + return self._connected and self.context is not None and self.socket is not None + + def connect(self, warmup: bool = True) -> None: + """Connect to ZMQ camera server.""" + if self.is_connected: + raise DeviceAlreadyConnectedError(f"{self} is already connected.") + + logger.info(f"Connecting to {self}...") + + try: + import zmq + + self.context = zmq.Context() + self.socket = self.context.socket(zmq.SUB) + self.socket.setsockopt_string(zmq.SUBSCRIBE, "") + self.socket.setsockopt(zmq.RCVTIMEO, self.timeout_ms) + self.socket.setsockopt(zmq.CONFLATE, True) + self.socket.connect(f"tcp://{self.server_address}:{self.port}") + self._connected = True + + # Auto-detect resolution + if self.width is None or self.height is None: + h, w = self.read().shape[:2] + self.height = h + self.width = w + logger.info(f"{self} resolution: {w}x{h}") + + logger.info(f"{self} connected.") + + if warmup: + time.sleep(0.1) + + except Exception as e: + self._cleanup() + raise RuntimeError(f"Failed to connect to {self}: {e}") from e + + def _cleanup(self): + """Clean up ZMQ resources.""" + self._connected = False + if self.socket: + self.socket.close() + self.socket = None + if self.context: + self.context.term() + self.context = None + + @staticmethod + def find_cameras() -> list[dict[str, Any]]: + """ZMQ cameras require manual configuration (server address/port).""" + return [] + + def read(self, color_mode: ColorMode | None = None) -> NDArray[Any]: + """ + Read a single frame from the ZMQ camera. + + Returns: + np.ndarray: Decoded frame (height, width, 3) + """ + if not self.is_connected or self.socket is None: + raise DeviceNotConnectedError(f"{self} is not connected.") + + try: + message = self.socket.recv_string() + except Exception as e: + if type(e).__name__ == "Again": + raise TimeoutError(f"{self} timeout after {self.timeout_ms}ms") from e + raise + + # Decode JSON message + data = json.loads(message) + + if "images" not in data: + raise RuntimeError(f"{self} invalid message: missing 'images' key") + + images = data["images"] + + # Get image by camera name or first available + if self.camera_name in images: + img_b64 = images[self.camera_name] + elif images: + img_b64 = next(iter(images.values())) + else: + raise RuntimeError(f"{self} no images in message") + + # Decode base64 JPEG + img_bytes = base64.b64decode(img_b64) + frame = cv2.imdecode(np.frombuffer(img_bytes, np.uint8), cv2.IMREAD_COLOR) + + if frame is None: + raise RuntimeError(f"{self} failed to decode image") + + return frame + + def _read_loop(self) -> None: + while self.stop_event and not self.stop_event.is_set(): + try: + frame = self.read() + with self.frame_lock: + self.latest_frame = frame + self.new_frame_event.set() + except DeviceNotConnectedError: + break + except TimeoutError: + pass + except Exception as e: + logger.warning(f"Read error: {e}") + + def _start_read_thread(self) -> None: + if self.thread and self.thread.is_alive(): + return + self.stop_event = Event() + self.thread = Thread(target=self._read_loop, daemon=True) + self.thread.start() + + def _stop_read_thread(self) -> None: + if self.stop_event: + self.stop_event.set() + if self.thread and self.thread.is_alive(): + self.thread.join(timeout=2.0) + self.thread = None + self.stop_event = None + + def async_read(self, timeout_ms: float = 10000) -> NDArray[Any]: + """Read latest frame asynchronously (non-blocking).""" + if not self.is_connected: + raise DeviceNotConnectedError(f"{self} is not connected.") + + if not self.thread or not self.thread.is_alive(): + self._start_read_thread() + + if not self.new_frame_event.wait(timeout=timeout_ms / 1000.0): + raise TimeoutError(f"{self} async_read timeout after {timeout_ms}ms") + + with self.frame_lock: + frame = self.latest_frame + self.new_frame_event.clear() + + if frame is None: + raise RuntimeError(f"{self} no frame available") + + return frame + + def disconnect(self) -> None: + """Disconnect from ZMQ camera.""" + if not self.is_connected and not self.thread: + raise DeviceNotConnectedError(f"{self} not connected.") + + self._stop_read_thread() + self._cleanup() + logger.info(f"{self} disconnected.") diff --git a/src/lerobot/cameras/zmq/configuration_zmq.py b/src/lerobot/cameras/zmq/configuration_zmq.py new file mode 100644 index 000000000..027ae12b5 --- /dev/null +++ b/src/lerobot/cameras/zmq/configuration_zmq.py @@ -0,0 +1,46 @@ +#!/usr/bin/env python + +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass + +from ..configs import CameraConfig, ColorMode + +__all__ = ["ZMQCameraConfig", "ColorMode"] + + +@CameraConfig.register_subclass("zmq") +@dataclass +class ZMQCameraConfig(CameraConfig): + server_address: str + port: int = 5555 + camera_name: str = "zmq_camera" + color_mode: ColorMode = ColorMode.RGB + timeout_ms: int = 5000 + + def __post_init__(self) -> None: + if self.color_mode not in (ColorMode.RGB, ColorMode.BGR): + raise ValueError( + f"`color_mode` is expected to be {ColorMode.RGB.value} or {ColorMode.BGR.value}, but {self.color_mode} is provided." + ) + + if self.timeout_ms <= 0: + raise ValueError(f"`timeout_ms` must be positive, but {self.timeout_ms} is provided.") + + if not self.server_address: + raise ValueError("`server_address` cannot be empty.") + + if self.port <= 0 or self.port > 65535: + raise ValueError(f"`port` must be between 1 and 65535, but {self.port} is provided.") diff --git a/src/lerobot/cameras/zmq/image_server.py b/src/lerobot/cameras/zmq/image_server.py new file mode 100644 index 000000000..2da366cef --- /dev/null +++ b/src/lerobot/cameras/zmq/image_server.py @@ -0,0 +1,114 @@ +#!/usr/bin/env python + +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Streams camera images over ZMQ. +Uses lerobot's OpenCVCamera for capture, encodes images to base64 and sends them over ZMQ. +""" + +import base64 +import contextlib +import json +import logging +import time +from collections import deque + +import cv2 +import numpy as np +import zmq + +from lerobot.cameras.configs import ColorMode +from lerobot.cameras.opencv import OpenCVCamera, OpenCVCameraConfig + +logger = logging.getLogger(__name__) + + +def encode_image(image: np.ndarray, quality: int = 80) -> str: + """Encode RGB image to base64 JPEG string.""" + _, buffer = cv2.imencode(".jpg", image, [int(cv2.IMWRITE_JPEG_QUALITY), quality]) + return base64.b64encode(buffer).decode("utf-8") + + +class ImageServer: + def __init__(self, config: dict, port: int = 5555): + self.fps = config.get("fps", 30) + self.cameras: dict[str, OpenCVCamera] = {} + + for name, cfg in config.get("cameras", {}).items(): + shape = cfg.get("shape", [480, 640]) + cam_config = OpenCVCameraConfig( + index_or_path=cfg.get("device_id", 0), + fps=self.fps, + width=shape[1], + height=shape[0], + color_mode=ColorMode.RGB, + ) + camera = OpenCVCamera(cam_config) + camera.connect() + self.cameras[name] = camera + logger.info(f"Camera {name}: {shape[1]}x{shape[0]}") + + # ZMQ PUB socket + self.context = zmq.Context() + self.socket = self.context.socket(zmq.PUB) + self.socket.setsockopt(zmq.SNDHWM, 20) + self.socket.setsockopt(zmq.LINGER, 0) + self.socket.bind(f"tcp://*:{port}") + + logger.info(f"ImageServer running on port {port}") + + def run(self): + frame_count = 0 + frame_times = deque(maxlen=60) + + try: + while True: + t0 = time.time() + + # Build message + message = {"timestamps": {}, "images": {}} + for name, cam in self.cameras.items(): + frame = cam.read() # Returns RGB + message["timestamps"][name] = time.time() + message["images"][name] = encode_image(frame) + + # Send as JSON string (suppress if buffer full) + with contextlib.suppress(zmq.Again): + self.socket.send_string(json.dumps(message), zmq.NOBLOCK) + + frame_count += 1 + frame_times.append(time.time() - t0) + + if frame_count % 60 == 0: + logger.debug(f"FPS: {len(frame_times) / sum(frame_times):.1f}") + + sleep = (1.0 / self.fps) - (time.time() - t0) + if sleep > 0: + time.sleep(sleep) + + except KeyboardInterrupt: + pass + finally: + for cam in self.cameras.values(): + cam.disconnect() + self.socket.close() + self.context.term() + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO) + config = {"fps": 30, "cameras": {"head_camera": {"device_id": 4, "shape": [480, 640]}}} + ImageServer(config, port=5555).run() diff --git a/src/lerobot/robots/unitree_g1/config_unitree_g1.py b/src/lerobot/robots/unitree_g1/config_unitree_g1.py index c66edbd1c..0b163019d 100644 --- a/src/lerobot/robots/unitree_g1/config_unitree_g1.py +++ b/src/lerobot/robots/unitree_g1/config_unitree_g1.py @@ -16,6 +16,8 @@ from dataclasses import dataclass, field +from lerobot.cameras import CameraConfig + from ..config import RobotConfig _GAINS: dict[str, dict[str, list[float]]] = { @@ -60,3 +62,6 @@ class UnitreeG1Config(RobotConfig): # Socket config for ZMQ bridge robot_ip: str = "192.168.123.164" # default G1 IP + + # Cameras (ZMQ-based remote cameras) + cameras: dict[str, CameraConfig] = field(default_factory=dict) diff --git a/src/lerobot/robots/unitree_g1/unitree_g1.py b/src/lerobot/robots/unitree_g1/unitree_g1.py index 5bc4f3110..1764f31b5 100644 --- a/src/lerobot/robots/unitree_g1/unitree_g1.py +++ b/src/lerobot/robots/unitree_g1/unitree_g1.py @@ -23,13 +23,8 @@ from functools import cached_property from typing import Any import numpy as np -from unitree_sdk2py.idl.default import unitree_hg_msg_dds__LowCmd_ -from unitree_sdk2py.idl.unitree_hg.msg.dds_ import ( - LowCmd_ as hg_LowCmd, - LowState_ as hg_LowState, -) -from unitree_sdk2py.utils.crc import CRC +from lerobot.cameras.utils import make_cameras_from_configs from lerobot.envs.factory import make_env from lerobot.robots.unitree_g1.g1_utils import G1_29_JointIndex @@ -43,8 +38,6 @@ logger = logging.getLogger(__name__) kTopicLowCommand_Debug = "rt/lowcmd" kTopicLowState = "rt/lowstate" -G1_29_Num_Motors = 29 - @dataclass class MotorState: @@ -66,28 +59,12 @@ class IMUState: # g1 observation class @dataclass class G1_29_LowState: # noqa: N801 - motor_state: list[MotorState] = field( - default_factory=lambda: [MotorState() for _ in range(G1_29_Num_Motors)] - ) + motor_state: list[MotorState] = field(default_factory=lambda: [MotorState() for _ in G1_29_JointIndex]) imu_state: IMUState = field(default_factory=IMUState) wireless_remote: Any = None # Raw wireless remote data mode_machine: int = 0 # Robot mode -class DataBuffer: - def __init__(self): - self.data = None - self.lock = threading.Lock() - - def get_data(self): - with self.lock: - return self.data - - def set_data(self, data): - with self.lock: - self.data = data - - class UnitreeG1(Robot): config_class = UnitreeG1Config name = "unitree_g1" @@ -117,9 +94,12 @@ class UnitreeG1(Robot): logger.info("Initialize UnitreeG1...") self.config = config - self.control_dt = config.control_dt + # Initialize cameras config (ZMQ-based) - actual connection in connect() + self._cameras = make_cameras_from_configs(config.cameras) + + # Import channel classes based on mode if config.is_simulation: from unitree_sdk2py.core.channel import ( ChannelFactoryInitialize, @@ -133,62 +113,33 @@ class UnitreeG1(Robot): ChannelSubscriber, ) - # connect robot - self.ChannelFactoryInitialize = ChannelFactoryInitialize - self.connect() + # Store for use in connect() + self._ChannelFactoryInitialize = ChannelFactoryInitialize + self._ChannelPublisher = ChannelPublisher + self._ChannelSubscriber = ChannelSubscriber - # initialize direct motor control interface - self.lowcmd_publisher = ChannelPublisher(kTopicLowCommand_Debug, hg_LowCmd) - self.lowcmd_publisher.Init() - self.lowstate_subscriber = ChannelSubscriber(kTopicLowState, hg_LowState) - self.lowstate_subscriber.Init() - self.lowstate_buffer = DataBuffer() - - # initialize subscribe thread to read robot state + # Initialize state variables + self.sim_env = None + self._env_wrapper = None + self._lowstate = None self._shutdown_event = threading.Event() - self.subscribe_thread = threading.Thread(target=self._subscribe_motor_state) - self.subscribe_thread.start() - - while not self.is_connected: - time.sleep(0.1) - - # initialize hg's lowcmd msg - self.crc = CRC() - self.msg = unitree_hg_msg_dds__LowCmd_() - self.msg.mode_pr = 0 - - # Wait for first state message to arrive - lowstate = None - while lowstate is None: - lowstate = self.lowstate_buffer.get_data() - if lowstate is None: - time.sleep(0.01) - logger.warning("[UnitreeG1] Waiting for robot state...") - logger.warning("[UnitreeG1] Connected to robot.") - self.msg.mode_machine = lowstate.mode_machine - - # initialize all motors with unified kp/kd from config - self.kp = np.array(config.kp, dtype=np.float32) - self.kd = np.array(config.kd, dtype=np.float32) - - for id in G1_29_JointIndex: - self.msg.motor_cmd[id].mode = 1 - self.msg.motor_cmd[id].kp = self.kp[id.value] - self.msg.motor_cmd[id].kd = self.kd[id.value] - self.msg.motor_cmd[id].q = lowstate.motor_state[id.value].q - - # Initialize remote controller + self.subscribe_thread = None self.remote_controller = self.RemoteController() def _subscribe_motor_state(self): # polls robot state @ 250Hz while not self._shutdown_event.is_set(): start_time = time.time() + + # Step simulation if in simulation mode + if self.config.is_simulation and self.sim_env is not None: + self.sim_env.step() + msg = self.lowstate_subscriber.Read() if msg is not None: lowstate = G1_29_LowState() - # Capture motor states - for id in range(G1_29_Num_Motors): + # Capture motor states using jointindex + for id in G1_29_JointIndex: lowstate.motor_state[id].q = msg.motor_state[id].q lowstate.motor_state[id].dq = msg.motor_state[id].dq lowstate.motor_state[id].tau_est = msg.motor_state[id].tau_est @@ -207,7 +158,7 @@ class UnitreeG1(Robot): # Capture mode_machine lowstate.mode_machine = msg.mode_machine - self.lowstate_buffer.set_data(lowstate) + self._lowstate = lowstate current_time = time.time() all_t_elapsed = current_time - start_time @@ -216,7 +167,7 @@ class UnitreeG1(Robot): @cached_property def action_features(self) -> dict[str, type]: - return {f"{G1_29_JointIndex(motor).name}.pos": float for motor in G1_29_JointIndex} + return {f"{G1_29_JointIndex(motor).name}.q": float for motor in G1_29_JointIndex} def calibrate(self) -> None: # robot is already calibrated pass @@ -225,20 +176,153 @@ class UnitreeG1(Robot): pass def connect(self, calibrate: bool = True) -> None: # connect to DDS + from unitree_sdk2py.idl.default import unitree_hg_msg_dds__LowCmd_ + from unitree_sdk2py.idl.unitree_hg.msg.dds_ import ( + LowCmd_ as hg_LowCmd, + LowState_ as hg_LowState, + ) + from unitree_sdk2py.utils.crc import CRC + + # Initialize DDS channel and simulation environment if self.config.is_simulation: - self.ChannelFactoryInitialize(0, "lo") - self.mujoco_env = make_env("lerobot/unitree-g1-mujoco", trust_remote_code=True) + self._ChannelFactoryInitialize(0, "lo") + self._env_wrapper = make_env("lerobot/unitree-g1-mujoco", trust_remote_code=True) + # Extract the actual gym env from the dict structure + self.sim_env = self._env_wrapper["hub_env"][0].envs[0] else: - self.ChannelFactoryInitialize(0) + self._ChannelFactoryInitialize(0) + + # Initialize direct motor control interface + self.lowcmd_publisher = self._ChannelPublisher(kTopicLowCommand_Debug, hg_LowCmd) + self.lowcmd_publisher.Init() + self.lowstate_subscriber = self._ChannelSubscriber(kTopicLowState, hg_LowState) + self.lowstate_subscriber.Init() + + # Start subscribe thread to read robot state + self.subscribe_thread = threading.Thread(target=self._subscribe_motor_state) + self.subscribe_thread.start() + + # Connect cameras + for cam in self._cameras.values(): + if not cam.is_connected: + cam.connect() + + logger.info(f"Connected {len(self._cameras)} camera(s).") + + # Initialize lowcmd message + self.crc = CRC() + self.msg = unitree_hg_msg_dds__LowCmd_() + self.msg.mode_pr = 0 + + # Wait for first state message to arrive + lowstate = None + while lowstate is None: + lowstate = self._lowstate + if lowstate is None: + time.sleep(0.01) + logger.warning("[UnitreeG1] Waiting for robot state...") + logger.warning("[UnitreeG1] Connected to robot.") + self.msg.mode_machine = lowstate.mode_machine + + # Initialize all motors with unified kp/kd from config + self.kp = np.array(self.config.kp, dtype=np.float32) + self.kd = np.array(self.config.kd, dtype=np.float32) + + for id in G1_29_JointIndex: + self.msg.motor_cmd[id].mode = 1 + self.msg.motor_cmd[id].kp = self.kp[id.value] + self.msg.motor_cmd[id].kd = self.kd[id.value] + self.msg.motor_cmd[id].q = lowstate.motor_state[id.value].q def disconnect(self): + # Signal thread to stop and unblock any waits self._shutdown_event.set() - self.subscribe_thread.join(timeout=2.0) - if self.config.is_simulation: - self.mujoco_env["hub_env"][0].envs[0].kill_sim() + + # Wait for subscribe thread to finish + if self.subscribe_thread is not None: + self.subscribe_thread.join(timeout=2.0) + if self.subscribe_thread.is_alive(): + logger.warning("Subscribe thread did not stop cleanly") + + # Close simulation environment + if self.config.is_simulation and self.sim_env is not None: + try: + # Force-kill the image publish subprocess first to avoid long waits + if hasattr(self.sim_env, "simulator") and hasattr(self.sim_env.simulator, "sim_env"): + sim_env_inner = self.sim_env.simulator.sim_env + if hasattr(sim_env_inner, "image_publish_process"): + proc = sim_env_inner.image_publish_process + if proc.process and proc.process.is_alive(): + logger.info("Force-terminating image publish subprocess...") + proc.stop_event.set() + proc.process.terminate() + proc.process.join(timeout=1) + if proc.process.is_alive(): + proc.process.kill() + self.sim_env.close() + except Exception as e: + logger.warning(f"Error closing sim_env: {e}") + self.sim_env = None + self._env_wrapper = None + + # Disconnect cameras + for cam in self._cameras.values(): + cam.disconnect() def get_observation(self) -> dict[str, Any]: - return self.lowstate_buffer.get_data() + lowstate = self._lowstate + if lowstate is None: + return {} + + obs = {} + + # Motors - q, dq, tau for all joints + for motor in G1_29_JointIndex: + name = motor.name + idx = motor.value + obs[f"{name}.q"] = lowstate.motor_state[idx].q + obs[f"{name}.dq"] = lowstate.motor_state[idx].dq + obs[f"{name}.tau"] = lowstate.motor_state[idx].tau_est + + # IMU - gyroscope + if lowstate.imu_state.gyroscope: + obs["imu.gyro.x"] = lowstate.imu_state.gyroscope[0] + obs["imu.gyro.y"] = lowstate.imu_state.gyroscope[1] + obs["imu.gyro.z"] = lowstate.imu_state.gyroscope[2] + + # IMU - accelerometer + if lowstate.imu_state.accelerometer: + obs["imu.accel.x"] = lowstate.imu_state.accelerometer[0] + obs["imu.accel.y"] = lowstate.imu_state.accelerometer[1] + obs["imu.accel.z"] = lowstate.imu_state.accelerometer[2] + + # IMU - quaternion + if lowstate.imu_state.quaternion: + obs["imu.quat.w"] = lowstate.imu_state.quaternion[0] + obs["imu.quat.x"] = lowstate.imu_state.quaternion[1] + obs["imu.quat.y"] = lowstate.imu_state.quaternion[2] + obs["imu.quat.z"] = lowstate.imu_state.quaternion[3] + + # IMU - rpy + if lowstate.imu_state.rpy: + obs["imu.rpy.roll"] = lowstate.imu_state.rpy[0] + obs["imu.rpy.pitch"] = lowstate.imu_state.rpy[1] + obs["imu.rpy.yaw"] = lowstate.imu_state.rpy[2] + + # Controller - parse wireless_remote and add to obs + if lowstate.wireless_remote and len(lowstate.wireless_remote) >= 24: + self.remote_controller.set(lowstate.wireless_remote) + obs["remote.buttons"] = self.remote_controller.button.copy() + obs["remote.lx"] = self.remote_controller.lx + obs["remote.ly"] = self.remote_controller.ly + obs["remote.rx"] = self.remote_controller.rx + obs["remote.ry"] = self.remote_controller.ry + + # Cameras - read images from ZMQ cameras + for cam_name, cam in self._cameras.items(): + obs[cam_name] = cam.async_read() + + return obs @property def is_calibrated(self) -> bool: @@ -246,11 +330,15 @@ class UnitreeG1(Robot): @property def is_connected(self) -> bool: - return self.lowstate_buffer.get_data() is not None + return self._lowstate is not None @property def _motors_ft(self) -> dict[str, type]: - return {f"{G1_29_JointIndex(motor).name}.pos": float for motor in G1_29_JointIndex} + return {f"{G1_29_JointIndex(motor).name}.q": float for motor in G1_29_JointIndex} + + @property + def cameras(self) -> dict: + return self._cameras @property def _cameras_ft(self) -> dict[str, tuple]: @@ -293,39 +381,51 @@ class UnitreeG1(Robot): self, control_dt: float | None = None, default_positions: list[float] | None = None, - ) -> None: # interpolate to default position + ) -> None: # move robot to default position if control_dt is None: control_dt = self.config.control_dt if default_positions is None: default_positions = np.array(self.config.default_positions, dtype=np.float32) - total_time = 3.0 - num_steps = int(total_time / control_dt) + if self.config.is_simulation and self.sim_env is not None: + self.sim_env.reset() - # get current state - robot_state = self.get_observation() - - # record current positions - init_dof_pos = np.zeros(29, dtype=np.float32) - for i in range(29): - init_dof_pos[i] = robot_state.motor_state[i].q - - # Interpolate to default position - for step in range(num_steps): - start_time = time.time() - - alpha = step / num_steps - action_dict = {} for motor in G1_29_JointIndex: - target_pos = default_positions[motor.value] - interp_pos = init_dof_pos[motor.value] * (1 - alpha) + target_pos * alpha - action_dict[f"{motor.name}.q"] = float(interp_pos) + self.msg.motor_cmd[motor.value].q = default_positions[motor.value] + self.msg.motor_cmd[motor.value].qd = 0 + self.msg.motor_cmd[motor.value].kp = self.kp[motor.value] + self.msg.motor_cmd[motor.value].kd = self.kd[motor.value] + self.msg.motor_cmd[motor.value].tau = 0 + self.msg.crc = self.crc.Crc(self.msg) + self.lowcmd_publisher.Write(self.msg) + else: + total_time = 3.0 + num_steps = int(total_time / control_dt) - self.send_action(action_dict) + # get current state + obs = self.get_observation() - # Maintain constant control rate - elapsed = time.time() - start_time - sleep_time = max(0, control_dt - elapsed) - time.sleep(sleep_time) + # record current positions + init_dof_pos = np.zeros(29, dtype=np.float32) + for motor in G1_29_JointIndex: + init_dof_pos[motor.value] = obs[f"{motor.name}.q"] + + # Interpolate to default position + for step in range(num_steps): + start_time = time.time() + + alpha = step / num_steps + action_dict = {} + for motor in G1_29_JointIndex: + target_pos = default_positions[motor.value] + interp_pos = init_dof_pos[motor.value] * (1 - alpha) + target_pos * alpha + action_dict[f"{motor.name}.q"] = float(interp_pos) + + self.send_action(action_dict) + + # Maintain constant control rate + elapsed = time.time() - start_time + sleep_time = max(0, control_dt - elapsed) + time.sleep(sleep_time) logger.info("Reached default position") diff --git a/src/lerobot/scripts/lerobot_record.py b/src/lerobot/scripts/lerobot_record.py index 8eafa8e6d..9a327d986 100644 --- a/src/lerobot/scripts/lerobot_record.py +++ b/src/lerobot/scripts/lerobot_record.py @@ -74,6 +74,7 @@ from lerobot.cameras import ( # noqa: F401 ) from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig # noqa: F401 from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig # noqa: F401 +from lerobot.cameras.zmq.configuration_zmq import ZMQCameraConfig # noqa: F401 from lerobot.configs import parser from lerobot.configs.policies import PreTrainedConfig from lerobot.datasets.image_writer import safe_stop_image_writer @@ -103,6 +104,7 @@ from lerobot.robots import ( # noqa: F401 make_robot_from_config, omx_follower, so_follower, + unitree_g1, ) from lerobot.teleoperators import ( # noqa: F401 Teleoperator, @@ -508,6 +510,11 @@ def record(cfg: RecordConfig) -> LeRobotDataset: (recorded_episodes < cfg.dataset.num_episodes - 1) or events["rerecord_episode"] ): log_say("Reset the environment", cfg.play_sounds) + + # reset g1 robot + if robot.name == "unitree_g1": + robot.reset() + record_loop( robot=robot, events=events, diff --git a/src/lerobot/scripts/lerobot_replay.py b/src/lerobot/scripts/lerobot_replay.py index 8e0d9cf6d..c16271932 100644 --- a/src/lerobot/scripts/lerobot_replay.py +++ b/src/lerobot/scripts/lerobot_replay.py @@ -60,6 +60,7 @@ from lerobot.robots import ( # noqa: F401 make_robot_from_config, omx_follower, so_follower, + unitree_g1, ) from lerobot.utils.constants import ACTION from lerobot.utils.import_utils import register_third_party_plugins From b8ec1152d447190e01fbe065f25076c730f20f10 Mon Sep 17 00:00:00 2001 From: Steven Palma Date: Mon, 12 Jan 2026 18:05:16 +0100 Subject: [PATCH 2/5] fix(robots): add reachy2 fixes (#2783) * fix(robots): add reachy2 fixes * tests(robots): remove reachy sdk stub --- docs/source/reachy2.mdx | 53 +++-- pyproject.toml | 2 +- .../configuration_reachy2_camera.py | 8 +- .../cameras/reachy2_camera/reachy2_camera.py | 183 +++++------------- .../robots/reachy2/configuration_reachy2.py | 34 ++-- src/lerobot/robots/reachy2/robot_reachy2.py | 10 +- src/lerobot/scripts/lerobot_record.py | 3 + src/lerobot/scripts/lerobot_replay.py | 1 + src/lerobot/scripts/lerobot_teleoperate.py | 2 + .../reachy2_teleoperator.py | 51 +++-- src/lerobot/utils/import_utils.py | 1 + tests/cameras/test_reachy2_camera.py | 14 +- tests/conftest.py | 1 - tests/plugins/reachy2_sdk.py | 46 ----- tests/robots/test_reachy2.py | 2 + 15 files changed, 165 insertions(+), 246 deletions(-) delete mode 100644 tests/plugins/reachy2_sdk.py diff --git a/docs/source/reachy2.mdx b/docs/source/reachy2.mdx index 7d3dc1b60..51b09acd2 100644 --- a/docs/source/reachy2.mdx +++ b/docs/source/reachy2.mdx @@ -38,6 +38,7 @@ docker run --rm -it \ start_rviz:=true start_sdk_server:=true mujoco:=true ``` +> [!NOTE] > If MuJoCo runs slowly (low simulation frequency), append `-e LD_LIBRARY_PATH="/opt/host-libs:$LD_LIBRARY_PATH" \` to the previous command to improve performance: > > ``` @@ -141,7 +142,7 @@ If you choose this option but still want to use the VR teleoperation application First add reachy2 and reachy2_teleoperator to the imports of the record script. Then you can use the following command: ```bash -python -m lerobot.record \ +lerobot-record \ --robot.type=reachy2 \ --robot.ip_address=192.168.0.200 \ --robot.id=r2-0000 \ @@ -150,6 +151,7 @@ python -m lerobot.record \ --teleop.type=reachy2_teleoperator \ --teleop.ip_address=192.168.0.200 \ --teleop.with_mobile_base=false \ + --robot.with_torso_camera=true \ --dataset.repo_id=pollen_robotics/record_test \ --dataset.single_task="Reachy 2 recording test" \ --dataset.num_episodes=1 \ @@ -165,7 +167,7 @@ python -m lerobot.record \ **Extended setup overview (all options included):** ```bash -python -m lerobot.record \ +lerobot-record \ --robot.type=reachy2 \ --robot.ip_address=192.168.0.200 \ --robot.use_external_commands=true \ @@ -177,6 +179,8 @@ python -m lerobot.record \ --robot.with_left_teleop_camera=true \ --robot.with_right_teleop_camera=true \ --robot.with_torso_camera=false \ + --robot.camera_width=640 \ + --robot.camera_height=480 \ --robot.disable_torque_on_disconnect=false \ --robot.max_relative_target=5.0 \ --teleop.type=reachy2_teleoperator \ @@ -212,9 +216,10 @@ Must be set to true if a compliant Reachy 2 is used to control another one. From our initial tests, recording **all** joints when only some are moving can reduce model quality with certain policies. To avoid this, you can exclude specific parts from recording and replay using: -```` +```bash --robot.with_=false -```, +``` + with `` being one of : `mobile_base`, `l_arm`, `r_arm", `neck`, `antennas`. It determine whether the corresponding part is recorded in the observations. True if not set. @@ -222,49 +227,60 @@ By default, **all parts are recorded**. The same per-part mechanism is available in `reachy2_teleoperator` as well. -```` - +```bash --teleop.with\_ - ``` + with `` being one of : `mobile_base`, `l_arm`, `r_arm", `neck`, `antennas`. Determine whether the corresponding part is recorded in the actions. True if not set. > **Important:** In a given session, the **enabled parts must match** on both the robot and the teleoperator. -For example, if the robot runs with `--robot.with_mobile_base=false`, the teleoperator must disable the same part `--teleoperator.with_mobile_base=false`. +> For example, if the robot runs with `--robot.with_mobile_base=false`, the teleoperator must disable the same part `--teleoperator.with_mobile_base=false`. ##### Use the relevant cameras -You can do the same for **cameras**. By default, only the **teleoperation cameras** are recorded (both `left_teleop_camera` and `right_teleop_camera`). Enable or disable each camera with: +You can do the same for **cameras**. Enable or disable each camera with default parameters using: +```bash +--robot.with_left_teleop_camera= \ +--robot.with_right_teleop_camera= \ +--robot.with_torso_camera= ``` ---robot.with_left_teleop_camera= ---robot.with_right_teleop_camera= ---robot.with_torso_camera= +By default, no camera is recorded, all camera arguments are set to `false`. +If you want to, you can use custom `width` and `height` parameters for Reachy 2's cameras using the `--robot.camera_width` & `--robot.camera_height` argument: -```` +```bash +--robot.camera_width=1920 \ +--robot.camera_height=1080 +``` +This will change the resolution of all 3 default robot cameras (enabled by the above bool arguments). + +If you want, you can add additional cameras other than the ones in the robot as usual with: + +```bash +--robot.cameras="{ extra: {type: opencv, index_or_path: 42, width: 640, height: 480, fps: 30}}" \ +``` ## Step 2: Replay Make sure the robot is configured with the same parts as the dataset: ```bash -python -m lerobot.replay \ +lerobot-replay \ --robot.type=reachy2 \ --robot.ip_address=192.168.0.200 \ --robot.use_external_commands=false \ --robot.with_mobile_base=false \ --dataset.repo_id=pollen_robotics/record_test \ --dataset.episode=0 - --display_data=true -```` +``` ## Step 3: Train ```bash -python -m lerobot.scripts.train \ +lerobot-train \ --dataset.repo_id=pollen_robotics/record_test \ --policy.type=act \ --output_dir=outputs/train/reachy2_test \ @@ -277,10 +293,9 @@ python -m lerobot.scripts.train \ ## Step 4: Evaluate ```bash -python -m lerobot.record \ +lerobot-eval \ --robot.type=reachy2 \ --robot.ip_address=192.168.0.200 \ - --display_data=false \ --dataset.repo_id=pollen_robotics/eval_record_test \ --dataset.single_task="Evaluate reachy2 policy" \ --dataset.num_episodes=10 \ diff --git a/pyproject.toml b/pyproject.toml index e8f334c77..067cb1df8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -111,7 +111,7 @@ unitree_g1 = [ "pyzmq>=26.2.1,<28.0.0", "onnxruntime>=1.16.0,<2.0.0" ] -reachy2 = ["reachy2_sdk>=1.0.14,<1.1.0"] +reachy2 = ["reachy2_sdk>=1.0.15,<1.1.0"] kinematics = ["lerobot[placo-dep]"] intelrealsense = [ "pyrealsense2>=2.55.1.6486,<2.57.0 ; sys_platform != 'darwin'", diff --git a/src/lerobot/cameras/reachy2_camera/configuration_reachy2_camera.py b/src/lerobot/cameras/reachy2_camera/configuration_reachy2_camera.py index f26cf2ad1..ca6db4f03 100644 --- a/src/lerobot/cameras/reachy2_camera/configuration_reachy2_camera.py +++ b/src/lerobot/cameras/reachy2_camera/configuration_reachy2_camera.py @@ -35,18 +35,19 @@ class Reachy2CameraConfig(CameraConfig): name="teleop", image_type="left", ip_address="192.168.0.200", # IP address of the robot - fps=15, + port=50065, # Port of the camera server width=640, height=480, + fps=30, # Not configurable for Reachy 2 cameras color_mode=ColorMode.RGB, - ) # Left teleop camera, 640x480 @ 15FPS + ) # Left teleop camera, 640x480 @ 30FPS ``` Attributes: name: Name of the camera device. Can be "teleop" or "depth". image_type: Type of image stream. For "teleop" camera, can be "left" or "right". For "depth" camera, can be "rgb" or "depth". (depth is not supported yet) - fps: Requested frames per second for the color stream. + fps: Requested frames per second for the color stream. Not configurable for Reachy 2 cameras. width: Requested frame width in pixels for the color stream. height: Requested frame height in pixels for the color stream. color_mode: Color mode for image output (RGB or BGR). Defaults to RGB. @@ -62,7 +63,6 @@ class Reachy2CameraConfig(CameraConfig): color_mode: ColorMode = ColorMode.RGB ip_address: str | None = "localhost" port: int = 50065 - # use_depth: bool = False def __post_init__(self) -> None: if self.name not in ["teleop", "depth"]: diff --git a/src/lerobot/cameras/reachy2_camera/reachy2_camera.py b/src/lerobot/cameras/reachy2_camera/reachy2_camera.py index 30e096767..c8916c5ee 100644 --- a/src/lerobot/cameras/reachy2_camera/reachy2_camera.py +++ b/src/lerobot/cameras/reachy2_camera/reachy2_camera.py @@ -16,12 +16,13 @@ Provides the Reachy2Camera class for capturing frames from Reachy 2 cameras using Reachy 2's CameraManager. """ +from __future__ import annotations + import logging import os import platform import time -from threading import Event, Lock, Thread -from typing import Any +from typing import TYPE_CHECKING, Any from numpy.typing import NDArray # type: ignore # TODO: add type stubs for numpy.typing @@ -30,10 +31,19 @@ if platform.system() == "Windows" and "OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS" os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0" import cv2 # type: ignore # TODO: add type stubs for OpenCV import numpy as np # type: ignore # TODO: add type stubs for numpy -from reachy2_sdk.media.camera import CameraView # type: ignore # TODO: add type stubs for reachy2_sdk -from reachy2_sdk.media.camera_manager import ( # type: ignore # TODO: add type stubs for reachy2_sdk - CameraManager, -) + +from lerobot.utils.import_utils import _reachy2_sdk_available + +if TYPE_CHECKING or _reachy2_sdk_available: + from reachy2_sdk.media.camera import CameraView + from reachy2_sdk.media.camera_manager import CameraManager +else: + CameraManager = None + + class CameraView: + LEFT = 0 + RIGHT = 1 + from lerobot.utils.errors import DeviceNotConnectedError @@ -69,17 +79,10 @@ class Reachy2Camera(Camera): self.config = config - self.fps = config.fps self.color_mode = config.color_mode self.cam_manager: CameraManager | None = None - self.thread: Thread | None = None - self.stop_event: Event | None = None - self.frame_lock: Lock = Lock() - self.latest_frame: NDArray[Any] | None = None - self.new_frame_event: Event = Event() - def __str__(self) -> str: return f"{self.__class__.__name__}({self.config.name}, {self.config.image_type})" @@ -100,44 +103,23 @@ class Reachy2Camera(Camera): def connect(self, warmup: bool = True) -> None: """ Connects to the Reachy2 CameraManager as specified in the configuration. + + Raises: + DeviceNotConnectedError: If the camera is not connected. """ self.cam_manager = CameraManager(host=self.config.ip_address, port=self.config.port) + if self.cam_manager is None: + raise DeviceNotConnectedError(f"Could not connect to {self}.") self.cam_manager.initialize_cameras() logger.info(f"{self} connected.") @staticmethod - def find_cameras(ip_address: str = "localhost", port: int = 50065) -> list[dict[str, Any]]: + def find_cameras() -> list[dict[str, Any]]: """ - Detects available Reachy 2 cameras. - - Returns: - List[Dict[str, Any]]: A list of dictionaries, - where each dictionary contains 'name', 'stereo', - and the default profile properties (width, height, fps). + Detection not implemented for Reachy2 cameras. """ - initialized_cameras = [] - camera_manager = CameraManager(host=ip_address, port=port) - - for camera in [camera_manager.teleop, camera_manager.depth]: - if camera is None: - continue - - height, width, _, _, _, _, _ = camera.get_parameters() - - camera_info = { - "name": camera._cam_info.name, - "stereo": camera._cam_info.stereo, - "default_profile": { - "width": width, - "height": height, - "fps": 30, - }, - } - initialized_cameras.append(camera_info) - - camera_manager.disconnect() - return initialized_cameras + raise NotImplementedError("Camera detection is not implemented for Reachy2 cameras.") def read(self, color_mode: ColorMode | None = None) -> NDArray[Any]: """ @@ -155,95 +137,49 @@ class Reachy2Camera(Camera): (height, width, channels), using the specified or default color mode and applying any configured rotation. """ + start_time = time.perf_counter() + if not self.is_connected: raise DeviceNotConnectedError(f"{self} is not connected.") - start_time = time.perf_counter() + if self.cam_manager is None: + raise DeviceNotConnectedError(f"{self} is not connected.") frame: NDArray[Any] = np.empty((0, 0, 3), dtype=np.uint8) - if self.cam_manager is None: - raise DeviceNotConnectedError(f"{self} is not connected.") + if self.config.name == "teleop" and hasattr(self.cam_manager, "teleop"): + if self.config.image_type == "left": + frame = self.cam_manager.teleop.get_frame( + CameraView.LEFT, size=(self.config.width, self.config.height) + )[0] + elif self.config.image_type == "right": + frame = self.cam_manager.teleop.get_frame( + CameraView.RIGHT, size=(self.config.width, self.config.height) + )[0] + elif self.config.name == "depth" and hasattr(self.cam_manager, "depth"): + if self.config.image_type == "depth": + frame = self.cam_manager.depth.get_depth_frame()[0] + elif self.config.image_type == "rgb": + frame = self.cam_manager.depth.get_frame(size=(self.config.width, self.config.height))[0] else: - if self.config.name == "teleop" and hasattr(self.cam_manager, "teleop"): - if self.config.image_type == "left": - frame = self.cam_manager.teleop.get_frame(CameraView.LEFT, size=(640, 480))[0] - elif self.config.image_type == "right": - frame = self.cam_manager.teleop.get_frame(CameraView.RIGHT, size=(640, 480))[0] - elif self.config.name == "depth" and hasattr(self.cam_manager, "depth"): - if self.config.image_type == "depth": - frame = self.cam_manager.depth.get_depth_frame()[0] - elif self.config.image_type == "rgb": - frame = self.cam_manager.depth.get_frame(size=(640, 480))[0] + raise ValueError(f"Invalid camera name '{self.config.name}'. Expected 'teleop' or 'depth'.") - if frame is None: - return np.empty((0, 0, 3), dtype=np.uint8) + if frame is None: + return np.empty((0, 0, 3), dtype=np.uint8) - if self.config.color_mode == "rgb": - frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + if self.config.color_mode == "rgb": + frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) read_duration_ms = (time.perf_counter() - start_time) * 1e3 logger.debug(f"{self} read took: {read_duration_ms:.1f}ms") return frame - def _read_loop(self) -> None: - """ - Internal loop run by the background thread for asynchronous reading. - - On each iteration: - 1. Reads a color frame - 2. Stores result in latest_frame (thread-safe) - 3. Sets new_frame_event to notify listeners - - Stops on DeviceNotConnectedError, logs other errors and continues. - """ - if self.stop_event is None: - raise RuntimeError(f"{self}: stop_event is not initialized before starting read loop.") - - while not self.stop_event.is_set(): - try: - color_image = self.read() - - with self.frame_lock: - self.latest_frame = color_image - self.new_frame_event.set() - - except DeviceNotConnectedError: - break - except Exception as e: - logger.warning(f"Error reading frame in background thread for {self}: {e}") - - def _start_read_thread(self) -> None: - """Starts or restarts the background read thread if it's not running.""" - if self.thread is not None and self.thread.is_alive(): - self.thread.join(timeout=0.1) - if self.stop_event is not None: - self.stop_event.set() - - self.stop_event = Event() - self.thread = Thread(target=self._read_loop, args=(), name=f"{self}_read_loop") - self.thread.daemon = True - self.thread.start() - - def _stop_read_thread(self) -> None: - """Signals the background read thread to stop and waits for it to join.""" - if self.stop_event is not None: - self.stop_event.set() - - if self.thread is not None and self.thread.is_alive(): - self.thread.join(timeout=2.0) - - self.thread = None - self.stop_event = None - def async_read(self, timeout_ms: float = 200) -> NDArray[Any]: """ - Reads the latest available frame asynchronously. + Reads the latest available frame. - This method retrieves the most recent frame captured by the background - read thread. It does not block waiting for the camera hardware directly, - but may wait up to timeout_ms for the background thread to provide a frame. + This method retrieves the most recent frame available in Reachy 2's low-level software. Args: timeout_ms (float): Maximum time in milliseconds to wait for a frame @@ -261,22 +197,10 @@ class Reachy2Camera(Camera): if not self.is_connected: raise DeviceNotConnectedError(f"{self} is not connected.") - if self.thread is None or not self.thread.is_alive(): - self._start_read_thread() - - if not self.new_frame_event.wait(timeout=timeout_ms / 1000.0): - thread_alive = self.thread is not None and self.thread.is_alive() - raise TimeoutError( - f"Timed out waiting for frame from camera {self} after {timeout_ms} ms. " - f"Read thread alive: {thread_alive}." - ) - - with self.frame_lock: - frame = self.latest_frame - self.new_frame_event.clear() + frame = self.read() if frame is None: - raise RuntimeError(f"Internal error: Event set but no frame available for {self}.") + raise RuntimeError(f"Internal error: No frame available for {self}.") return frame @@ -287,12 +211,9 @@ class Reachy2Camera(Camera): Raises: DeviceNotConnectedError: If the camera is already disconnected. """ - if not self.is_connected and self.thread is None: + if not self.is_connected: raise DeviceNotConnectedError(f"{self} not connected.") - if self.thread is not None: - self._stop_read_thread() - if self.cam_manager is not None: self.cam_manager.disconnect() diff --git a/src/lerobot/robots/reachy2/configuration_reachy2.py b/src/lerobot/robots/reachy2/configuration_reachy2.py index aa25351c6..63293e675 100644 --- a/src/lerobot/robots/reachy2/configuration_reachy2.py +++ b/src/lerobot/robots/reachy2/configuration_reachy2.py @@ -30,6 +30,8 @@ class Reachy2RobotConfig(RobotConfig): # IP address of the Reachy 2 robot ip_address: str | None = "localhost" + # Port of the Reachy 2 robot + port: int = 50065 # If True, turn_off_smoothly() will be sent to the robot before disconnecting. disable_torque_on_disconnect: bool = False @@ -51,11 +53,16 @@ class Reachy2RobotConfig(RobotConfig): # Robot cameras # Set to True if you want to use the corresponding cameras in the observations. - # By default, only the teleop cameras are used. - with_left_teleop_camera: bool = True - with_right_teleop_camera: bool = True + # By default, no camera is used. + with_left_teleop_camera: bool = False + with_right_teleop_camera: bool = False with_torso_camera: bool = False + # Camera parameters + camera_width: int = 640 + camera_height: int = 480 + + # For cameras other than the 3 default Reachy 2 cameras. cameras: dict[str, CameraConfig] = field(default_factory=dict) def __post_init__(self) -> None: @@ -65,9 +72,10 @@ class Reachy2RobotConfig(RobotConfig): name="teleop", image_type="left", ip_address=self.ip_address, - fps=15, - width=640, - height=480, + port=self.port, + width=self.camera_width, + height=self.camera_height, + fps=30, # Not configurable for Reachy 2 cameras color_mode=ColorMode.RGB, ) if self.with_right_teleop_camera: @@ -75,9 +83,10 @@ class Reachy2RobotConfig(RobotConfig): name="teleop", image_type="right", ip_address=self.ip_address, - fps=15, - width=640, - height=480, + port=self.port, + width=self.camera_width, + height=self.camera_height, + fps=30, # Not configurable for Reachy 2 cameras color_mode=ColorMode.RGB, ) if self.with_torso_camera: @@ -85,9 +94,10 @@ class Reachy2RobotConfig(RobotConfig): name="depth", image_type="rgb", ip_address=self.ip_address, - fps=15, - width=640, - height=480, + port=self.port, + width=self.camera_width, + height=self.camera_height, + fps=30, # Not configurable for Reachy 2 cameras color_mode=ColorMode.RGB, ) diff --git a/src/lerobot/robots/reachy2/robot_reachy2.py b/src/lerobot/robots/reachy2/robot_reachy2.py index ecc488a79..74742ee8d 100644 --- a/src/lerobot/robots/reachy2/robot_reachy2.py +++ b/src/lerobot/robots/reachy2/robot_reachy2.py @@ -13,19 +13,25 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations import time -from typing import Any +from typing import TYPE_CHECKING, Any import numpy as np -from reachy2_sdk import ReachySDK from lerobot.cameras.utils import make_cameras_from_configs +from lerobot.utils.import_utils import _reachy2_sdk_available from ..robot import Robot from ..utils import ensure_safe_goal_position from .configuration_reachy2 import Reachy2RobotConfig +if TYPE_CHECKING or _reachy2_sdk_available: + from reachy2_sdk import ReachySDK +else: + ReachySDK = None + # {lerobot_keys: reachy2_sdk_keys} REACHY2_NECK_JOINTS = { "neck_yaw.pos": "head.neck.yaw", diff --git a/src/lerobot/scripts/lerobot_record.py b/src/lerobot/scripts/lerobot_record.py index 9a327d986..f03776989 100644 --- a/src/lerobot/scripts/lerobot_record.py +++ b/src/lerobot/scripts/lerobot_record.py @@ -73,6 +73,7 @@ from lerobot.cameras import ( # noqa: F401 CameraConfig, # noqa: F401 ) from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig # noqa: F401 +from lerobot.cameras.reachy2_camera.configuration_reachy2_camera import Reachy2CameraConfig # noqa: F401 from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig # noqa: F401 from lerobot.cameras.zmq.configuration_zmq import ZMQCameraConfig # noqa: F401 from lerobot.configs import parser @@ -103,6 +104,7 @@ from lerobot.robots import ( # noqa: F401 koch_follower, make_robot_from_config, omx_follower, + reachy2, so_follower, unitree_g1, ) @@ -114,6 +116,7 @@ from lerobot.teleoperators import ( # noqa: F401 koch_leader, make_teleoperator_from_config, omx_leader, + reachy2_teleoperator, so_leader, ) from lerobot.teleoperators.keyboard.teleop_keyboard import KeyboardTeleop diff --git a/src/lerobot/scripts/lerobot_replay.py b/src/lerobot/scripts/lerobot_replay.py index c16271932..49c06d643 100644 --- a/src/lerobot/scripts/lerobot_replay.py +++ b/src/lerobot/scripts/lerobot_replay.py @@ -59,6 +59,7 @@ from lerobot.robots import ( # noqa: F401 koch_follower, make_robot_from_config, omx_follower, + reachy2, so_follower, unitree_g1, ) diff --git a/src/lerobot/scripts/lerobot_teleoperate.py b/src/lerobot/scripts/lerobot_teleoperate.py index f2392bb51..05d4534d4 100644 --- a/src/lerobot/scripts/lerobot_teleoperate.py +++ b/src/lerobot/scripts/lerobot_teleoperate.py @@ -76,6 +76,7 @@ from lerobot.robots import ( # noqa: F401 koch_follower, make_robot_from_config, omx_follower, + reachy2, so_follower, ) from lerobot.teleoperators import ( # noqa: F401 @@ -88,6 +89,7 @@ from lerobot.teleoperators import ( # noqa: F401 koch_leader, make_teleoperator_from_config, omx_leader, + reachy2_teleoperator, so_leader, ) from lerobot.utils.import_utils import register_third_party_plugins diff --git a/src/lerobot/teleoperators/reachy2_teleoperator/reachy2_teleoperator.py b/src/lerobot/teleoperators/reachy2_teleoperator/reachy2_teleoperator.py index 5a427dd71..578aaa7b2 100644 --- a/src/lerobot/teleoperators/reachy2_teleoperator/reachy2_teleoperator.py +++ b/src/lerobot/teleoperators/reachy2_teleoperator/reachy2_teleoperator.py @@ -13,11 +13,20 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations import logging import time +from typing import TYPE_CHECKING -from reachy2_sdk import ReachySDK +from lerobot.utils.import_utils import _reachy2_sdk_available + +if TYPE_CHECKING or _reachy2_sdk_available: + from reachy2_sdk import ReachySDK +else: + ReachySDK = None + +from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError from ..teleoperator import Teleoperator from .config_reachy2_teleoperator import Reachy2TeleoperatorConfig @@ -75,6 +84,7 @@ class Reachy2Teleoperator(Teleoperator): def __init__(self, config: Reachy2TeleoperatorConfig): super().__init__(config) + self.config = config self.reachy: None | ReachySDK = None @@ -117,9 +127,13 @@ class Reachy2Teleoperator(Teleoperator): return self.reachy.is_connected() if self.reachy is not None else False def connect(self, calibrate: bool = True) -> None: + if self.is_connected: + raise DeviceAlreadyConnectedError(f"{self} already connected") + self.reachy = ReachySDK(self.config.ip_address) + if not self.is_connected: - raise ConnectionError() + raise DeviceNotConnectedError() logger.info(f"{self} connected.") @property @@ -135,23 +149,24 @@ class Reachy2Teleoperator(Teleoperator): def get_action(self) -> dict[str, float]: start = time.perf_counter() - if self.reachy and self.is_connected: - if self.config.use_present_position: - joint_action = { - k: self.reachy.joints[v].present_position for k, v in self.joints_dict.items() - } - else: - joint_action = {k: self.reachy.joints[v].goal_position for k, v in self.joints_dict.items()} + if not self.is_connected: + raise DeviceNotConnectedError(f"{self} is not connected.") - if not self.config.with_mobile_base: - dt_ms = (time.perf_counter() - start) * 1e3 - logger.debug(f"{self} read action: {dt_ms:.1f}ms") - return joint_action + joint_action: dict[str, float] = {} + vel_action: dict[str, float] = {} - if self.config.use_present_position: - vel_action = {k: self.reachy.mobile_base.odometry[v] for k, v in REACHY2_VEL.items()} - else: - vel_action = {k: self.reachy.mobile_base.last_cmd_vel[v] for k, v in REACHY2_VEL.items()} + if self.config.use_present_position: + joint_action = {k: self.reachy.joints[v].present_position for k, v in self.joints_dict.items()} + else: + joint_action = {k: self.reachy.joints[v].goal_position for k, v in self.joints_dict.items()} + if not self.config.with_mobile_base: + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read action: {dt_ms:.1f}ms") + return joint_action + if self.config.use_present_position: + vel_action = {k: self.reachy.mobile_base.odometry[v] for k, v in REACHY2_VEL.items()} + else: + vel_action = {k: self.reachy.mobile_base.last_cmd_vel[v] for k, v in REACHY2_VEL.items()} dt_ms = (time.perf_counter() - start) * 1e3 logger.debug(f"{self} read action: {dt_ms:.1f}ms") return {**joint_action, **vel_action} @@ -160,5 +175,5 @@ class Reachy2Teleoperator(Teleoperator): raise NotImplementedError def disconnect(self) -> None: - if self.reachy and self.is_connected: + if self.is_connected: self.reachy.disconnect() diff --git a/src/lerobot/utils/import_utils.py b/src/lerobot/utils/import_utils.py index e6817ba6c..0206a8ac9 100644 --- a/src/lerobot/utils/import_utils.py +++ b/src/lerobot/utils/import_utils.py @@ -64,6 +64,7 @@ def is_package_available(pkg_name: str, return_version: bool = False) -> tuple[b _transformers_available = is_package_available("transformers") _peft_available = is_package_available("peft") _scipy_available = is_package_available("scipy") +_reachy2_sdk_available = is_package_available("reachy2_sdk") def make_device_from_device_class(config: ChoiceRegistry) -> Any: diff --git a/tests/cameras/test_reachy2_camera.py b/tests/cameras/test_reachy2_camera.py index 0b38e8b0b..14774bf38 100644 --- a/tests/cameras/test_reachy2_camera.py +++ b/tests/cameras/test_reachy2_camera.py @@ -20,6 +20,8 @@ from unittest.mock import MagicMock, patch import numpy as np import pytest +pytest.importorskip("reachy2_sdk") + from lerobot.cameras.reachy2_camera import Reachy2Camera, Reachy2CameraConfig from lerobot.utils.errors import DeviceNotConnectedError @@ -127,24 +129,12 @@ def test_async_read(camera): try: img = camera.async_read() - assert camera.thread is not None - assert camera.thread.is_alive() assert isinstance(img, np.ndarray) finally: if camera.is_connected: camera.disconnect() -def test_async_read_timeout(camera): - camera.connect() - try: - with pytest.raises(TimeoutError): - camera.async_read(timeout_ms=0) - finally: - if camera.is_connected: - camera.disconnect() - - def test_read_before_connect(camera): with pytest.raises(DeviceNotConnectedError): _ = camera.read() diff --git a/tests/conftest.py b/tests/conftest.py index b14e9aed5..2fcf878ab 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -28,7 +28,6 @@ pytest_plugins = [ "tests.fixtures.files", "tests.fixtures.hub", "tests.fixtures.optimizers", - "tests.plugins.reachy2_sdk", ] diff --git a/tests/plugins/reachy2_sdk.py b/tests/plugins/reachy2_sdk.py deleted file mode 100644 index 457fcf0f9..000000000 --- a/tests/plugins/reachy2_sdk.py +++ /dev/null @@ -1,46 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2025 The HuggingFace Inc. team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import sys -import types -from unittest.mock import MagicMock - - -def _install_reachy2_sdk_stub(): - sdk = types.ModuleType("reachy2_sdk") - sdk.__path__ = [] - sdk.ReachySDK = MagicMock(name="ReachySDK") - - media = types.ModuleType("reachy2_sdk.media") - media.__path__ = [] - camera = types.ModuleType("reachy2_sdk.media.camera") - camera.CameraView = MagicMock(name="CameraView") - camera_manager = types.ModuleType("reachy2_sdk.media.camera_manager") - camera_manager.CameraManager = MagicMock(name="CameraManager") - - sdk.media = media - media.camera = camera - media.camera_manager = camera_manager - - # Register in sys.modules - sys.modules.setdefault("reachy2_sdk", sdk) - sys.modules.setdefault("reachy2_sdk.media", media) - sys.modules.setdefault("reachy2_sdk.media.camera", camera) - sys.modules.setdefault("reachy2_sdk.media.camera_manager", camera_manager) - - -def pytest_sessionstart(session): - _install_reachy2_sdk_stub() diff --git a/tests/robots/test_reachy2.py b/tests/robots/test_reachy2.py index 94152ea38..d3c44bf5a 100644 --- a/tests/robots/test_reachy2.py +++ b/tests/robots/test_reachy2.py @@ -19,6 +19,8 @@ from unittest.mock import MagicMock, patch import numpy as np import pytest +pytest.importorskip("reachy2_sdk") + from lerobot.robots.reachy2 import ( REACHY2_ANTENNAS_JOINTS, REACHY2_L_ARM_JOINTS, From d0f57f58d1e24688664863040ef24cb8a1b37374 Mon Sep 17 00:00:00 2001 From: samet-rob Date: Mon, 12 Jan 2026 21:24:19 +0300 Subject: [PATCH 3/5] Move cfg.validate() earlier to fix NoneType error with --policy.path (#2782) --- src/lerobot/scripts/lerobot_train.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index 927e9659a..55737d5d8 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -251,6 +251,8 @@ def train(cfg: TrainPipelineConfig, accelerator: Accelerator | None = None): cfg: A `TrainPipelineConfig` object containing all training configurations. accelerator: Optional Accelerator instance. If None, one will be created automatically. """ + cfg.validate() + # Create Accelerator if not provided # It will automatically detect if running in distributed mode or single-process mode # We set step_scheduler_with_optimizer=False to prevent accelerate from adjusting the lr_scheduler steps based on the num_processes @@ -274,8 +276,6 @@ def train(cfg: TrainPipelineConfig, accelerator: Accelerator | None = None): # When using accelerate, only the main process should log to avoid duplicate outputs is_main_process = accelerator.is_main_process - cfg.validate() - # Only log on main process if is_main_process: logging.info(pformat(cfg.to_dict())) From 2cdd9f43f7e104702ed74c9b6e809d079527314e Mon Sep 17 00:00:00 2001 From: Jade Choghari Date: Tue, 13 Jan 2026 01:42:53 +0100 Subject: [PATCH 4/5] fix: train tokenizer CLI entry point (#2784) --- .../scripts/lerobot_train_tokenizer.py | 177 +++++++++++------- 1 file changed, 105 insertions(+), 72 deletions(-) diff --git a/src/lerobot/scripts/lerobot_train_tokenizer.py b/src/lerobot/scripts/lerobot_train_tokenizer.py index 03bfcaaf8..296447bad 100644 --- a/src/lerobot/scripts/lerobot_train_tokenizer.py +++ b/src/lerobot/scripts/lerobot_train_tokenizer.py @@ -14,11 +14,14 @@ """Train FAST tokenizer for action encoding. This script: -1. Loads action chunks from LeRobotDataset (with sampling) -2. Applies delta transforms and per-timestamp normalization -3. Trains FAST tokenizer on specified action dimensions -4. Saves tokenizer to assets directory -5. Reports compression statistics +1. Loads action chunks from LeRobotDataset (with episode sampling) +2. Optionally applies delta transforms (relative vs absolute actions) +3. Extracts specified action dimensions for encoding +4. Applies normalization (MEAN_STD, MIN_MAX, QUANTILES, or other modes) +5. Trains FAST tokenizer (BPE on DCT coefficients) on the action chunks +6. Saves tokenizer to output directory +7. Optionally pushes tokenizer to Hugging Face Hub +8. Reports compression statistics Example: @@ -42,18 +45,64 @@ lerobot-train-tokenizer \ """ import json +from dataclasses import dataclass from pathlib import Path +from typing import TYPE_CHECKING import numpy as np import torch -import tyro from huggingface_hub import HfApi -from transformers import AutoProcessor +from lerobot.utils.import_utils import _transformers_available + +if TYPE_CHECKING or _transformers_available: + from transformers import AutoProcessor +else: + AutoProcessor = None + +from lerobot.configs import parser from lerobot.configs.types import NormalizationMode from lerobot.datasets.lerobot_dataset import LeRobotDataset +@dataclass +class TokenizerTrainingConfig: + """Configuration for training FAST tokenizer.""" + + # LeRobot dataset repository ID + repo_id: str + # Root directory for dataset (default: ~/.cache/huggingface/lerobot) + root: str | None = None + # Number of future actions in each chunk + action_horizon: int = 10 + # Max episodes to use (None = all episodes in dataset) + max_episodes: int | None = None + # Fraction of chunks to sample per episode + sample_fraction: float = 0.1 + # Comma-separated dimension ranges to encode (e.g., "0:6,7:23") + encoded_dims: str = "0:6,7:23" + # Comma-separated dimension indices for delta transform (e.g., "0,1,2,3,4,5") + delta_dims: str | None = None + # Whether to apply delta transform (relative actions vs absolute actions) + use_delta_transform: bool = False + # Dataset key for state observations (default: "observation.state") + state_key: str = "observation.state" + # Normalization mode (MEAN_STD, MIN_MAX, QUANTILES, QUANTILE10, IDENTITY) + normalization_mode: str = "QUANTILES" + # FAST vocabulary size (BPE vocab size) + vocab_size: int = 1024 + # DCT scaling factor (default: 10.0) + scale: float = 10.0 + # Directory to save tokenizer (default: ./fast_tokenizer_{repo_id}) + output_dir: str | None = None + # Whether to push the tokenizer to Hugging Face Hub + push_to_hub: bool = False + # Hub repository ID (e.g., "username/tokenizer-name"). If None, uses output_dir name + hub_repo_id: str | None = None + # Whether to create a private repository on the Hub + hub_private: bool = False + + def apply_delta_transform(state: np.ndarray, actions: np.ndarray, delta_dims: list[int] | None) -> np.ndarray: """Apply delta transform to specified dimensions. @@ -327,88 +376,57 @@ def compute_compression_stats(tokenizer, action_chunks: np.ndarray): return stats -def main( - repo_id: str, - root: str | None = None, - action_horizon: int = 10, - max_episodes: int | None = None, - sample_fraction: float = 0.1, - encoded_dims: str = "0:6,7:23", - delta_dims: str | None = None, - use_delta_transform: bool = False, - state_key: str = "observation.state", - normalization_mode: str = "QUANTILES", - vocab_size: int = 1024, - scale: float = 10.0, - output_dir: str | None = None, - push_to_hub: bool = False, - hub_repo_id: str | None = None, - hub_private: bool = False, -): +@parser.wrap() +def train_tokenizer(cfg: TokenizerTrainingConfig): """ Train FAST tokenizer for action encoding. Args: - repo_id: LeRobot dataset repository ID - root: Root directory for dataset (default: ~/.cache/huggingface/lerobot) - action_horizon: Number of future actions in each chunk - max_episodes: Max episodes to use (None = all episodes in dataset) - sample_fraction: Fraction of chunks to sample per episode - encoded_dims: Comma-separated dimension ranges to encode (e.g., "0:6,7:23") - delta_dims: Comma-separated dimension indices for delta transform (e.g., "0,1,2,3,4,5") - use_delta_transform: Whether to apply delta transform (relative actions vs absolute actions) - state_key: Dataset key for state observations (default: "observation.state") - normalization_mode: Normalization mode (MEAN_STD, MIN_MAX, QUANTILES, QUANTILE10, IDENTITY) - vocab_size: FAST vocabulary size (BPE vocab size) - scale: DCT scaling factor (default: 10.0) - output_dir: Directory to save tokenizer (default: ./fast_tokenizer_{repo_id}) - push_to_hub: Whether to push the tokenizer to Hugging Face Hub - hub_repo_id: Hub repository ID (e.g., "username/tokenizer-name"). If None, uses output_dir name - hub_private: Whether to create a private repository on the Hub + cfg: TokenizerTrainingConfig dataclass with all configuration parameters """ # load dataset - print(f"Loading dataset: {repo_id}") - dataset = LeRobotDataset(repo_id=repo_id, root=root) + print(f"Loading dataset: {cfg.repo_id}") + dataset = LeRobotDataset(repo_id=cfg.repo_id, root=cfg.root) print(f"Dataset loaded: {dataset.num_episodes} episodes, {dataset.num_frames} frames") # parse normalization mode try: - norm_mode = NormalizationMode(normalization_mode) + norm_mode = NormalizationMode(cfg.normalization_mode) except ValueError as err: raise ValueError( - f"Invalid normalization_mode: {normalization_mode}. " + f"Invalid normalization_mode: {cfg.normalization_mode}. " f"Must be one of: {', '.join([m.value for m in NormalizationMode])}" ) from err print(f"Normalization mode: {norm_mode.value}") # parse encoded dimensions encoded_dim_ranges = [] - for range_str in encoded_dims.split(","): + for range_str in cfg.encoded_dims.split(","): start, end = map(int, range_str.strip().split(":")) encoded_dim_ranges.append((start, end)) total_encoded_dims = sum(end - start for start, end in encoded_dim_ranges) - print(f"Encoding {total_encoded_dims} dimensions: {encoded_dims}") + print(f"Encoding {total_encoded_dims} dimensions: {cfg.encoded_dims}") # parse delta dimensions delta_dim_list = None - if delta_dims is not None and delta_dims.strip(): - delta_dim_list = [int(d.strip()) for d in delta_dims.split(",")] + if cfg.delta_dims is not None and cfg.delta_dims.strip(): + delta_dim_list = [int(d.strip()) for d in cfg.delta_dims.split(",")] print(f"Delta dimensions: {delta_dim_list}") else: print("No delta dimensions specified") - print(f"Use delta transform: {use_delta_transform}") - if use_delta_transform and (delta_dim_list is None or len(delta_dim_list) == 0): + print(f"Use delta transform: {cfg.use_delta_transform}") + if cfg.use_delta_transform and (delta_dim_list is None or len(delta_dim_list) == 0): print("Warning: use_delta_transform=True but no delta_dims specified. No delta will be applied.") - print(f"Action horizon: {action_horizon}") - print(f"State key: {state_key}") + print(f"Action horizon: {cfg.action_horizon}") + print(f"State key: {cfg.state_key}") # determine episodes to process num_episodes = dataset.num_episodes - if max_episodes is not None: - num_episodes = min(max_episodes, num_episodes) + if cfg.max_episodes is not None: + num_episodes = min(cfg.max_episodes, num_episodes) print(f"Processing {num_episodes} episodes...") @@ -419,7 +437,15 @@ def main( print(f" Processing episode {ep_idx}/{num_episodes}...") chunks = process_episode( - (dataset, ep_idx, action_horizon, delta_dim_list, sample_fraction, state_key, use_delta_transform) + ( + dataset, + ep_idx, + cfg.action_horizon, + delta_dim_list, + cfg.sample_fraction, + cfg.state_key, + cfg.use_delta_transform, + ) ) if chunks is not None: all_chunks.append(chunks) @@ -495,16 +521,17 @@ def main( # train FAST tokenizer tokenizer = train_fast_tokenizer( encoded_chunks, - vocab_size=vocab_size, - scale=scale, + vocab_size=cfg.vocab_size, + scale=cfg.scale, ) # compute compression statistics compression_stats = compute_compression_stats(tokenizer, encoded_chunks) # save tokenizer + output_dir = cfg.output_dir if output_dir is None: - output_dir = f"fast_tokenizer_{repo_id.replace('/', '_')}" + output_dir = f"fast_tokenizer_{cfg.repo_id.replace('/', '_')}" output_path = Path(output_dir) output_path.mkdir(parents=True, exist_ok=True) @@ -512,18 +539,18 @@ def main( # save metadata metadata = { - "repo_id": repo_id, - "vocab_size": vocab_size, - "scale": scale, - "encoded_dims": encoded_dims, + "repo_id": cfg.repo_id, + "vocab_size": cfg.vocab_size, + "scale": cfg.scale, + "encoded_dims": cfg.encoded_dims, "encoded_dim_ranges": encoded_dim_ranges, "total_encoded_dims": total_encoded_dims, - "delta_dims": delta_dims, + "delta_dims": cfg.delta_dims, "delta_dim_list": delta_dim_list, - "use_delta_transform": use_delta_transform, - "state_key": state_key, + "use_delta_transform": cfg.use_delta_transform, + "state_key": cfg.state_key, "normalization_mode": norm_mode.value, - "action_horizon": action_horizon, + "action_horizon": cfg.action_horizon, "num_training_chunks": len(encoded_chunks), "compression_stats": compression_stats, } @@ -535,21 +562,22 @@ def main( print(f"Metadata: {json.dumps(metadata, indent=2)}") # push to Hugging Face Hub if requested - if push_to_hub: + if cfg.push_to_hub: # determine the hub repository ID + hub_repo_id = cfg.hub_repo_id if hub_repo_id is None: hub_repo_id = output_path.name print(f"\nNo hub_repo_id provided, using: {hub_repo_id}") print(f"\nPushing tokenizer to Hugging Face Hub: {hub_repo_id}") - print(f" Private: {hub_private}") + print(f" Private: {cfg.hub_private}") try: # use the tokenizer's push_to_hub method tokenizer.push_to_hub( repo_id=hub_repo_id, - private=hub_private, - commit_message=f"Upload FAST tokenizer trained on {repo_id}", + private=cfg.hub_private, + commit_message=f"Upload FAST tokenizer trained on {cfg.repo_id}", ) # also upload the metadata.json file separately @@ -568,5 +596,10 @@ def main( print(" Make sure you're logged in with `huggingface-cli login`") +def main(): + """CLI entry point that parses arguments and runs the tokenizer training.""" + train_tokenizer() + + if __name__ == "__main__": - tyro.cli(main) + main() From 15724826dd5b18c18656151a02a24ce5fd690c46 Mon Sep 17 00:00:00 2001 From: Steven Palma Date: Tue, 13 Jan 2026 09:49:46 +0100 Subject: [PATCH 5/5] chore: use alias & constants (#2785) * chore: use alias and constants * fix(rl): solve circular dependecy * chore: nit right constant * chore: pre-commit * chore(script): conflict tokenizer train --------- Signed-off-by: Steven Palma --- src/lerobot/async_inference/robot_client.py | 4 +-- src/lerobot/datasets/pipeline_features.py | 4 +-- src/lerobot/envs/libero.py | 6 ++-- src/lerobot/envs/metaworld.py | 12 ++++---- src/lerobot/envs/utils.py | 3 +- src/lerobot/policies/factory.py | 8 ++++-- .../policies/groot/configuration_groot.py | 13 +++++---- src/lerobot/policies/groot/groot_n1.py | 6 ++-- src/lerobot/policies/groot/modeling_groot.py | 3 +- src/lerobot/policies/groot/processor_groot.py | 28 +++++++++++-------- src/lerobot/policies/pi0/configuration_pi0.py | 10 +++---- .../policies/pi05/configuration_pi05.py | 11 ++++---- .../pi0_fast/configuration_pi0_fast.py | 11 ++++---- .../policies/sarm/configuration_sarm.py | 5 ++-- src/lerobot/policies/sarm/modeling_sarm.py | 3 +- src/lerobot/policies/utils.py | 3 +- .../policies/wall_x/configuration_wall_x.py | 13 +++++---- .../policies/wall_x/modeling_wall_x.py | 6 ++-- src/lerobot/policies/xvla/processor_xvla.py | 16 ++++++----- src/lerobot/processor/__init__.py | 3 -- src/lerobot/processor/converters.py | 6 ++-- src/lerobot/processor/core.py | 2 +- src/lerobot/processor/env_processor.py | 11 ++++---- src/lerobot/processor/normalize_processor.py | 4 +-- src/lerobot/processor/pipeline.py | 6 ++-- src/lerobot/processor/tokenizer_processor.py | 4 +-- src/lerobot/rl/gym_manipulator.py | 11 ++++---- .../joint_observations_processor.py | 0 .../robots/bi_so_follower/bi_so_follower.py | 6 ++-- .../robot_earthrover_mini_plus.py | 11 ++++---- src/lerobot/robots/hope_jr/hope_jr_arm.py | 6 ++-- src/lerobot/robots/hope_jr/hope_jr_hand.py | 6 ++-- .../robots/koch_follower/koch_follower.py | 10 +++---- src/lerobot/robots/lekiwi/lekiwi.py | 7 +++-- src/lerobot/robots/lekiwi/lekiwi_client.py | 19 ++++++------- .../robots/omx_follower/omx_follower.py | 10 +++---- src/lerobot/robots/reachy2/robot_reachy2.py | 9 +++--- src/lerobot/robots/robot.py | 12 ++++---- .../so_follower/robot_kinematic_processor.py | 3 +- src/lerobot/robots/so_follower/so_follower.py | 9 +++--- src/lerobot/robots/unitree_g1/unitree_g1.py | 5 ++-- src/lerobot/scripts/lerobot_eval.py | 4 +-- src/lerobot/scripts/lerobot_teleoperate.py | 2 +- .../scripts/lerobot_train_tokenizer.py | 13 ++++----- .../teleoperators/gamepad/teleop_gamepad.py | 4 ++- .../teleoperators/keyboard/teleop_keyboard.py | 9 +++--- src/lerobot/teleoperators/teleoperator.py | 5 ++-- src/lerobot/utils/constants.py | 2 ++ src/lerobot/utils/visualization_utils.py | 11 ++++---- tests/mocks/mock_robot.py | 6 ++-- tests/mocks/mock_teleop.py | 3 +- 51 files changed, 206 insertions(+), 178 deletions(-) rename src/lerobot/{processor => rl}/joint_observations_processor.py (100%) diff --git a/src/lerobot/async_inference/robot_client.py b/src/lerobot/async_inference/robot_client.py index eea5585b0..f26639dc1 100644 --- a/src/lerobot/async_inference/robot_client.py +++ b/src/lerobot/async_inference/robot_client.py @@ -40,7 +40,6 @@ from collections.abc import Callable from dataclasses import asdict from pprint import pformat from queue import Queue -from typing import Any import draccus import grpc @@ -48,6 +47,7 @@ import torch from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig # noqa: F401 from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig # noqa: F401 +from lerobot.processor import RobotAction from lerobot.robots import ( # noqa: F401 Robot, RobotConfig, @@ -351,7 +351,7 @@ class RobotClient: action = {key: action_tensor[i].item() for i, key in enumerate(self.robot.action_features)} return action - def control_loop_action(self, verbose: bool = False) -> dict[str, Any]: + def control_loop_action(self, verbose: bool = False) -> RobotAction: """Reading and performing actions in local queue""" # Lock only for queue operations diff --git a/src/lerobot/datasets/pipeline_features.py b/src/lerobot/datasets/pipeline_features.py index 4fad7bd20..161633f26 100644 --- a/src/lerobot/datasets/pipeline_features.py +++ b/src/lerobot/datasets/pipeline_features.py @@ -18,12 +18,12 @@ from typing import Any from lerobot.configs.types import PipelineFeatureType from lerobot.datasets.utils import hw_to_dataset_features -from lerobot.processor import DataProcessorPipeline +from lerobot.processor import DataProcessorPipeline, RobotAction, RobotObservation from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE, OBS_STR def create_initial_features( - action: dict[str, Any] | None = None, observation: dict[str, Any] | None = None + action: RobotAction | None = None, observation: RobotObservation | None = None ) -> dict[PipelineFeatureType, dict[str, Any]]: """ Creates the initial features dict for the dataset from action and observation specs. diff --git a/src/lerobot/envs/libero.py b/src/lerobot/envs/libero.py index b1eb37377..74882ad18 100644 --- a/src/lerobot/envs/libero.py +++ b/src/lerobot/envs/libero.py @@ -29,6 +29,8 @@ from gymnasium import spaces from libero.libero import benchmark, get_libero_path from libero.libero.envs import OffScreenRenderEnv +from lerobot.processor import RobotObservation + def _parse_camera_names(camera_name: str | Sequence[str]) -> list[str]: """Normalize camera_name into a non-empty list of strings.""" @@ -237,7 +239,7 @@ class LiberoEnv(gym.Env): env.reset() return env - def _format_raw_obs(self, raw_obs: dict[str, Any]) -> dict[str, Any]: + def _format_raw_obs(self, raw_obs: RobotObservation) -> RobotObservation: images = {} for camera_name in self.camera_name: image = raw_obs[camera_name] @@ -313,7 +315,7 @@ class LiberoEnv(gym.Env): info = {"is_success": False} return observation, info - def step(self, action: np.ndarray) -> tuple[dict[str, Any], float, bool, bool, dict[str, Any]]: + def step(self, action: np.ndarray) -> tuple[RobotObservation, float, bool, bool, dict[str, Any]]: if action.ndim != 1: raise ValueError( f"Expected action to be 1-D (shape (action_dim,)), " diff --git a/src/lerobot/envs/metaworld.py b/src/lerobot/envs/metaworld.py index 9190f33ad..4d91e002d 100644 --- a/src/lerobot/envs/metaworld.py +++ b/src/lerobot/envs/metaworld.py @@ -25,6 +25,8 @@ import metaworld.policies as policies import numpy as np from gymnasium import spaces +from lerobot.processor import RobotObservation + # ---- Load configuration data from the external JSON file ---- CONFIG_PATH = Path(__file__).parent / "metaworld_config.json" try: @@ -161,7 +163,7 @@ class MetaworldEnv(gym.Env): env._freeze_rand_vec = False # otherwise no randomization return env - def _format_raw_obs(self, raw_obs: np.ndarray) -> dict[str, Any]: + def _format_raw_obs(self, raw_obs: np.ndarray) -> RobotObservation: image = None if self._env is not None: image = self._env.render() @@ -196,7 +198,7 @@ class MetaworldEnv(gym.Env): self, seed: int | None = None, **kwargs, - ) -> tuple[dict[str, Any], dict[str, Any]]: + ) -> tuple[RobotObservation, dict[str, Any]]: """ Reset the environment to its initial state. @@ -204,7 +206,7 @@ class MetaworldEnv(gym.Env): seed (Optional[int]): Random seed for environment initialization. Returns: - observation (Dict[str, Any]): The initial formatted observation. + observation (RobotObservation): The initial formatted observation. info (Dict[str, Any]): Additional info about the reset state. """ super().reset(seed=seed) @@ -216,7 +218,7 @@ class MetaworldEnv(gym.Env): info = {"is_success": False} return observation, info - def step(self, action: np.ndarray) -> tuple[dict[str, Any], float, bool, bool, dict[str, Any]]: + def step(self, action: np.ndarray) -> tuple[RobotObservation, float, bool, bool, dict[str, Any]]: """ Perform one environment step. @@ -224,7 +226,7 @@ class MetaworldEnv(gym.Env): action (np.ndarray): The action to execute, must be 1-D with shape (action_dim,). Returns: - observation (Dict[str, Any]): The formatted observation after the step. + observation (RobotObservation): The formatted observation after the step. reward (float): The scalar reward for this step. terminated (bool): Whether the episode terminated successfully. truncated (bool): Whether the episode was truncated due to a time limit. diff --git a/src/lerobot/envs/utils.py b/src/lerobot/envs/utils.py index af814c92a..09431a18d 100644 --- a/src/lerobot/envs/utils.py +++ b/src/lerobot/envs/utils.py @@ -29,6 +29,7 @@ from torch import Tensor from lerobot.configs.types import FeatureType, PolicyFeature from lerobot.envs.configs import EnvConfig +from lerobot.processor import RobotObservation from lerobot.utils.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE, OBS_STR from lerobot.utils.utils import get_channel_first_image_shape @@ -152,7 +153,7 @@ def check_env_attributes_and_types(env: gym.vector.VectorEnv) -> None: ) -def add_envs_task(env: gym.vector.VectorEnv, observation: dict[str, Any]) -> dict[str, Any]: +def add_envs_task(env: gym.vector.VectorEnv, observation: RobotObservation) -> RobotObservation: """Adds task feature to the observation dict with respect to the first environment attribute.""" if hasattr(env.envs[0], "task_description"): task_result = env.call("task_description") diff --git a/src/lerobot/policies/factory.py b/src/lerobot/policies/factory.py index fff08ad37..a593e5bcb 100644 --- a/src/lerobot/policies/factory.py +++ b/src/lerobot/policies/factory.py @@ -51,7 +51,11 @@ from lerobot.processor.converters import ( transition_to_batch, transition_to_policy_action, ) -from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME +from lerobot.utils.constants import ( + ACTION, + POLICY_POSTPROCESSOR_DEFAULT_NAME, + POLICY_PREPROCESSOR_DEFAULT_NAME, +) def get_policy_class(name: str) -> type[PreTrainedPolicy]: @@ -250,7 +254,7 @@ def make_pre_post_processors( } # Also ensure postprocessing slices to env action dim and unnormalizes with dataset stats - env_action_dim = policy_cfg.output_features["action"].shape[0] + env_action_dim = policy_cfg.output_features[ACTION].shape[0] postprocessor_overrides["groot_action_unpack_unnormalize_v1"] = { "stats": kwargs.get("dataset_stats"), "normalize_min_max": True, diff --git a/src/lerobot/policies/groot/configuration_groot.py b/src/lerobot/policies/groot/configuration_groot.py index 8002c69ea..4f3d78222 100644 --- a/src/lerobot/policies/groot/configuration_groot.py +++ b/src/lerobot/policies/groot/configuration_groot.py @@ -20,6 +20,7 @@ from lerobot.configs.policies import PreTrainedConfig from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature from lerobot.optim.optimizers import AdamWConfig from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig +from lerobot.utils.constants import ACTION, OBS_STATE @PreTrainedConfig.register_subclass("groot") @@ -137,14 +138,14 @@ class GrootConfig(PreTrainedConfig): "No features of type FeatureType.VISUAL found in input_features." ) - if "observation.state" not in self.input_features: + if OBS_STATE not in self.input_features: state_feature = PolicyFeature( type=FeatureType.STATE, shape=(self.max_state_dim,), ) - self.input_features["observation.state"] = state_feature + self.input_features[OBS_STATE] = state_feature else: - state_shape = self.input_features["observation.state"].shape + state_shape = self.input_features[OBS_STATE].shape state_dim = state_shape[0] if state_shape else 0 if state_dim > self.max_state_dim: raise ValueError( @@ -152,14 +153,14 @@ class GrootConfig(PreTrainedConfig): f"Either reduce state dimension or increase max_state_dim in config." ) - if "action" not in self.output_features: + if ACTION not in self.output_features: action_feature = PolicyFeature( type=FeatureType.ACTION, shape=(self.max_action_dim,), ) - self.output_features["action"] = action_feature + self.output_features[ACTION] = action_feature else: - action_shape = self.output_features["action"].shape + action_shape = self.output_features[ACTION].shape action_dim = action_shape[0] if action_shape else 0 if action_dim > self.max_action_dim: raise ValueError( diff --git a/src/lerobot/policies/groot/groot_n1.py b/src/lerobot/policies/groot/groot_n1.py index 0da26874e..06ff5a04d 100644 --- a/src/lerobot/policies/groot/groot_n1.py +++ b/src/lerobot/policies/groot/groot_n1.py @@ -46,7 +46,7 @@ from lerobot.policies.groot.action_head.flow_matching_action_head import ( FlowmatchingActionHeadConfig, ) from lerobot.policies.groot.utils import ensure_eagle_cache_ready -from lerobot.utils.constants import HF_LEROBOT_HOME +from lerobot.utils.constants import ACTION, HF_LEROBOT_HOME DEFAULT_VENDOR_EAGLE_PATH = str((Path(__file__).resolve().parent / "eagle2_hg_model").resolve()) DEFAULT_TOKENIZER_ASSETS_REPO = "lerobot/eagle2hg-processor-groot-n1p5" @@ -227,8 +227,8 @@ class GR00TN15(PreTrainedModel): detected_error = False error_msg = ERROR_MSG - if "action" in inputs: - action = inputs["action"] + if ACTION in inputs: + action = inputs[ACTION] # In inference, action may be omitted or None; validate only when it's a tensor. if action is None: pass # allow None during inference diff --git a/src/lerobot/policies/groot/modeling_groot.py b/src/lerobot/policies/groot/modeling_groot.py index bdaef37b9..fd9baa9b1 100644 --- a/src/lerobot/policies/groot/modeling_groot.py +++ b/src/lerobot/policies/groot/modeling_groot.py @@ -41,6 +41,7 @@ from torch import Tensor from lerobot.policies.groot.configuration_groot import GrootConfig from lerobot.policies.groot.groot_n1 import GR00TN15 from lerobot.policies.pretrained import PreTrainedPolicy +from lerobot.utils.constants import ACTION class GrootPolicy(PreTrainedPolicy): @@ -147,7 +148,7 @@ class GrootPolicy(PreTrainedPolicy): actions = outputs.get("action_pred") - original_action_dim = self.config.output_features["action"].shape[0] + original_action_dim = self.config.output_features[ACTION].shape[0] actions = actions[:, :, :original_action_dim] return actions diff --git a/src/lerobot/policies/groot/processor_groot.py b/src/lerobot/policies/groot/processor_groot.py index d87e43c11..14149cf2f 100644 --- a/src/lerobot/policies/groot/processor_groot.py +++ b/src/lerobot/policies/groot/processor_groot.py @@ -51,7 +51,11 @@ from lerobot.processor.converters import ( ) from lerobot.processor.core import EnvTransition, TransitionKey from lerobot.utils.constants import ( + ACTION, HF_LEROBOT_HOME, + OBS_IMAGE, + OBS_IMAGES, + OBS_STATE, POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME, ) @@ -107,9 +111,9 @@ def make_groot_pre_post_processors( # Define feature specs for optional normalization steps _features: dict[str, PolicyFeature] = { # Observation features (only add those we may normalize) - "observation.state": PolicyFeature(type=FeatureType.STATE, shape=(state_horizon, max_state_dim)), + OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(state_horizon, max_state_dim)), # Action feature - "action": PolicyFeature(type=FeatureType.ACTION, shape=(action_horizon, max_action_dim)), + ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(action_horizon, max_action_dim)), } # Normalize STATE and ACTION with min_max (SO100-like default) @@ -120,7 +124,7 @@ def make_groot_pre_post_processors( # Determine env action dimension from config (simple, object-like PolicyFeature) try: - env_action_dim = int(config.output_features["action"].shape[0]) + env_action_dim = int(config.output_features[ACTION].shape[0]) except Exception: env_action_dim = 0 @@ -268,9 +272,9 @@ class GrootPackInputsStep(ProcessorStep): return torch.where(mask, mapped, torch.zeros_like(mapped)) # 1) Video (B, T=1, V, H, W, C) uint8 - img_keys = sorted([k for k in obs if k.startswith("observation.images.")]) - if not img_keys and "observation.image" in obs: - img_keys = ["observation.image"] + img_keys = sorted([k for k in obs if k.startswith(OBS_IMAGES)]) + if not img_keys and OBS_IMAGE in obs: + img_keys = [OBS_IMAGE] if img_keys: cams = [_to_uint8_np_bhwc(obs[k]) for k in img_keys] video = np.stack(cams, axis=1) # (B, V, H, W, C) @@ -294,14 +298,14 @@ class GrootPackInputsStep(ProcessorStep): comp["language"] = lang # 3) State/state_mask -> (B, 1, max_state_dim) - if "observation.state" in obs: - state = obs["observation.state"] # (B, D) + if OBS_STATE in obs: + state = obs[OBS_STATE] # (B, D) if state.dim() != 2: raise ValueError(f"state must be (B, D), got {tuple(state.shape)}") bsz, d = state.shape # Normalize BEFORE padding if self.normalize_min_max: - state = _min_max_norm(state, "observation.state") + state = _min_max_norm(state, OBS_STATE) state = state.unsqueeze(1) # (B, 1, D) if d > self.max_state_dim: state = state[:, :, : self.max_state_dim] @@ -320,11 +324,11 @@ class GrootPackInputsStep(ProcessorStep): # Normalize BEFORE temporal expansion/padding if self.normalize_min_max: if action.dim() == 2: - action = _min_max_norm(action, "action") + action = _min_max_norm(action, ACTION) elif action.dim() == 3: b, t, d = action.shape flat = action.reshape(b * t, d) - flat = _min_max_norm(flat, "action") + flat = _min_max_norm(flat, ACTION) action = flat.view(b, t, d) if action.dim() == 2: action = action.unsqueeze(1).repeat(1, self.action_horizon, 1) @@ -590,7 +594,7 @@ class GrootActionUnpackUnnormalizeStep(ProcessorStep): # forward: y = 2 * (x - min) / denom - 1, with y=0 when denom==0 # inverse: x = (y+1)/2 * denom + min, and when denom==0 -> x = min if self.normalize_min_max and self.stats is not None: - stats_k = self.stats.get("action", {}) + stats_k = self.stats.get(ACTION, {}) d = action.shape[-1] min_v = torch.as_tensor( stats_k.get("min", torch.zeros(d)), dtype=action.dtype, device=action.device diff --git a/src/lerobot/policies/pi0/configuration_pi0.py b/src/lerobot/policies/pi0/configuration_pi0.py index d91145aa7..be9b4530f 100644 --- a/src/lerobot/policies/pi0/configuration_pi0.py +++ b/src/lerobot/policies/pi0/configuration_pi0.py @@ -21,7 +21,7 @@ from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature from lerobot.optim.optimizers import AdamWConfig from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig from lerobot.policies.rtc.configuration_rtc import RTCConfig -from lerobot.utils.constants import OBS_IMAGES +from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE DEFAULT_IMAGE_SIZE = 224 @@ -124,19 +124,19 @@ class PI0Config(PreTrainedConfig): ) self.input_features[key] = empty_camera - if "observation.state" not in self.input_features: + if OBS_STATE not in self.input_features: state_feature = PolicyFeature( type=FeatureType.STATE, shape=(self.max_state_dim,), # Padded to max_state_dim ) - self.input_features["observation.state"] = state_feature + self.input_features[OBS_STATE] = state_feature - if "action" not in self.output_features: + if ACTION not in self.output_features: action_feature = PolicyFeature( type=FeatureType.ACTION, shape=(self.max_action_dim,), # Padded to max_action_dim ) - self.output_features["action"] = action_feature + self.output_features[ACTION] = action_feature def get_optimizer_preset(self) -> AdamWConfig: return AdamWConfig( diff --git a/src/lerobot/policies/pi05/configuration_pi05.py b/src/lerobot/policies/pi05/configuration_pi05.py index 4a0e23039..b96e6d196 100644 --- a/src/lerobot/policies/pi05/configuration_pi05.py +++ b/src/lerobot/policies/pi05/configuration_pi05.py @@ -21,6 +21,7 @@ from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature from lerobot.optim.optimizers import AdamWConfig from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig from lerobot.policies.rtc.configuration_rtc import RTCConfig +from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE DEFAULT_IMAGE_SIZE = 224 @@ -117,26 +118,26 @@ class PI05Config(PreTrainedConfig): def validate_features(self) -> None: """Validate and set up input/output features.""" for i in range(self.empty_cameras): - key = f"observation.images.empty_camera_{i}" + key = OBS_IMAGES + f".empty_camera_{i}" empty_camera = PolicyFeature( type=FeatureType.VISUAL, shape=(3, *self.image_resolution), # Use configured image resolution ) self.input_features[key] = empty_camera - if "observation.state" not in self.input_features: + if OBS_STATE not in self.input_features: state_feature = PolicyFeature( type=FeatureType.STATE, shape=(self.max_state_dim,), # Padded to max_state_dim ) - self.input_features["observation.state"] = state_feature + self.input_features[OBS_STATE] = state_feature - if "action" not in self.output_features: + if ACTION not in self.output_features: action_feature = PolicyFeature( type=FeatureType.ACTION, shape=(self.max_action_dim,), # Padded to max_action_dim ) - self.output_features["action"] = action_feature + self.output_features[ACTION] = action_feature def get_optimizer_preset(self) -> AdamWConfig: return AdamWConfig( diff --git a/src/lerobot/policies/pi0_fast/configuration_pi0_fast.py b/src/lerobot/policies/pi0_fast/configuration_pi0_fast.py index 42aa4a132..96137e91f 100644 --- a/src/lerobot/policies/pi0_fast/configuration_pi0_fast.py +++ b/src/lerobot/policies/pi0_fast/configuration_pi0_fast.py @@ -21,6 +21,7 @@ from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature from lerobot.optim.optimizers import AdamWConfig from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig from lerobot.policies.rtc.configuration_rtc import RTCConfig +from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE DEFAULT_IMAGE_SIZE = 224 @@ -110,26 +111,26 @@ class PI0FastConfig(PreTrainedConfig): def validate_features(self) -> None: """Validate and set up input/output features.""" for i in range(self.empty_cameras): - key = f"observation.images.empty_camera_{i}" + key = OBS_IMAGES + f".empty_camera_{i}" empty_camera = PolicyFeature( type=FeatureType.VISUAL, shape=(3, *self.image_resolution), # Use configured image resolution ) self.input_features[key] = empty_camera - if "observation.state" not in self.input_features: + if OBS_STATE not in self.input_features: state_feature = PolicyFeature( type=FeatureType.STATE, shape=(self.max_state_dim,), # Padded to max_state_dim ) - self.input_features["observation.state"] = state_feature + self.input_features[OBS_STATE] = state_feature - if "action" not in self.output_features: + if ACTION not in self.output_features: action_feature = PolicyFeature( type=FeatureType.ACTION, shape=(self.max_action_dim,), # Padded to max_action_dim ) - self.output_features["action"] = action_feature + self.output_features[ACTION] = action_feature def get_optimizer_preset(self) -> AdamWConfig: return AdamWConfig( diff --git a/src/lerobot/policies/sarm/configuration_sarm.py b/src/lerobot/policies/sarm/configuration_sarm.py index 59cb352d5..673422fe2 100644 --- a/src/lerobot/policies/sarm/configuration_sarm.py +++ b/src/lerobot/policies/sarm/configuration_sarm.py @@ -26,6 +26,7 @@ from lerobot.configs.policies import PreTrainedConfig from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature from lerobot.optim.optimizers import AdamWConfig from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig +from lerobot.utils.constants import OBS_IMAGES, OBS_STATE @PreTrainedConfig.register_subclass("sarm") @@ -86,8 +87,8 @@ class SARMConfig(PreTrainedConfig): pretrained_model_path: str | None = None device: str | None = None - image_key: str = "observation.images.top" # Key for image used from the dataset - state_key: str = "observation.state" + image_key: str = OBS_IMAGES + ".top" # Key for image used from the dataset + state_key: str = OBS_STATE # Populated by the processor (video_features, state_features, text_features) input_features: dict = field(default_factory=lambda: {}) diff --git a/src/lerobot/policies/sarm/modeling_sarm.py b/src/lerobot/policies/sarm/modeling_sarm.py index a88b2ad64..6051d90f8 100644 --- a/src/lerobot/policies/sarm/modeling_sarm.py +++ b/src/lerobot/policies/sarm/modeling_sarm.py @@ -40,6 +40,7 @@ from lerobot.policies.sarm.sarm_utils import ( normalize_stage_tau, pad_state_to_max_dim, ) +from lerobot.utils.constants import OBS_STR class StageTransformer(nn.Module): @@ -721,7 +722,7 @@ class SARMRewardModel(PreTrainedPolicy): Returns: Tuple of (total_loss, output_dict with loss components) """ - observation = batch.get("observation", batch) + observation = batch.get(OBS_STR, batch) # Extract features video_features = observation["video_features"].to(self.device) diff --git a/src/lerobot/policies/utils.py b/src/lerobot/policies/utils.py index bfbe2bf1d..1a14b2925 100644 --- a/src/lerobot/policies/utils.py +++ b/src/lerobot/policies/utils.py @@ -16,7 +16,6 @@ import logging from collections import deque -from typing import Any import numpy as np import torch @@ -140,7 +139,7 @@ def prepare_observation_for_inference( def build_inference_frame( - observation: dict[str, Any], + observation: RobotObservation, device: torch.device, ds_features: dict[str, dict], task: str | None = None, diff --git a/src/lerobot/policies/wall_x/configuration_wall_x.py b/src/lerobot/policies/wall_x/configuration_wall_x.py index 0d10a8f98..3962b56f6 100644 --- a/src/lerobot/policies/wall_x/configuration_wall_x.py +++ b/src/lerobot/policies/wall_x/configuration_wall_x.py @@ -18,6 +18,7 @@ from lerobot.configs.policies import PreTrainedConfig from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature from lerobot.optim.optimizers import AdamWConfig from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig +from lerobot.utils.constants import ACTION, OBS_STATE @PreTrainedConfig.register_subclass("wall_x") @@ -105,14 +106,14 @@ class WallXConfig(PreTrainedConfig): "No features of type FeatureType.VISUAL found in input_features." ) - if "observation.state" not in self.input_features: + if OBS_STATE not in self.input_features: state_feature = PolicyFeature( type=FeatureType.STATE, shape=(self.max_state_dim,), # Padded to max_state_dim ) - self.input_features["observation.state"] = state_feature + self.input_features[OBS_STATE] = state_feature else: - state_shape = self.input_features["observation.state"].shape + state_shape = self.input_features[OBS_STATE].shape state_dim = state_shape[0] if state_shape else 0 if state_dim > self.max_state_dim: raise ValueError( @@ -120,14 +121,14 @@ class WallXConfig(PreTrainedConfig): f"Either reduce state dimension or increase max_state_dim in config." ) - if "action" not in self.output_features: + if ACTION not in self.output_features: action_feature = PolicyFeature( type=FeatureType.ACTION, shape=(self.max_action_dim,), # Padded to max_action_dim ) - self.output_features["action"] = action_feature + self.output_features[ACTION] = action_feature else: - action_shape = self.output_features["action"].shape + action_shape = self.output_features[ACTION].shape action_dim = action_shape[0] if action_shape else 0 if action_dim > self.max_action_dim: raise ValueError( diff --git a/src/lerobot/policies/wall_x/modeling_wall_x.py b/src/lerobot/policies/wall_x/modeling_wall_x.py index 94ee7897e..ef99bad89 100644 --- a/src/lerobot/policies/wall_x/modeling_wall_x.py +++ b/src/lerobot/policies/wall_x/modeling_wall_x.py @@ -1861,7 +1861,7 @@ class WallXPolicy(PreTrainedPolicy): dim=-1, ) else: - action_dim = self.config.output_features["action"].shape[0] + action_dim = self.config.output_features[ACTION].shape[0] dof_mask = torch.cat( [ torch.ones( @@ -1977,7 +1977,7 @@ class WallXPolicy(PreTrainedPolicy): elif self.config.prediction_mode == "fast": output = self.model( **batch, - action_dim=self.config.output_features["action"].shape[0], + action_dim=self.config.output_features[ACTION].shape[0], pred_horizon=self.config.chunk_size, mode="predict", predict_mode="fast", @@ -1989,7 +1989,7 @@ class WallXPolicy(PreTrainedPolicy): actions = output["predict_action"] # Unpad actions to actual action dimension - action_dim = self.config.output_features["action"].shape[0] + action_dim = self.config.output_features[ACTION].shape[0] actions = actions[:, :, :action_dim] return actions diff --git a/src/lerobot/policies/xvla/processor_xvla.py b/src/lerobot/policies/xvla/processor_xvla.py index 7f7297b9a..c4e3f2d6f 100644 --- a/src/lerobot/policies/xvla/processor_xvla.py +++ b/src/lerobot/policies/xvla/processor_xvla.py @@ -41,6 +41,7 @@ from lerobot.processor.converters import policy_action_to_transition, transition from lerobot.processor.core import EnvTransition, TransitionKey from lerobot.utils.constants import ( OBS_IMAGES, + OBS_PREFIX, OBS_STATE, POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME, @@ -137,8 +138,9 @@ class LiberoProcessorStep(ObservationProcessorStep): processed_obs[key] = img # Process robot_state into a flat state vector - if "observation.robot_state" in processed_obs: - robot_state = processed_obs.pop("observation.robot_state") + robot_state_str = OBS_PREFIX + "robot_state" + if robot_state_str in processed_obs: + robot_state = processed_obs.pop(robot_state_str) # Extract components eef_pos = robot_state["eef"]["pos"] # (B, 3,) @@ -174,8 +176,8 @@ class LiberoProcessorStep(ObservationProcessorStep): state_feats = {} # add our new flattened state - state_feats["observation.state"] = PolicyFeature( - key="observation.state", + state_feats[OBS_STATE] = PolicyFeature( + key=OBS_STATE, shape=(20,), dtype="float32", ) @@ -247,7 +249,7 @@ class XVLAImageScaleProcessorStep(ProcessorStep): keys_to_scale = self.image_keys if keys_to_scale is None: # Auto-detect image keys - keys_to_scale = [k for k in obs if k.startswith("observation.images.")] + keys_to_scale = [k for k in obs if k.startswith(OBS_IMAGES)] # Scale each image for key in keys_to_scale: @@ -303,7 +305,7 @@ class XVLAImageToFloatProcessorStep(ProcessorStep): keys_to_convert = self.image_keys if keys_to_convert is None: # Auto-detect image keys - keys_to_convert = [k for k in obs if k.startswith("observation.images.")] + keys_to_convert = [k for k in obs if k.startswith(OBS_IMAGES)] # Convert each image for key in keys_to_convert: @@ -376,7 +378,7 @@ class XVLAImageNetNormalizeProcessorStep(ProcessorStep): keys_to_normalize = self.image_keys if keys_to_normalize is None: # Auto-detect image keys - keys_to_normalize = [k for k in obs if k.startswith("observation.images.")] + keys_to_normalize = [k for k in obs if k.startswith(OBS_IMAGES)] # Normalize each image for key in keys_to_normalize: diff --git a/src/lerobot/processor/__init__.py b/src/lerobot/processor/__init__.py index 676ba29ee..164f7da03 100644 --- a/src/lerobot/processor/__init__.py +++ b/src/lerobot/processor/__init__.py @@ -49,7 +49,6 @@ from .hil_processor import ( RewardClassifierProcessorStep, TimeLimitProcessorStep, ) -from .joint_observations_processor import JointVelocityProcessorStep, MotorCurrentProcessorStep from .normalize_processor import NormalizerProcessorStep, UnnormalizerProcessorStep, hotswap_stats from .observation_processor import VanillaObservationProcessorStep from .pipeline import ( @@ -94,14 +93,12 @@ __all__ = [ "ImageCropResizeProcessorStep", "InfoProcessorStep", "InterventionActionProcessorStep", - "JointVelocityProcessorStep", "make_default_processors", "make_default_teleop_action_processor", "make_default_robot_action_processor", "make_default_robot_observation_processor", "MapDeltaActionToRobotActionStep", "MapTensorToDeltaActionDictStep", - "MotorCurrentProcessorStep", "NormalizerProcessorStep", "Numpy2TorchActionProcessorStep", "ObservationProcessorStep", diff --git a/src/lerobot/processor/converters.py b/src/lerobot/processor/converters.py index 126be0e36..4f9485fee 100644 --- a/src/lerobot/processor/converters.py +++ b/src/lerobot/processor/converters.py @@ -23,7 +23,7 @@ from typing import Any import numpy as np import torch -from lerobot.utils.constants import ACTION, DONE, OBS_PREFIX, REWARD, TRUNCATED +from lerobot.utils.constants import ACTION, DONE, INFO, OBS_PREFIX, REWARD, TRUNCATED from .core import EnvTransition, PolicyAction, RobotAction, RobotObservation, TransitionKey @@ -176,7 +176,7 @@ def _extract_complementary_data(batch: dict[str, Any]) -> dict[str, Any]: def create_transition( - observation: dict[str, Any] | None = None, + observation: RobotObservation | None = None, action: PolicyAction | RobotAction | None = None, reward: float = 0.0, done: bool = False, @@ -384,7 +384,7 @@ def transition_to_batch(transition: EnvTransition) -> dict[str, Any]: REWARD: transition.get(TransitionKey.REWARD, 0.0), DONE: transition.get(TransitionKey.DONE, False), TRUNCATED: transition.get(TransitionKey.TRUNCATED, False), - "info": transition.get(TransitionKey.INFO, {}), + INFO: transition.get(TransitionKey.INFO, {}), } # Add complementary data. diff --git a/src/lerobot/processor/core.py b/src/lerobot/processor/core.py index 679ba8c54..0b293c9b0 100644 --- a/src/lerobot/processor/core.py +++ b/src/lerobot/processor/core.py @@ -45,7 +45,7 @@ RobotObservation: TypeAlias = dict[str, Any] EnvTransition = TypedDict( "EnvTransition", { - TransitionKey.OBSERVATION.value: dict[str, Any] | None, + TransitionKey.OBSERVATION.value: RobotObservation | None, TransitionKey.ACTION.value: PolicyAction | RobotAction | EnvAction | None, TransitionKey.REWARD.value: float | torch.Tensor | None, TransitionKey.DONE.value: bool | torch.Tensor | None, diff --git a/src/lerobot/processor/env_processor.py b/src/lerobot/processor/env_processor.py index a5210af30..8d42bfdb7 100644 --- a/src/lerobot/processor/env_processor.py +++ b/src/lerobot/processor/env_processor.py @@ -18,7 +18,7 @@ from dataclasses import dataclass import torch from lerobot.configs.types import PipelineFeatureType, PolicyFeature -from lerobot.utils.constants import OBS_IMAGES, OBS_STATE, OBS_STR +from lerobot.utils.constants import OBS_IMAGES, OBS_PREFIX, OBS_STATE, OBS_STR from .pipeline import ObservationProcessorStep, ProcessorStepRegistry @@ -60,8 +60,9 @@ class LiberoProcessorStep(ObservationProcessorStep): processed_obs[key] = img # Process robot_state into a flat state vector - if "observation.robot_state" in processed_obs: - robot_state = processed_obs.pop("observation.robot_state") + observation_robot_state_str = OBS_PREFIX + "robot_state" + if observation_robot_state_str in processed_obs: + robot_state = processed_obs.pop(observation_robot_state_str) # Extract components eef_pos = robot_state["eef"]["pos"] # (B, 3,) @@ -98,8 +99,8 @@ class LiberoProcessorStep(ObservationProcessorStep): state_feats = {} # add our new flattened state - state_feats["observation.state"] = PolicyFeature( - key="observation.state", + state_feats[OBS_STATE] = PolicyFeature( + key=OBS_STATE, shape=(8,), # [eef_pos(3), axis_angle(3), gripper(2)] dtype="float32", description=("Concatenated end-effector position (3), axis-angle (3), and gripper qpos (2)."), diff --git a/src/lerobot/processor/normalize_processor.py b/src/lerobot/processor/normalize_processor.py index 368c9b270..4769b91ac 100644 --- a/src/lerobot/processor/normalize_processor.py +++ b/src/lerobot/processor/normalize_processor.py @@ -30,7 +30,7 @@ from lerobot.utils.constants import ACTION from .converters import from_tensor_to_numpy, to_tensor from .core import EnvTransition, PolicyAction, TransitionKey -from .pipeline import PolicyProcessorPipeline, ProcessorStep, ProcessorStepRegistry +from .pipeline import PolicyProcessorPipeline, ProcessorStep, ProcessorStepRegistry, RobotObservation @dataclass @@ -239,7 +239,7 @@ class _NormalizationMixin: config["normalize_observation_keys"] = sorted(self.normalize_observation_keys) return config - def _normalize_observation(self, observation: dict[str, Any], inverse: bool) -> dict[str, Tensor]: + def _normalize_observation(self, observation: RobotObservation, inverse: bool) -> dict[str, Tensor]: """ Applies (un)normalization to all relevant features in an observation dictionary. diff --git a/src/lerobot/processor/pipeline.py b/src/lerobot/processor/pipeline.py index e14d8b0b9..97ec716ff 100644 --- a/src/lerobot/processor/pipeline.py +++ b/src/lerobot/processor/pipeline.py @@ -49,7 +49,7 @@ from lerobot.configs.types import PipelineFeatureType, PolicyFeature from lerobot.utils.hub import HubMixin from .converters import batch_to_transition, create_transition, transition_to_batch -from .core import EnvAction, EnvTransition, PolicyAction, RobotAction, TransitionKey +from .core import EnvAction, EnvTransition, PolicyAction, RobotAction, RobotObservation, TransitionKey # Generic type variables for pipeline input and output. TInput = TypeVar("TInput") @@ -1337,7 +1337,7 @@ class DataProcessorPipeline(HubMixin, Generic[TInput, TOutput]): return features # Convenience methods for processing individual parts of a transition. - def process_observation(self, observation: dict[str, Any]) -> dict[str, Any]: + def process_observation(self, observation: RobotObservation) -> RobotObservation: """Processes only the observation part of a transition through the pipeline. Args: @@ -1440,7 +1440,7 @@ class ObservationProcessorStep(ProcessorStep, ABC): """An abstract `ProcessorStep` that specifically targets the observation in a transition.""" @abstractmethod - def observation(self, observation: dict[str, Any]) -> dict[str, Any]: + def observation(self, observation: RobotObservation) -> RobotObservation: """Processes an observation dictionary. Subclasses must implement this method. Args: diff --git a/src/lerobot/processor/tokenizer_processor.py b/src/lerobot/processor/tokenizer_processor.py index 93e0395b9..5cd1bebb0 100644 --- a/src/lerobot/processor/tokenizer_processor.py +++ b/src/lerobot/processor/tokenizer_processor.py @@ -38,7 +38,7 @@ from lerobot.utils.constants import ( ) from lerobot.utils.import_utils import _transformers_available -from .core import EnvTransition, TransitionKey +from .core import EnvTransition, RobotObservation, TransitionKey from .pipeline import ActionProcessorStep, ObservationProcessorStep, ProcessorStepRegistry # Conditional import for type checking and lazy loading @@ -139,7 +139,7 @@ class TokenizerProcessorStep(ObservationProcessorStep): return None - def observation(self, observation: dict[str, Any]) -> dict[str, Any]: + def observation(self, observation: RobotObservation) -> RobotObservation: """ Tokenizes the task description and adds it to the observation dictionary. diff --git a/src/lerobot/rl/gym_manipulator.py b/src/lerobot/rl/gym_manipulator.py index 604adb931..3d58ae18f 100644 --- a/src/lerobot/rl/gym_manipulator.py +++ b/src/lerobot/rl/gym_manipulator.py @@ -38,13 +38,12 @@ from lerobot.processor import ( GripperPenaltyProcessorStep, ImageCropResizeProcessorStep, InterventionActionProcessorStep, - JointVelocityProcessorStep, MapDeltaActionToRobotActionStep, MapTensorToDeltaActionDictStep, - MotorCurrentProcessorStep, Numpy2TorchActionProcessorStep, RewardClassifierProcessorStep, RobotActionToPolicyActionProcessorStep, + RobotObservation, TimeLimitProcessorStep, Torch2NumpyActionProcessorStep, TransitionKey, @@ -77,6 +76,8 @@ from lerobot.utils.constants import ACTION, DONE, OBS_IMAGES, OBS_STATE, REWARD from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.utils import log_say +from .joint_observations_processor import JointVelocityProcessorStep, MotorCurrentProcessorStep + logging.basicConfig(level=logging.INFO) @@ -163,7 +164,7 @@ class RobotEnv(gym.Env): self._setup_spaces() - def _get_observation(self) -> dict[str, Any]: + def _get_observation(self) -> RobotObservation: """Get current robot observation including joint positions and camera images.""" obs_dict = self.robot.get_observation() raw_joint_joint_position = {f"{name}.pos": obs_dict[f"{name}.pos"] for name in self._joint_names} @@ -220,7 +221,7 @@ class RobotEnv(gym.Env): def reset( self, *, seed: int | None = None, options: dict[str, Any] | None = None - ) -> tuple[dict[str, Any], dict[str, Any]]: + ) -> tuple[RobotObservation, dict[str, Any]]: """Reset environment to initial state. Args: @@ -249,7 +250,7 @@ class RobotEnv(gym.Env): self._raw_joint_positions = {f"{key}.pos": obs[f"{key}.pos"] for key in self._joint_names} return obs, {TeleopEvents.IS_INTERVENTION: False} - def step(self, action) -> tuple[dict[str, np.ndarray], float, bool, bool, dict[str, Any]]: + def step(self, action) -> tuple[RobotObservation, float, bool, bool, dict[str, Any]]: """Execute one environment step with given action.""" joint_targets_dict = {f"{key}.pos": action[i] for i, key in enumerate(self.robot.bus.motors.keys())} diff --git a/src/lerobot/processor/joint_observations_processor.py b/src/lerobot/rl/joint_observations_processor.py similarity index 100% rename from src/lerobot/processor/joint_observations_processor.py rename to src/lerobot/rl/joint_observations_processor.py diff --git a/src/lerobot/robots/bi_so_follower/bi_so_follower.py b/src/lerobot/robots/bi_so_follower/bi_so_follower.py index fa81e7d09..09f849772 100644 --- a/src/lerobot/robots/bi_so_follower/bi_so_follower.py +++ b/src/lerobot/robots/bi_so_follower/bi_so_follower.py @@ -16,8 +16,8 @@ import logging from functools import cached_property -from typing import Any +from lerobot.processor import RobotAction, RobotObservation from lerobot.robots.so_follower import SOFollower, SOFollowerRobotConfig from ..robot import Robot @@ -116,7 +116,7 @@ class BiSOFollower(Robot): self.left_arm.setup_motors() self.right_arm.setup_motors() - def get_observation(self) -> dict[str, Any]: + def get_observation(self) -> RobotObservation: obs_dict = {} # Add "left_" prefix @@ -129,7 +129,7 @@ class BiSOFollower(Robot): return obs_dict - def send_action(self, action: dict[str, Any]) -> dict[str, Any]: + def send_action(self, action: RobotAction) -> RobotAction: # Remove "left_" prefix left_action = { key.removeprefix("left_"): value for key, value in action.items() if key.startswith("left_") diff --git a/src/lerobot/robots/earthrover_mini_plus/robot_earthrover_mini_plus.py b/src/lerobot/robots/earthrover_mini_plus/robot_earthrover_mini_plus.py index 48cb09215..784a95577 100644 --- a/src/lerobot/robots/earthrover_mini_plus/robot_earthrover_mini_plus.py +++ b/src/lerobot/robots/earthrover_mini_plus/robot_earthrover_mini_plus.py @@ -18,12 +18,12 @@ import base64 import logging from functools import cached_property -from typing import Any import cv2 import numpy as np import requests +from lerobot.processor import RobotAction, RobotObservation from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError from ..robot import Robot @@ -197,11 +197,11 @@ class EarthRoverMiniPlus(Robot): ACTION_ANGULAR_VEL: float, } - def get_observation(self) -> dict[str, Any]: + def get_observation(self) -> RobotObservation: """Get current robot observation from SDK. Returns: - dict: Observation containing: + RobotObservation: Observation containing: - front: Front camera image (480, 640, 3) in RGB format - rear: Rear camera image (480, 640, 3) in RGB format - linear.vel: Current speed (0-1, SDK reports only positive speeds) @@ -255,7 +255,7 @@ class EarthRoverMiniPlus(Robot): return observation - def send_action(self, action: dict[str, Any]) -> dict[str, Any]: + def send_action(self, action: RobotAction) -> RobotAction: """Send action to robot via SDK. Args: @@ -264,8 +264,7 @@ class EarthRoverMiniPlus(Robot): - angular.vel: Target angular velocity (-1 to 1) Returns: - dict: The action that was sent (matches action_features keys) - + RobotAction: The action that was sent (matches action_features keys) Raises: DeviceNotConnectedError: If robot is not connected diff --git a/src/lerobot/robots/hope_jr/hope_jr_arm.py b/src/lerobot/robots/hope_jr/hope_jr_arm.py index 220a29f8c..4be8a0b17 100644 --- a/src/lerobot/robots/hope_jr/hope_jr_arm.py +++ b/src/lerobot/robots/hope_jr/hope_jr_arm.py @@ -17,7 +17,6 @@ import logging import time from functools import cached_property -from typing import Any from lerobot.cameras.utils import make_cameras_from_configs from lerobot.motors import Motor, MotorNormMode @@ -25,6 +24,7 @@ from lerobot.motors.calibration_gui import RangeFinderGUI from lerobot.motors.feetech import ( FeetechMotorsBus, ) +from lerobot.processor import RobotAction, RobotObservation from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError from ..robot import Robot @@ -128,7 +128,7 @@ class HopeJrArm(Robot): self.bus.setup_motor(motor) print(f"'{motor}' motor id set to {self.bus.motors[motor].id}") - def get_observation(self) -> dict[str, Any]: + def get_observation(self) -> RobotObservation: if not self.is_connected: raise DeviceNotConnectedError(f"{self} is not connected.") @@ -149,7 +149,7 @@ class HopeJrArm(Robot): return obs_dict - def send_action(self, action: dict[str, Any]) -> dict[str, Any]: + def send_action(self, action: RobotAction) -> RobotAction: if not self.is_connected: raise DeviceNotConnectedError(f"{self} is not connected.") diff --git a/src/lerobot/robots/hope_jr/hope_jr_hand.py b/src/lerobot/robots/hope_jr/hope_jr_hand.py index 9e960642b..73fb4464f 100644 --- a/src/lerobot/robots/hope_jr/hope_jr_hand.py +++ b/src/lerobot/robots/hope_jr/hope_jr_hand.py @@ -17,7 +17,6 @@ import logging import time from functools import cached_property -from typing import Any from lerobot.cameras.utils import make_cameras_from_configs from lerobot.motors import Motor, MotorNormMode @@ -25,6 +24,7 @@ from lerobot.motors.calibration_gui import RangeFinderGUI from lerobot.motors.feetech import ( FeetechMotorsBus, ) +from lerobot.processor import RobotAction, RobotObservation from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError from ..robot import Robot @@ -159,7 +159,7 @@ class HopeJrHand(Robot): self.bus.setup_motor(motor) print(f"'{motor}' motor id set to {self.bus.motors[motor].id}") - def get_observation(self) -> dict[str, Any]: + def get_observation(self) -> RobotObservation: if not self.is_connected: raise DeviceNotConnectedError(f"{self} is not connected.") @@ -181,7 +181,7 @@ class HopeJrHand(Robot): return obs_dict - def send_action(self, action: dict[str, Any]) -> dict[str, Any]: + def send_action(self, action: RobotAction) -> RobotAction: if not self.is_connected: raise DeviceNotConnectedError(f"{self} is not connected.") diff --git a/src/lerobot/robots/koch_follower/koch_follower.py b/src/lerobot/robots/koch_follower/koch_follower.py index 41a57828b..a1d001ba8 100644 --- a/src/lerobot/robots/koch_follower/koch_follower.py +++ b/src/lerobot/robots/koch_follower/koch_follower.py @@ -17,7 +17,6 @@ import logging import time from functools import cached_property -from typing import Any from lerobot.cameras.utils import make_cameras_from_configs from lerobot.motors import Motor, MotorCalibration, MotorNormMode @@ -25,6 +24,7 @@ from lerobot.motors.dynamixel import ( DynamixelMotorsBus, OperatingMode, ) +from lerobot.processor import RobotAction, RobotObservation from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError from ..robot import Robot @@ -182,7 +182,7 @@ class KochFollower(Robot): self.bus.setup_motor(motor) print(f"'{motor}' motor id set to {self.bus.motors[motor].id}") - def get_observation(self) -> dict[str, Any]: + def get_observation(self) -> RobotObservation: if not self.is_connected: raise DeviceNotConnectedError(f"{self} is not connected.") @@ -202,7 +202,7 @@ class KochFollower(Robot): return obs_dict - def send_action(self, action: dict[str, float]) -> dict[str, float]: + def send_action(self, action: RobotAction) -> RobotAction: """Command arm to move to a target joint configuration. The relative action magnitude may be clipped depending on the configuration parameter @@ -210,10 +210,10 @@ class KochFollower(Robot): Thus, this function always returns the action actually sent. Args: - action (dict[str, float]): The goal positions for the motors. + action (RobotAction): The goal positions for the motors. Returns: - dict[str, float]: The action sent to the motors, potentially clipped. + RobotAction: The action sent to the motors, potentially clipped. """ if not self.is_connected: raise DeviceNotConnectedError(f"{self} is not connected.") diff --git a/src/lerobot/robots/lekiwi/lekiwi.py b/src/lerobot/robots/lekiwi/lekiwi.py index 86fe017d6..c84e81001 100644 --- a/src/lerobot/robots/lekiwi/lekiwi.py +++ b/src/lerobot/robots/lekiwi/lekiwi.py @@ -28,6 +28,7 @@ from lerobot.motors.feetech import ( FeetechMotorsBus, OperatingMode, ) +from lerobot.processor import RobotAction, RobotObservation from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError from ..robot import Robot @@ -338,7 +339,7 @@ class LeKiwi(Robot): "theta.vel": theta, } # m/s and deg/s - def get_observation(self) -> dict[str, Any]: + def get_observation(self) -> RobotObservation: if not self.is_connected: raise DeviceNotConnectedError(f"{self} is not connected.") @@ -369,7 +370,7 @@ class LeKiwi(Robot): return obs_dict - def send_action(self, action: dict[str, Any]) -> dict[str, Any]: + def send_action(self, action: RobotAction) -> RobotAction: """Command lekiwi to move to a target joint configuration. The relative action magnitude may be clipped depending on the configuration parameter @@ -380,7 +381,7 @@ class LeKiwi(Robot): RobotDeviceNotConnectedError: if robot is not connected. Returns: - np.ndarray: the action sent to the motors, potentially clipped. + RobotAction: the action sent to the motors, potentially clipped. """ if not self.is_connected: raise DeviceNotConnectedError(f"{self} is not connected.") diff --git a/src/lerobot/robots/lekiwi/lekiwi_client.py b/src/lerobot/robots/lekiwi/lekiwi_client.py index 19744e244..bb865dc10 100644 --- a/src/lerobot/robots/lekiwi/lekiwi_client.py +++ b/src/lerobot/robots/lekiwi/lekiwi_client.py @@ -18,11 +18,11 @@ import base64 import json import logging from functools import cached_property -from typing import Any import cv2 import numpy as np +from lerobot.processor import RobotAction, RobotObservation from lerobot.utils.constants import ACTION, OBS_STATE from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError @@ -172,7 +172,7 @@ class LeKiwiClient(Robot): return last_msg - def _parse_observation_json(self, obs_string: str) -> dict[str, Any] | None: + def _parse_observation_json(self, obs_string: str) -> RobotObservation | None: """Parses the JSON observation string.""" try: return json.loads(obs_string) @@ -196,15 +196,15 @@ class LeKiwiClient(Robot): return None def _remote_state_from_obs( - self, observation: dict[str, Any] - ) -> tuple[dict[str, np.ndarray], dict[str, Any]]: + self, observation: RobotObservation + ) -> tuple[dict[str, np.ndarray], RobotObservation]: """Extracts frames, and state from the parsed observation.""" flat_state = {key: observation.get(key, 0.0) for key in self._state_order} state_vec = np.array([flat_state[key] for key in self._state_order], dtype=np.float32) - obs_dict: dict[str, Any] = {**flat_state, OBS_STATE: state_vec} + obs_dict: RobotObservation = {**flat_state, OBS_STATE: state_vec} # Decode images current_frames: dict[str, np.ndarray] = {} @@ -217,7 +217,7 @@ class LeKiwiClient(Robot): return current_frames, obs_dict - def _get_data(self) -> tuple[dict[str, np.ndarray], dict[str, Any], dict[str, Any]]: + def _get_data(self) -> tuple[dict[str, np.ndarray], RobotObservation]: """ Polls the video socket for the latest observation data. @@ -252,7 +252,7 @@ class LeKiwiClient(Robot): return new_frames, new_state - def get_observation(self) -> dict[str, Any]: + def get_observation(self) -> RobotObservation: """ Capture observations from the remote robot: current follower arm positions, present wheel speeds (converted to body-frame velocities: x, y, theta), @@ -307,12 +307,11 @@ class LeKiwiClient(Robot): def configure(self): pass - def send_action(self, action: dict[str, Any]) -> dict[str, Any]: + def send_action(self, action: RobotAction) -> RobotAction: """Command lekiwi to move to a target joint configuration. Translates to motor space + sends over ZMQ Args: - action (np.ndarray): array containing the goal positions for the motors. - + action (RobotAction): array containing the goal positions for the motors. Raises: RobotDeviceNotConnectedError: if robot is not connected. diff --git a/src/lerobot/robots/omx_follower/omx_follower.py b/src/lerobot/robots/omx_follower/omx_follower.py index 2dd851377..14668b3a7 100644 --- a/src/lerobot/robots/omx_follower/omx_follower.py +++ b/src/lerobot/robots/omx_follower/omx_follower.py @@ -17,7 +17,6 @@ import logging import time from functools import cached_property -from typing import Any from lerobot.cameras.utils import make_cameras_from_configs from lerobot.motors import Motor, MotorCalibration, MotorNormMode @@ -26,6 +25,7 @@ from lerobot.motors.dynamixel import ( DynamixelMotorsBus, OperatingMode, ) +from lerobot.processor import RobotAction, RobotObservation from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError from ..robot import Robot @@ -165,7 +165,7 @@ class OmxFollower(Robot): self.bus.setup_motor(motor) print(f"'{motor}' motor id set to {self.bus.motors[motor].id}") - def get_observation(self) -> dict[str, Any]: + def get_observation(self) -> RobotObservation: if not self.is_connected: raise DeviceNotConnectedError(f"{self} is not connected.") @@ -185,7 +185,7 @@ class OmxFollower(Robot): return obs_dict - def send_action(self, action: dict[str, float]) -> dict[str, float]: + def send_action(self, action: RobotAction) -> RobotAction: """Command arm to move to a target joint configuration. The relative action magnitude may be clipped depending on the configuration parameter @@ -193,10 +193,10 @@ class OmxFollower(Robot): Thus, this function always returns the action actually sent. Args: - action (dict[str, float]): The goal positions for the motors. + action (RobotAction): The goal positions for the motors. Returns: - dict[str, float]: The action sent to the motors, potentially clipped. + RobotAction: The action sent to the motors, potentially clipped. """ if not self.is_connected: raise DeviceNotConnectedError(f"{self} is not connected.") diff --git a/src/lerobot/robots/reachy2/robot_reachy2.py b/src/lerobot/robots/reachy2/robot_reachy2.py index 74742ee8d..6f4eef56c 100644 --- a/src/lerobot/robots/reachy2/robot_reachy2.py +++ b/src/lerobot/robots/reachy2/robot_reachy2.py @@ -18,9 +18,8 @@ from __future__ import annotations import time from typing import TYPE_CHECKING, Any -import numpy as np - from lerobot.cameras.utils import make_cameras_from_configs +from lerobot.processor import RobotAction, RobotObservation from lerobot.utils.import_utils import _reachy2_sdk_available from ..robot import Robot @@ -171,8 +170,8 @@ class Reachy2Robot(Robot): else: return {} - def get_observation(self) -> dict[str, np.ndarray]: - obs_dict: dict[str, Any] = {} + def get_observation(self) -> RobotObservation: + obs_dict: RobotObservation = {} # Read Reachy 2 state before_read_t = time.perf_counter() @@ -185,7 +184,7 @@ class Reachy2Robot(Robot): return obs_dict - def send_action(self, action: dict[str, Any]) -> dict[str, Any]: + def send_action(self, action: RobotAction) -> RobotAction: if self.reachy is not None: if not self.is_connected: raise ConnectionError() diff --git a/src/lerobot/robots/robot.py b/src/lerobot/robots/robot.py index 5e88b915b..d1021daf4 100644 --- a/src/lerobot/robots/robot.py +++ b/src/lerobot/robots/robot.py @@ -15,11 +15,11 @@ import abc import builtins from pathlib import Path -from typing import Any import draccus from lerobot.motors import MotorCalibration +from lerobot.processor import RobotAction, RobotObservation from lerobot.utils.constants import HF_LEROBOT_CALIBRATION, ROBOTS from .config import RobotConfig @@ -153,28 +153,28 @@ class Robot(abc.ABC): pass @abc.abstractmethod - def get_observation(self) -> dict[str, Any]: + def get_observation(self) -> RobotObservation: """ Retrieve the current observation from the robot. Returns: - dict[str, Any]: A flat dictionary representing the robot's current sensory state. Its structure + RobotObservation: A flat dictionary representing the robot's current sensory state. Its structure should match :pymeth:`observation_features`. """ pass @abc.abstractmethod - def send_action(self, action: dict[str, Any]) -> dict[str, Any]: + def send_action(self, action: RobotAction) -> RobotAction: """ Send an action command to the robot. Args: - action (dict[str, Any]): Dictionary representing the desired action. Its structure should match + action (RobotAction): Dictionary representing the desired action. Its structure should match :pymeth:`action_features`. Returns: - dict[str, Any]: The action actually sent to the motors potentially clipped or modified, e.g. by + RobotAction: The action actually sent to the motors potentially clipped or modified, e.g. by safety limits on velocity. """ pass diff --git a/src/lerobot/robots/so_follower/robot_kinematic_processor.py b/src/lerobot/robots/so_follower/robot_kinematic_processor.py index 87e832db6..2aa60e12a 100644 --- a/src/lerobot/robots/so_follower/robot_kinematic_processor.py +++ b/src/lerobot/robots/so_follower/robot_kinematic_processor.py @@ -28,6 +28,7 @@ from lerobot.processor import ( ProcessorStepRegistry, RobotAction, RobotActionProcessorStep, + RobotObservation, TransitionKey, ) from lerobot.utils.rotation import Rotation @@ -438,7 +439,7 @@ class ForwardKinematicsJointsToEEObservation(ObservationProcessorStep): kinematics: RobotKinematics motor_names: list[str] - def observation(self, observation: dict[str, Any]) -> dict[str, Any]: + def observation(self, observation: RobotObservation) -> RobotObservation: return compute_forward_kinematics_joints_to_ee(observation, self.kinematics, self.motor_names) def transform_features( diff --git a/src/lerobot/robots/so_follower/so_follower.py b/src/lerobot/robots/so_follower/so_follower.py index 5e99b33a1..011a0061e 100644 --- a/src/lerobot/robots/so_follower/so_follower.py +++ b/src/lerobot/robots/so_follower/so_follower.py @@ -17,7 +17,7 @@ import logging import time from functools import cached_property -from typing import Any, TypeAlias +from typing import TypeAlias from lerobot.cameras.utils import make_cameras_from_configs from lerobot.motors import Motor, MotorCalibration, MotorNormMode @@ -25,6 +25,7 @@ from lerobot.motors.feetech import ( FeetechMotorsBus, OperatingMode, ) +from lerobot.processor import RobotAction, RobotObservation from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError from ..robot import Robot @@ -175,7 +176,7 @@ class SOFollower(Robot): self.bus.setup_motor(motor) print(f"'{motor}' motor id set to {self.bus.motors[motor].id}") - def get_observation(self) -> dict[str, Any]: + def get_observation(self) -> RobotObservation: if not self.is_connected: raise DeviceNotConnectedError(f"{self} is not connected.") @@ -195,7 +196,7 @@ class SOFollower(Robot): return obs_dict - def send_action(self, action: dict[str, Any]) -> dict[str, Any]: + def send_action(self, action: RobotAction) -> RobotAction: """Command arm to move to a target joint configuration. The relative action magnitude may be clipped depending on the configuration parameter @@ -206,7 +207,7 @@ class SOFollower(Robot): RobotDeviceNotConnectedError: if robot is not connected. Returns: - the action sent to the motors, potentially clipped. + RobotAction: the action sent to the motors, potentially clipped. """ if not self.is_connected: raise DeviceNotConnectedError(f"{self} is not connected.") diff --git a/src/lerobot/robots/unitree_g1/unitree_g1.py b/src/lerobot/robots/unitree_g1/unitree_g1.py index 1764f31b5..fa6e0da85 100644 --- a/src/lerobot/robots/unitree_g1/unitree_g1.py +++ b/src/lerobot/robots/unitree_g1/unitree_g1.py @@ -26,6 +26,7 @@ import numpy as np from lerobot.cameras.utils import make_cameras_from_configs from lerobot.envs.factory import make_env +from lerobot.processor import RobotAction, RobotObservation from lerobot.robots.unitree_g1.g1_utils import G1_29_JointIndex from ..robot import Robot @@ -269,7 +270,7 @@ class UnitreeG1(Robot): for cam in self._cameras.values(): cam.disconnect() - def get_observation(self) -> dict[str, Any]: + def get_observation(self) -> RobotObservation: lowstate = self._lowstate if lowstate is None: return {} @@ -350,7 +351,7 @@ class UnitreeG1(Robot): def observation_features(self) -> dict[str, type | tuple]: return {**self._motors_ft, **self._cameras_ft} - def send_action(self, action: dict[str, Any]) -> dict[str, Any]: + def send_action(self, action: RobotAction) -> RobotAction: for motor in G1_29_JointIndex: key = f"{motor.name}.q" if key in action: diff --git a/src/lerobot/scripts/lerobot_eval.py b/src/lerobot/scripts/lerobot_eval.py index 7b45e88e1..e32b80404 100644 --- a/src/lerobot/scripts/lerobot_eval.py +++ b/src/lerobot/scripts/lerobot_eval.py @@ -177,9 +177,9 @@ def rollout( action = policy.select_action(observation) action = postprocessor(action) - action_transition = {"action": action} + action_transition = {ACTION: action} action_transition = env_postprocessor(action_transition) - action = action_transition["action"] + action = action_transition[ACTION] # Convert to CPU / numpy. action_numpy: np.ndarray = action.to("cpu").numpy() diff --git a/src/lerobot/scripts/lerobot_teleoperate.py b/src/lerobot/scripts/lerobot_teleoperate.py index 05d4534d4..18d8863d6 100644 --- a/src/lerobot/scripts/lerobot_teleoperate.py +++ b/src/lerobot/scripts/lerobot_teleoperate.py @@ -165,7 +165,7 @@ def teleop_loop( # Process action for robot through pipeline robot_action_to_send = robot_action_processor((teleop_action, obs)) - # Send processed action to robot (robot_action_processor.to_output should return dict[str, Any]) + # Send processed action to robot (robot_action_processor.to_output should return RobotAction) _ = robot.send_action(robot_action_to_send) if display_data: diff --git a/src/lerobot/scripts/lerobot_train_tokenizer.py b/src/lerobot/scripts/lerobot_train_tokenizer.py index 296447bad..1d8f4644b 100644 --- a/src/lerobot/scripts/lerobot_train_tokenizer.py +++ b/src/lerobot/scripts/lerobot_train_tokenizer.py @@ -63,6 +63,7 @@ else: from lerobot.configs import parser from lerobot.configs.types import NormalizationMode from lerobot.datasets.lerobot_dataset import LeRobotDataset +from lerobot.utils.constants import ACTION, OBS_STATE @dataclass @@ -86,7 +87,7 @@ class TokenizerTrainingConfig: # Whether to apply delta transform (relative actions vs absolute actions) use_delta_transform: bool = False # Dataset key for state observations (default: "observation.state") - state_key: str = "observation.state" + state_key: str = OBS_STATE # Normalization mode (MEAN_STD, MIN_MAX, QUANTILES, QUANTILE10, IDENTITY) normalization_mode: str = "QUANTILES" # FAST vocabulary size (BPE vocab size) @@ -223,12 +224,10 @@ def process_episode(args): else: # if no state key, use zeros (no delta transform) state = np.zeros_like( - frame["action"].numpy() if torch.is_tensor(frame["action"]) else np.array(frame["action"]) + frame[ACTION].numpy() if torch.is_tensor(frame[ACTION]) else np.array(frame[ACTION]) ) - action = ( - frame["action"].numpy() if torch.is_tensor(frame["action"]) else np.array(frame["action"]) - ) + action = frame[ACTION].numpy() if torch.is_tensor(frame[ACTION]) else np.array(frame[ACTION]) states.append(state) actions.append(action) @@ -468,8 +467,8 @@ def train_tokenizer(cfg: TokenizerTrainingConfig): # get normalization stats from dataset norm_stats = dataset.meta.stats - if norm_stats is not None and "action" in norm_stats: - action_stats = norm_stats["action"] + if norm_stats is not None and ACTION in norm_stats: + action_stats = norm_stats[ACTION] # build encoded dimension indices encoded_dim_indices = [] diff --git a/src/lerobot/teleoperators/gamepad/teleop_gamepad.py b/src/lerobot/teleoperators/gamepad/teleop_gamepad.py index c7072f4a7..4dbb49c1d 100644 --- a/src/lerobot/teleoperators/gamepad/teleop_gamepad.py +++ b/src/lerobot/teleoperators/gamepad/teleop_gamepad.py @@ -20,6 +20,8 @@ from typing import Any import numpy as np +from lerobot.processor import RobotAction + from ..teleoperator import Teleoperator from ..utils import TeleopEvents from .configuration_gamepad import GamepadTeleopConfig @@ -83,7 +85,7 @@ class GamepadTeleop(Teleoperator): self.gamepad = Gamepad() self.gamepad.start() - def get_action(self) -> dict[str, Any]: + def get_action(self) -> RobotAction: # Update the controller to get fresh inputs self.gamepad.update() diff --git a/src/lerobot/teleoperators/keyboard/teleop_keyboard.py b/src/lerobot/teleoperators/keyboard/teleop_keyboard.py index ec8ea18f4..55c158da8 100644 --- a/src/lerobot/teleoperators/keyboard/teleop_keyboard.py +++ b/src/lerobot/teleoperators/keyboard/teleop_keyboard.py @@ -21,6 +21,7 @@ import time from queue import Queue from typing import Any +from lerobot.processor import RobotAction from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError from ..teleoperator import Teleoperator @@ -124,7 +125,7 @@ class KeyboardTeleop(Teleoperator): def configure(self): pass - def get_action(self) -> dict[str, Any]: + def get_action(self) -> RobotAction: before_read_t = time.perf_counter() if not self.is_connected: @@ -181,7 +182,7 @@ class KeyboardEndEffectorTeleop(KeyboardTeleop): "names": {"delta_x": 0, "delta_y": 1, "delta_z": 2}, } - def get_action(self) -> dict[str, Any]: + def get_action(self) -> RobotAction: if not self.is_connected: raise DeviceNotConnectedError( "KeyboardTeleop is not connected. You need to run `connect()` before `get_action()`." @@ -374,12 +375,12 @@ class KeyboardRoverTeleop(KeyboardTeleop): # Only remove key if it's being released self.current_pressed.pop(key_char, None) - def get_action(self) -> dict[str, Any]: + def get_action(self) -> RobotAction: """ Get the current action based on pressed keys. Returns: - dict with 'linear.vel' and 'angular.vel' keys + RobotAction with 'linear.vel' and 'angular.vel' keys """ before_read_t = time.perf_counter() diff --git a/src/lerobot/teleoperators/teleoperator.py b/src/lerobot/teleoperators/teleoperator.py index 95020a962..cd9e3a53d 100644 --- a/src/lerobot/teleoperators/teleoperator.py +++ b/src/lerobot/teleoperators/teleoperator.py @@ -20,6 +20,7 @@ from typing import Any import draccus from lerobot.motors.motors_bus import MotorCalibration +from lerobot.processor import RobotAction from lerobot.utils.constants import HF_LEROBOT_CALIBRATION, TELEOPERATORS from .config import TeleoperatorConfig @@ -150,12 +151,12 @@ class Teleoperator(abc.ABC): pass @abc.abstractmethod - def get_action(self) -> dict[str, Any]: + def get_action(self) -> RobotAction: """ Retrieve the current action from the teleoperator. Returns: - dict[str, Any]: A flat dictionary representing the teleoperator's current actions. Its + RobotAction: A flat dictionary representing the teleoperator's current actions. Its structure should match :pymeth:`observation_features`. """ pass diff --git a/src/lerobot/utils/constants.py b/src/lerobot/utils/constants.py index a96e1596d..43a61b4f7 100644 --- a/src/lerobot/utils/constants.py +++ b/src/lerobot/utils/constants.py @@ -28,11 +28,13 @@ OBS_LANGUAGE_TOKENS = OBS_LANGUAGE + ".tokens" OBS_LANGUAGE_ATTENTION_MASK = OBS_LANGUAGE + ".attention_mask" ACTION = "action" +ACTION_PREFIX = ACTION + "." ACTION_TOKENS = ACTION + ".tokens" ACTION_TOKEN_MASK = ACTION + ".token_mask" REWARD = "next.reward" TRUNCATED = "next.truncated" DONE = "next.done" +INFO = "info" ROBOTS = "robots" TELEOPERATORS = "teleoperators" diff --git a/src/lerobot/utils/visualization_utils.py b/src/lerobot/utils/visualization_utils.py index 9143d0f66..31ca8d247 100644 --- a/src/lerobot/utils/visualization_utils.py +++ b/src/lerobot/utils/visualization_utils.py @@ -14,12 +14,13 @@ import numbers import os -from typing import Any import numpy as np import rerun as rr -from .constants import OBS_PREFIX, OBS_STR +from lerobot.processor import RobotAction, RobotObservation + +from .constants import ACTION, ACTION_PREFIX, OBS_PREFIX, OBS_STR def init_rerun( @@ -50,8 +51,8 @@ def _is_scalar(x): def log_rerun_data( - observation: dict[str, Any] | None = None, - action: dict[str, Any] | None = None, + observation: RobotObservation | None = None, + action: RobotAction | None = None, compress_images: bool = False, ) -> None: """ @@ -96,7 +97,7 @@ def log_rerun_data( for k, v in action.items(): if v is None: continue - key = k if str(k).startswith("action.") else f"action.{k}" + key = k if str(k).startswith(ACTION_PREFIX) else f"{ACTION}.{k}" if _is_scalar(v): rr.log(key, rr.Scalars(float(v))) diff --git a/tests/mocks/mock_robot.py b/tests/mocks/mock_robot.py index b0513fd38..d997cb6d4 100644 --- a/tests/mocks/mock_robot.py +++ b/tests/mocks/mock_robot.py @@ -17,10 +17,10 @@ import random from dataclasses import dataclass, field from functools import cached_property -from typing import Any from lerobot.cameras import CameraConfig, make_cameras_from_configs from lerobot.motors.motors_bus import Motor, MotorNormMode +from lerobot.processor import RobotAction, RobotObservation from lerobot.robots import Robot, RobotConfig from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError from tests.mocks.mock_motors_bus import MockMotorsBus @@ -119,7 +119,7 @@ class MockRobot(Robot): def configure(self) -> None: pass - def get_observation(self) -> dict[str, Any]: + def get_observation(self) -> RobotObservation: if not self.is_connected: raise DeviceNotConnectedError(f"{self} is not connected.") @@ -130,7 +130,7 @@ class MockRobot(Robot): f"{motor}.pos": val for motor, val in zip(self.motors, self.config.static_values, strict=True) } - def send_action(self, action: dict[str, Any]) -> dict[str, Any]: + def send_action(self, action: RobotAction) -> RobotAction: if not self.is_connected: raise DeviceNotConnectedError(f"{self} is not connected.") diff --git a/tests/mocks/mock_teleop.py b/tests/mocks/mock_teleop.py index 71b49947c..04479bad9 100644 --- a/tests/mocks/mock_teleop.py +++ b/tests/mocks/mock_teleop.py @@ -19,6 +19,7 @@ from dataclasses import dataclass from functools import cached_property from typing import Any +from lerobot.processor import RobotAction from lerobot.teleoperators import Teleoperator, TeleoperatorConfig from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError @@ -88,7 +89,7 @@ class MockTeleop(Teleoperator): def configure(self) -> None: pass - def get_action(self) -> dict[str, Any]: + def get_action(self) -> RobotAction: if not self.is_connected: raise DeviceNotConnectedError(f"{self} is not connected.")