admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-25 10:46:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' into refactor/lerobot_train_rabc

Browse Source
This commit is contained in:
Michel Aractingi
2026-01-14 16:09:26 +01:00
committed by GitHub
parent 25e3119a33 15724826dd
commit fd7cb9a5c5
74 changed files with 1156 additions and 646 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/source/reachy2.mdx
+34 -19
View File
@@ -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_<part>=false
```,
```
with `<part>` 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\_<part>
```
with `<part>` 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=<true|false> \
--robot.with_right_teleop_camera=<true|false> \
--robot.with_torso_camera=<true|false>
```
--robot.with_left_teleop_camera=<true|false>
--robot.with_right_teleop_camera=<true|false>
--robot.with_torso_camera=<true|false>
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 \
docs/source/unitree_g1.mdx
+3 -3
View File
@@ -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@<YOUR_ROBOT_IP>
# Password: 123
```
Replace `<YOUR_ROBOT_IP>` with your robot's actual WiFi IP address (e.g., `172.18.129.215`).
Replace `<YOUR_ROBOT_IP>` 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
examples/unitree_g1/gr00t_locomotion.py
+18 -21
View File
@@ -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
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 self.robot.remote_controller.button[4]: # R2 - lower waist
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)
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
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
examples/unitree_g1/holosoma_locomotion.py
+14 -14
View File
@@ -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
pyproject.toml
+1 -1
View File
@@ -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'",
src/lerobot/async_inference/robot_client.py
+2 -2
View File
@@ -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
src/lerobot/cameras/reachy2_camera/configuration_reachy2_camera.py
+4 -4
View File
@@ -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"]:
src/lerobot/cameras/reachy2_camera/reachy2_camera.py
+42 -121
View File
@@ -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,26 +137,32 @@ 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.")
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]
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=(640, 480))[0]
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=(640, 480))[0]
frame = self.cam_manager.depth.get_frame(size=(self.config.width, self.config.height))[0]
else:
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)
@@ -187,63 +175,11 @@ class Reachy2Camera(Camera):
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()
src/lerobot/cameras/utils.py
+5
View File
@@ -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))
src/lerobot/cameras/zmq/__init__.py
+20
View File
@@ -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"]
src/lerobot/cameras/zmq/camera_zmq.py
+235
View File
@@ -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": "<base64-jpeg>"}
}
"""
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.")
src/lerobot/cameras/zmq/configuration_zmq.py
+46
View File
@@ -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.")
src/lerobot/cameras/zmq/image_server.py
+114
View File
@@ -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()
src/lerobot/datasets/pipeline_features.py
+2 -2
View File
@@ -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.
src/lerobot/envs/libero.py
+4 -2
View File
@@ -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,)), "
src/lerobot/envs/metaworld.py
+7 -5
View File
@@ -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.
src/lerobot/envs/utils.py
+2 -1
View File
@@ -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")
src/lerobot/policies/factory.py
+6 -2
View File
@@ -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,
src/lerobot/policies/groot/configuration_groot.py
+7 -6
View File
@@ -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(
src/lerobot/policies/groot/groot_n1.py
+3 -3
View File
@@ -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
src/lerobot/policies/groot/modeling_groot.py
+2 -1
View File
@@ -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
src/lerobot/policies/groot/processor_groot.py
+16 -12
View File
@@ -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
src/lerobot/policies/pi0/configuration_pi0.py
+5 -5
View File
@@ -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(
src/lerobot/policies/pi05/configuration_pi05.py
+6 -5
View File
@@ -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(
src/lerobot/policies/pi0_fast/configuration_pi0_fast.py
+6 -5
View File
@@ -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(
src/lerobot/policies/sarm/configuration_sarm.py
+3 -2
View File
@@ -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: {})
src/lerobot/policies/sarm/modeling_sarm.py
+2 -1
View File
@@ -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)
src/lerobot/policies/utils.py
+1 -2
View File
@@ -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,
src/lerobot/policies/wall_x/configuration_wall_x.py
+7 -6
View File
@@ -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(
src/lerobot/policies/wall_x/modeling_wall_x.py
+3 -3
View File
@@ -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
src/lerobot/policies/xvla/processor_xvla.py
+9 -7
View File
@@ -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:
src/lerobot/processor/__init__.py
-3
View File
@@ -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",
src/lerobot/processor/converters.py
+3 -3
View File
@@ -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.
src/lerobot/processor/core.py
+1 -1
View File
@@ -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,
src/lerobot/processor/env_processor.py
+6 -5
View File
@@ -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)."),
src/lerobot/processor/normalize_processor.py
+2 -2
View File
@@ -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.
src/lerobot/processor/pipeline.py
+3 -3
View File
@@ -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:
src/lerobot/processor/tokenizer_processor.py
+2 -2
View File
@@ -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.
src/lerobot/rl/gym_manipulator.py
+6 -5
View File
@@ -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())}
src/lerobot/processor/joint_observations_processor.py → src/lerobot/rl/joint_observations_processor.py
View File
src/lerobot/robots/bi_so_follower/bi_so_follower.py
+3 -3
View File
@@ -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_")
src/lerobot/robots/earthrover_mini_plus/robot_earthrover_mini_plus.py
+5 -6
View File
@@ -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
src/lerobot/robots/hope_jr/hope_jr_arm.py
+3 -3
View File
@@ -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.")
src/lerobot/robots/hope_jr/hope_jr_hand.py
+3 -3
View File
@@ -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.")
src/lerobot/robots/koch_follower/koch_follower.py
+5 -5
View File
@@ -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.")
src/lerobot/robots/lekiwi/lekiwi.py
+4 -3
View File
@@ -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.")
src/lerobot/robots/lekiwi/lekiwi_client.py
+9 -10
View File
@@ -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.
src/lerobot/robots/omx_follower/omx_follower.py
+5 -5
View File
@@ -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.")
src/lerobot/robots/reachy2/configuration_reachy2.py
+22 -12
View File
@@ -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,
)
src/lerobot/robots/reachy2/robot_reachy2.py
+12 -7
View File
@@ -13,19 +13,24 @@
# 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
import numpy as np
from reachy2_sdk import ReachySDK
from typing import TYPE_CHECKING, Any
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
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",
@@ -165,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()
@@ -179,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()
src/lerobot/robots/robot.py
+6 -6
View File
@@ -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
src/lerobot/robots/so_follower/robot_kinematic_processor.py
+2 -1
View File
@@ -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(
src/lerobot/robots/so_follower/so_follower.py
+5 -4
View File
@@ -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.")
src/lerobot/robots/unitree_g1/config_unitree_g1.py
+5
View File
@@ -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)
src/lerobot/robots/unitree_g1/unitree_g1.py
+189 -88
View File
@@ -23,14 +23,10 @@ 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.processor import RobotAction, RobotObservation
from lerobot.robots.unitree_g1.g1_utils import G1_29_JointIndex
from ..robot import Robot
@@ -43,8 +39,6 @@ logger = logging.getLogger(__name__)
kTopicLowCommand_Debug = "rt/lowcmd"
kTopicLowState = "rt/lowstate"
G1_29_Num_Motors = 29
@dataclass
class MotorState:
@@ -66,28 +60,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 +95,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 +114,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 +159,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 +168,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 +177,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()
def get_observation(self) -> dict[str, Any]:
return self.lowstate_buffer.get_data()
# 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) -> RobotObservation:
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 +331,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]:
@@ -262,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:
@@ -293,22 +382,34 @@ 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)
if self.config.is_simulation and self.sim_env is not None:
self.sim_env.reset()
for motor in G1_29_JointIndex:
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)
# get current state
robot_state = self.get_observation()
obs = 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
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):
src/lerobot/scripts/lerobot_eval.py
+2 -2
View File
@@ -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()
src/lerobot/scripts/lerobot_record.py
+10
View File
@@ -73,7 +73,9 @@ 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
from lerobot.configs.policies import PreTrainedConfig
from lerobot.datasets.image_writer import safe_stop_image_writer
@@ -102,7 +104,9 @@ from lerobot.robots import ( # noqa: F401
koch_follower,
make_robot_from_config,
omx_follower,
reachy2,
so_follower,
unitree_g1,
)
from lerobot.teleoperators import ( # noqa: F401
Teleoperator,
@@ -112,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
@@ -508,6 +513,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,
src/lerobot/scripts/lerobot_replay.py
+2
View File
@@ -59,7 +59,9 @@ from lerobot.robots import ( # noqa: F401
koch_follower,
make_robot_from_config,
omx_follower,
reachy2,
so_follower,
unitree_g1,
)
from lerobot.utils.constants import ACTION
from lerobot.utils.import_utils import register_third_party_plugins
src/lerobot/scripts/lerobot_teleoperate.py
+3 -1
View File
@@ -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
@@ -163,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:
src/lerobot/scripts/lerobot_train.py
+2 -2
View File
@@ -253,6 +253,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
@@ -276,8 +278,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()))
src/lerobot/scripts/lerobot_train_tokenizer.py
+110 -78
View File
@@ -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,16 +45,63 @@ 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
from lerobot.utils.constants import ACTION, OBS_STATE
@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 = OBS_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:
@@ -174,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)
@@ -327,88 +375,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 +436,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)
@@ -442,8 +467,8 @@ def main(
# 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 = []
@@ -495,16 +520,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 +538,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 +561,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 +595,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()
src/lerobot/teleoperators/gamepad/teleop_gamepad.py
+3 -1
View File
@@ -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()
src/lerobot/teleoperators/keyboard/teleop_keyboard.py
+5 -4
View File
@@ -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()
src/lerobot/teleoperators/reachy2_teleoperator/reachy2_teleoperator.py
+23 -8
View File
@@ -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 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,19 +149,20 @@ class Reachy2Teleoperator(Teleoperator):
def get_action(self) -> dict[str, float]:
start = time.perf_counter()
if self.reachy and self.is_connected:
if not self.is_connected:
raise DeviceNotConnectedError(f"{self} is not connected.")
joint_action: dict[str, float] = {}
vel_action: dict[str, float] = {}
if self.config.use_present_position:
joint_action = {
k: self.reachy.joints[v].present_position for k, v in self.joints_dict.items()
}
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:
@@ -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()
src/lerobot/teleoperators/teleoperator.py
+3 -2
View File
@@ -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
src/lerobot/utils/constants.py
+2
View File
@@ -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"
src/lerobot/utils/import_utils.py
+1
View File
@@ -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:
src/lerobot/utils/visualization_utils.py
+6 -5
View File
@@ -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)))
tests/cameras/test_reachy2_camera.py
+2 -12
View File
@@ -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()
tests/conftest.py
-1
View File
@@ -28,7 +28,6 @@ pytest_plugins = [
"tests.fixtures.files",
"tests.fixtures.hub",
"tests.fixtures.optimizers",
"tests.plugins.reachy2_sdk",
]
tests/mocks/mock_robot.py
+3 -3
View File
@@ -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.")
tests/mocks/mock_teleop.py
+2 -1
View File
@@ -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.")
tests/plugins/reachy2_sdk.py
-46
View File
@@ -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()
tests/robots/test_reachy2.py
+2
View File
@@ -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,