Merge branch 'main' into refactor/lerobot_train_rabc

This commit is contained in:
Michel Aractingi
2026-01-14 16:09:26 +01:00
committed by GitHub
74 changed files with 1156 additions and 646 deletions
+34 -19
View File
@@ -38,6 +38,7 @@ docker run --rm -it \
start_rviz:=true start_sdk_server:=true mujoco:=true 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: > 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: First add reachy2 and reachy2_teleoperator to the imports of the record script. Then you can use the following command:
```bash ```bash
python -m lerobot.record \ lerobot-record \
--robot.type=reachy2 \ --robot.type=reachy2 \
--robot.ip_address=192.168.0.200 \ --robot.ip_address=192.168.0.200 \
--robot.id=r2-0000 \ --robot.id=r2-0000 \
@@ -150,6 +151,7 @@ python -m lerobot.record \
--teleop.type=reachy2_teleoperator \ --teleop.type=reachy2_teleoperator \
--teleop.ip_address=192.168.0.200 \ --teleop.ip_address=192.168.0.200 \
--teleop.with_mobile_base=false \ --teleop.with_mobile_base=false \
--robot.with_torso_camera=true \
--dataset.repo_id=pollen_robotics/record_test \ --dataset.repo_id=pollen_robotics/record_test \
--dataset.single_task="Reachy 2 recording test" \ --dataset.single_task="Reachy 2 recording test" \
--dataset.num_episodes=1 \ --dataset.num_episodes=1 \
@@ -165,7 +167,7 @@ python -m lerobot.record \
**Extended setup overview (all options included):** **Extended setup overview (all options included):**
```bash ```bash
python -m lerobot.record \ lerobot-record \
--robot.type=reachy2 \ --robot.type=reachy2 \
--robot.ip_address=192.168.0.200 \ --robot.ip_address=192.168.0.200 \
--robot.use_external_commands=true \ --robot.use_external_commands=true \
@@ -177,6 +179,8 @@ python -m lerobot.record \
--robot.with_left_teleop_camera=true \ --robot.with_left_teleop_camera=true \
--robot.with_right_teleop_camera=true \ --robot.with_right_teleop_camera=true \
--robot.with_torso_camera=false \ --robot.with_torso_camera=false \
--robot.camera_width=640 \
--robot.camera_height=480 \
--robot.disable_torque_on_disconnect=false \ --robot.disable_torque_on_disconnect=false \
--robot.max_relative_target=5.0 \ --robot.max_relative_target=5.0 \
--teleop.type=reachy2_teleoperator \ --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. 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: To avoid this, you can exclude specific parts from recording and replay using:
```` ```bash
--robot.with_<part>=false --robot.with_<part>=false
```, ```
with `<part>` being one of : `mobile_base`, `l_arm`, `r_arm", `neck`, `antennas`. 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. 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. The same per-part mechanism is available in `reachy2_teleoperator` as well.
```` ```bash
--teleop.with\_<part> --teleop.with\_<part>
``` ```
with `<part>` being one of : `mobile_base`, `l_arm`, `r_arm", `neck`, `antennas`. 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. 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. > **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 ##### 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> By default, no camera is recorded, all camera arguments are set to `false`.
--robot.with_right_teleop_camera=<true|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:
--robot.with_torso_camera=<true|false>
```` ```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 ## Step 2: Replay
Make sure the robot is configured with the same parts as the dataset: Make sure the robot is configured with the same parts as the dataset:
```bash ```bash
python -m lerobot.replay \ lerobot-replay \
--robot.type=reachy2 \ --robot.type=reachy2 \
--robot.ip_address=192.168.0.200 \ --robot.ip_address=192.168.0.200 \
--robot.use_external_commands=false \ --robot.use_external_commands=false \
--robot.with_mobile_base=false \ --robot.with_mobile_base=false \
--dataset.repo_id=pollen_robotics/record_test \ --dataset.repo_id=pollen_robotics/record_test \
--dataset.episode=0 --dataset.episode=0
--display_data=true ```
````
## Step 3: Train ## Step 3: Train
```bash ```bash
python -m lerobot.scripts.train \ lerobot-train \
--dataset.repo_id=pollen_robotics/record_test \ --dataset.repo_id=pollen_robotics/record_test \
--policy.type=act \ --policy.type=act \
--output_dir=outputs/train/reachy2_test \ --output_dir=outputs/train/reachy2_test \
@@ -277,10 +293,9 @@ python -m lerobot.scripts.train \
## Step 4: Evaluate ## Step 4: Evaluate
```bash ```bash
python -m lerobot.record \ lerobot-eval \
--robot.type=reachy2 \ --robot.type=reachy2 \
--robot.ip_address=192.168.0.200 \ --robot.ip_address=192.168.0.200 \
--display_data=false \
--dataset.repo_id=pollen_robotics/eval_record_test \ --dataset.repo_id=pollen_robotics/eval_record_test \
--dataset.single_task="Evaluate reachy2 policy" \ --dataset.single_task="Evaluate reachy2 policy" \
--dataset.num_episodes=10 \ --dataset.num_episodes=10 \
+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: 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** - **`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 - **Locomotion policies** from NVIDIA gr00t and Amazon FAR Holosoma
- **Simulation mode** for testing policies without the physical robot in mujoco - **Simulation mode** for testing policies without the physical robot in mujoco
@@ -110,7 +110,7 @@ ssh unitree@<YOUR_ROBOT_IP>
# Password: 123 # 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) ## 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 ## Additional Resources
+18 -21
View File
@@ -111,34 +111,29 @@ class GrootLocomotionController:
def run_step(self): def run_step(self):
# Get current observation # Get current observation
robot_state = self.robot.get_observation() obs = self.robot.get_observation()
if robot_state is None: if not obs:
return return
# Get command from remote controller # Get command from remote controller
if robot_state.wireless_remote is not None: if obs["remote.buttons"][0]: # R1 - raise waist
self.robot.remote_controller.set(robot_state.wireless_remote)
if self.robot.remote_controller.button[0]: # R1 - raise waist
self.groot_height_cmd += 0.001 self.groot_height_cmd += 0.001
self.groot_height_cmd = np.clip(self.groot_height_cmd, 0.50, 1.00) 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 -= 0.001
self.groot_height_cmd = np.clip(self.groot_height_cmd, 0.50, 1.00) 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[0] = obs["remote.ly"] # Forward/backward
self.cmd[1] = self.robot.remote_controller.lx * -1 # Left/right self.cmd[1] = obs["remote.lx"] * -1 # Left/right
self.cmd[2] = self.robot.remote_controller.rx * -1 # Rotation rate self.cmd[2] = obs["remote.rx"] * -1 # Rotation rate
# Get joint positions and velocities # Get joint positions and velocities from flat dict
for i in range(29): for motor in G1_29_JointIndex:
self.groot_qj_all[i] = robot_state.motor_state[i].q name = motor.name
self.groot_dqj_all[i] = robot_state.motor_state[i].dq 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 # Adapt observation for g1_23dof
for idx in MISSING_JOINTS: for idx in MISSING_JOINTS:
@@ -150,8 +145,8 @@ class GrootLocomotionController:
dqj_obs = self.groot_dqj_all.copy() dqj_obs = self.groot_dqj_all.copy()
# Express IMU data in gravity frame of reference # Express IMU data in gravity frame of reference
quat = robot_state.imu_state.quaternion quat = [obs["imu.quat.w"], obs["imu.quat.x"], obs["imu.quat.y"], obs["imu.quat.z"]]
ang_vel = np.array(robot_state.imu_state.gyroscope, dtype=np.float32) 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) gravity_orientation = self.robot.get_gravity_orientation(quat)
# Scale joint positions and velocities before policy inference # Scale joint positions and velocities before policy inference
@@ -219,6 +214,8 @@ def run(repo_id: str = DEFAULT_GROOT_REPO_ID) -> None:
config = UnitreeG1Config() config = UnitreeG1Config()
robot = UnitreeG1(config) robot = UnitreeG1(config)
robot.connect()
# Initialize gr00T locomotion controller # Initialize gr00T locomotion controller
groot_controller = GrootLocomotionController( groot_controller = GrootLocomotionController(
policy_balance=policy_balance, 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") logger.info("Press Ctrl+C to stop")
# Run step # Run step
while True: while not robot._shutdown_event.is_set():
start_time = time.time() start_time = time.time()
groot_controller.run_step() groot_controller.run_step()
elapsed = time.time() - start_time elapsed = time.time() - start_time
+14 -14
View File
@@ -126,24 +126,23 @@ class HolosomaLocomotionController:
def run_step(self): def run_step(self):
# Get current observation # Get current observation
robot_state = self.robot.get_observation() obs = self.robot.get_observation()
if robot_state is None: if not obs:
return return
# Get command from remote controller # Get command from remote controller
if robot_state.wireless_remote is not None: ly = obs["remote.ly"] if abs(obs["remote.ly"]) > 0.1 else 0.0
self.robot.remote_controller.set(robot_state.wireless_remote) 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
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
self.cmd[:] = [ly, -lx, -rx] self.cmd[:] = [ly, -lx, -rx]
# Get joint positions and velocities # Get joint positions and velocities
for i in range(29): for motor in G1_29_JointIndex:
self.qj[i] = robot_state.motor_state[i].q name = motor.name
self.dqj[i] = robot_state.motor_state[i].dq idx = motor.value
self.qj[idx] = obs[f"{name}.q"]
self.dqj[idx] = obs[f"{name}.dq"]
# Adapt observation for g1_23dof # Adapt observation for g1_23dof
for idx in MISSING_JOINTS: for idx in MISSING_JOINTS:
@@ -151,8 +150,8 @@ class HolosomaLocomotionController:
self.dqj[idx] = 0.0 self.dqj[idx] = 0.0
# Express IMU data in gravity frame of reference # Express IMU data in gravity frame of reference
quat = robot_state.imu_state.quaternion quat = [obs["imu.quat.w"], obs["imu.quat.x"], obs["imu.quat.y"], obs["imu.quat.z"]]
ang_vel = np.array(robot_state.imu_state.gyroscope, dtype=np.float32) 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) gravity = self.robot.get_gravity_orientation(quat)
# Scale joint positions and velocities before policy inference # 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 # Initialize robot
config = UnitreeG1Config() config = UnitreeG1Config()
robot = UnitreeG1(config) robot = UnitreeG1(config)
robot.connect()
holosoma_controller = HolosomaLocomotionController(policy, robot, kp, kd) 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") logger.info("Press Ctrl+C to stop")
# Run step # Run step
while True: while not robot._shutdown_event.is_set():
start_time = time.time() start_time = time.time()
holosoma_controller.run_step() holosoma_controller.run_step()
elapsed = time.time() - start_time elapsed = time.time() - start_time
+1 -1
View File
@@ -111,7 +111,7 @@ unitree_g1 = [
"pyzmq>=26.2.1,<28.0.0", "pyzmq>=26.2.1,<28.0.0",
"onnxruntime>=1.16.0,<2.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]"] kinematics = ["lerobot[placo-dep]"]
intelrealsense = [ intelrealsense = [
"pyrealsense2>=2.55.1.6486,<2.57.0 ; sys_platform != 'darwin'", "pyrealsense2>=2.55.1.6486,<2.57.0 ; sys_platform != 'darwin'",
+2 -2
View File
@@ -40,7 +40,6 @@ from collections.abc import Callable
from dataclasses import asdict from dataclasses import asdict
from pprint import pformat from pprint import pformat
from queue import Queue from queue import Queue
from typing import Any
import draccus import draccus
import grpc import grpc
@@ -48,6 +47,7 @@ import torch
from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig # noqa: F401 from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig # noqa: F401
from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig # noqa: F401 from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig # noqa: F401
from lerobot.processor import RobotAction
from lerobot.robots import ( # noqa: F401 from lerobot.robots import ( # noqa: F401
Robot, Robot,
RobotConfig, RobotConfig,
@@ -351,7 +351,7 @@ class RobotClient:
action = {key: action_tensor[i].item() for i, key in enumerate(self.robot.action_features)} action = {key: action_tensor[i].item() for i, key in enumerate(self.robot.action_features)}
return action 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""" """Reading and performing actions in local queue"""
# Lock only for queue operations # Lock only for queue operations
@@ -35,18 +35,19 @@ class Reachy2CameraConfig(CameraConfig):
name="teleop", name="teleop",
image_type="left", image_type="left",
ip_address="192.168.0.200", # IP address of the robot ip_address="192.168.0.200", # IP address of the robot
fps=15, port=50065, # Port of the camera server
width=640, width=640,
height=480, height=480,
fps=30, # Not configurable for Reachy 2 cameras
color_mode=ColorMode.RGB, color_mode=ColorMode.RGB,
) # Left teleop camera, 640x480 @ 15FPS ) # Left teleop camera, 640x480 @ 30FPS
``` ```
Attributes: Attributes:
name: Name of the camera device. Can be "teleop" or "depth". 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". 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) 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. width: Requested frame width in pixels for the color stream.
height: Requested frame height 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. 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 color_mode: ColorMode = ColorMode.RGB
ip_address: str | None = "localhost" ip_address: str | None = "localhost"
port: int = 50065 port: int = 50065
# use_depth: bool = False
def __post_init__(self) -> None: def __post_init__(self) -> None:
if self.name not in ["teleop", "depth"]: if self.name not in ["teleop", "depth"]:
@@ -16,12 +16,13 @@
Provides the Reachy2Camera class for capturing frames from Reachy 2 cameras using Reachy 2's CameraManager. Provides the Reachy2Camera class for capturing frames from Reachy 2 cameras using Reachy 2's CameraManager.
""" """
from __future__ import annotations
import logging import logging
import os import os
import platform import platform
import time import time
from threading import Event, Lock, Thread from typing import TYPE_CHECKING, Any
from typing import Any
from numpy.typing import NDArray # type: ignore # TODO: add type stubs for numpy.typing 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" os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0"
import cv2 # type: ignore # TODO: add type stubs for OpenCV import cv2 # type: ignore # TODO: add type stubs for OpenCV
import numpy as np # type: ignore # TODO: add type stubs for numpy 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 from lerobot.utils.import_utils import _reachy2_sdk_available
CameraManager,
) 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 from lerobot.utils.errors import DeviceNotConnectedError
@@ -69,17 +79,10 @@ class Reachy2Camera(Camera):
self.config = config self.config = config
self.fps = config.fps
self.color_mode = config.color_mode self.color_mode = config.color_mode
self.cam_manager: CameraManager | None = None 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: def __str__(self) -> str:
return f"{self.__class__.__name__}({self.config.name}, {self.config.image_type})" 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: def connect(self, warmup: bool = True) -> None:
""" """
Connects to the Reachy2 CameraManager as specified in the configuration. 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) 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() self.cam_manager.initialize_cameras()
logger.info(f"{self} connected.") logger.info(f"{self} connected.")
@staticmethod @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. Detection not implemented for Reachy2 cameras.
Returns:
List[Dict[str, Any]]: A list of dictionaries,
where each dictionary contains 'name', 'stereo',
and the default profile properties (width, height, fps).
""" """
initialized_cameras = [] raise NotImplementedError("Camera detection is not implemented for Reachy2 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
def read(self, color_mode: ColorMode | None = None) -> NDArray[Any]: 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 (height, width, channels), using the specified or default
color mode and applying any configured rotation. color mode and applying any configured rotation.
""" """
start_time = time.perf_counter()
if not self.is_connected: if not self.is_connected:
raise DeviceNotConnectedError(f"{self} is not 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) 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.name == "teleop" and hasattr(self.cam_manager, "teleop"):
if self.config.image_type == "left": 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": 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"): elif self.config.name == "depth" and hasattr(self.cam_manager, "depth"):
if self.config.image_type == "depth": if self.config.image_type == "depth":
frame = self.cam_manager.depth.get_depth_frame()[0] frame = self.cam_manager.depth.get_depth_frame()[0]
elif self.config.image_type == "rgb": 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: if frame is None:
return np.empty((0, 0, 3), dtype=np.uint8) return np.empty((0, 0, 3), dtype=np.uint8)
@@ -187,63 +175,11 @@ class Reachy2Camera(Camera):
return frame 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]: 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 This method retrieves the most recent frame available in Reachy 2's low-level software.
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.
Args: Args:
timeout_ms (float): Maximum time in milliseconds to wait for a frame timeout_ms (float): Maximum time in milliseconds to wait for a frame
@@ -261,22 +197,10 @@ class Reachy2Camera(Camera):
if not self.is_connected: if not self.is_connected:
raise DeviceNotConnectedError(f"{self} is not connected.") raise DeviceNotConnectedError(f"{self} is not connected.")
if self.thread is None or not self.thread.is_alive(): frame = self.read()
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()
if frame is None: 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 return frame
@@ -287,12 +211,9 @@ class Reachy2Camera(Camera):
Raises: Raises:
DeviceNotConnectedError: If the camera is already disconnected. 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.") raise DeviceNotConnectedError(f"{self} not connected.")
if self.thread is not None:
self._stop_read_thread()
if self.cam_manager is not None: if self.cam_manager is not None:
self.cam_manager.disconnect() self.cam_manager.disconnect()
+5
View File
@@ -43,6 +43,11 @@ def make_cameras_from_configs(camera_configs: dict[str, CameraConfig]) -> dict[s
cameras[key] = Reachy2Camera(cfg) cameras[key] = Reachy2Camera(cfg)
elif cfg.type == "zmq":
from .zmq.camera_zmq import ZMQCamera
cameras[key] = ZMQCamera(cfg)
else: else:
try: try:
cameras[key] = cast(Camera, make_device_from_device_class(cfg)) cameras[key] = cast(Camera, make_device_from_device_class(cfg))
+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"]
+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.")
@@ -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.")
+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()
+2 -2
View File
@@ -18,12 +18,12 @@ from typing import Any
from lerobot.configs.types import PipelineFeatureType from lerobot.configs.types import PipelineFeatureType
from lerobot.datasets.utils import hw_to_dataset_features 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 from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE, OBS_STR
def create_initial_features( 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]]: ) -> dict[PipelineFeatureType, dict[str, Any]]:
""" """
Creates the initial features dict for the dataset from action and observation specs. Creates the initial features dict for the dataset from action and observation specs.
+4 -2
View File
@@ -29,6 +29,8 @@ from gymnasium import spaces
from libero.libero import benchmark, get_libero_path from libero.libero import benchmark, get_libero_path
from libero.libero.envs import OffScreenRenderEnv from libero.libero.envs import OffScreenRenderEnv
from lerobot.processor import RobotObservation
def _parse_camera_names(camera_name: str | Sequence[str]) -> list[str]: def _parse_camera_names(camera_name: str | Sequence[str]) -> list[str]:
"""Normalize camera_name into a non-empty list of strings.""" """Normalize camera_name into a non-empty list of strings."""
@@ -237,7 +239,7 @@ class LiberoEnv(gym.Env):
env.reset() env.reset()
return env 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 = {} images = {}
for camera_name in self.camera_name: for camera_name in self.camera_name:
image = raw_obs[camera_name] image = raw_obs[camera_name]
@@ -313,7 +315,7 @@ class LiberoEnv(gym.Env):
info = {"is_success": False} info = {"is_success": False}
return observation, info 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: if action.ndim != 1:
raise ValueError( raise ValueError(
f"Expected action to be 1-D (shape (action_dim,)), " f"Expected action to be 1-D (shape (action_dim,)), "
+7 -5
View File
@@ -25,6 +25,8 @@ import metaworld.policies as policies
import numpy as np import numpy as np
from gymnasium import spaces from gymnasium import spaces
from lerobot.processor import RobotObservation
# ---- Load configuration data from the external JSON file ---- # ---- Load configuration data from the external JSON file ----
CONFIG_PATH = Path(__file__).parent / "metaworld_config.json" CONFIG_PATH = Path(__file__).parent / "metaworld_config.json"
try: try:
@@ -161,7 +163,7 @@ class MetaworldEnv(gym.Env):
env._freeze_rand_vec = False # otherwise no randomization env._freeze_rand_vec = False # otherwise no randomization
return env 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 image = None
if self._env is not None: if self._env is not None:
image = self._env.render() image = self._env.render()
@@ -196,7 +198,7 @@ class MetaworldEnv(gym.Env):
self, self,
seed: int | None = None, seed: int | None = None,
**kwargs, **kwargs,
) -> tuple[dict[str, Any], dict[str, Any]]: ) -> tuple[RobotObservation, dict[str, Any]]:
""" """
Reset the environment to its initial state. Reset the environment to its initial state.
@@ -204,7 +206,7 @@ class MetaworldEnv(gym.Env):
seed (Optional[int]): Random seed for environment initialization. seed (Optional[int]): Random seed for environment initialization.
Returns: 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. info (Dict[str, Any]): Additional info about the reset state.
""" """
super().reset(seed=seed) super().reset(seed=seed)
@@ -216,7 +218,7 @@ class MetaworldEnv(gym.Env):
info = {"is_success": False} info = {"is_success": False}
return observation, info 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. 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,). action (np.ndarray): The action to execute, must be 1-D with shape (action_dim,).
Returns: 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. reward (float): The scalar reward for this step.
terminated (bool): Whether the episode terminated successfully. terminated (bool): Whether the episode terminated successfully.
truncated (bool): Whether the episode was truncated due to a time limit. truncated (bool): Whether the episode was truncated due to a time limit.
+2 -1
View File
@@ -29,6 +29,7 @@ from torch import Tensor
from lerobot.configs.types import FeatureType, PolicyFeature from lerobot.configs.types import FeatureType, PolicyFeature
from lerobot.envs.configs import EnvConfig 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.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE, OBS_STR
from lerobot.utils.utils import get_channel_first_image_shape 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.""" """Adds task feature to the observation dict with respect to the first environment attribute."""
if hasattr(env.envs[0], "task_description"): if hasattr(env.envs[0], "task_description"):
task_result = env.call("task_description") task_result = env.call("task_description")
+6 -2
View File
@@ -51,7 +51,11 @@ from lerobot.processor.converters import (
transition_to_batch, transition_to_batch,
transition_to_policy_action, 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]: 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 # 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"] = { postprocessor_overrides["groot_action_unpack_unnormalize_v1"] = {
"stats": kwargs.get("dataset_stats"), "stats": kwargs.get("dataset_stats"),
"normalize_min_max": True, "normalize_min_max": True,
@@ -20,6 +20,7 @@ from lerobot.configs.policies import PreTrainedConfig
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
from lerobot.optim.optimizers import AdamWConfig from lerobot.optim.optimizers import AdamWConfig
from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig
from lerobot.utils.constants import ACTION, OBS_STATE
@PreTrainedConfig.register_subclass("groot") @PreTrainedConfig.register_subclass("groot")
@@ -137,14 +138,14 @@ class GrootConfig(PreTrainedConfig):
"No features of type FeatureType.VISUAL found in input_features." "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( state_feature = PolicyFeature(
type=FeatureType.STATE, type=FeatureType.STATE,
shape=(self.max_state_dim,), shape=(self.max_state_dim,),
) )
self.input_features["observation.state"] = state_feature self.input_features[OBS_STATE] = state_feature
else: 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 state_dim = state_shape[0] if state_shape else 0
if state_dim > self.max_state_dim: if state_dim > self.max_state_dim:
raise ValueError( raise ValueError(
@@ -152,14 +153,14 @@ class GrootConfig(PreTrainedConfig):
f"Either reduce state dimension or increase max_state_dim in config." 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( action_feature = PolicyFeature(
type=FeatureType.ACTION, type=FeatureType.ACTION,
shape=(self.max_action_dim,), shape=(self.max_action_dim,),
) )
self.output_features["action"] = action_feature self.output_features[ACTION] = action_feature
else: 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 action_dim = action_shape[0] if action_shape else 0
if action_dim > self.max_action_dim: if action_dim > self.max_action_dim:
raise ValueError( raise ValueError(
+3 -3
View File
@@ -46,7 +46,7 @@ from lerobot.policies.groot.action_head.flow_matching_action_head import (
FlowmatchingActionHeadConfig, FlowmatchingActionHeadConfig,
) )
from lerobot.policies.groot.utils import ensure_eagle_cache_ready 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_VENDOR_EAGLE_PATH = str((Path(__file__).resolve().parent / "eagle2_hg_model").resolve())
DEFAULT_TOKENIZER_ASSETS_REPO = "lerobot/eagle2hg-processor-groot-n1p5" DEFAULT_TOKENIZER_ASSETS_REPO = "lerobot/eagle2hg-processor-groot-n1p5"
@@ -227,8 +227,8 @@ class GR00TN15(PreTrainedModel):
detected_error = False detected_error = False
error_msg = ERROR_MSG error_msg = ERROR_MSG
if "action" in inputs: if ACTION in inputs:
action = inputs["action"] action = inputs[ACTION]
# In inference, action may be omitted or None; validate only when it's a tensor. # In inference, action may be omitted or None; validate only when it's a tensor.
if action is None: if action is None:
pass # allow None during inference pass # allow None during inference
+2 -1
View File
@@ -41,6 +41,7 @@ from torch import Tensor
from lerobot.policies.groot.configuration_groot import GrootConfig from lerobot.policies.groot.configuration_groot import GrootConfig
from lerobot.policies.groot.groot_n1 import GR00TN15 from lerobot.policies.groot.groot_n1 import GR00TN15
from lerobot.policies.pretrained import PreTrainedPolicy from lerobot.policies.pretrained import PreTrainedPolicy
from lerobot.utils.constants import ACTION
class GrootPolicy(PreTrainedPolicy): class GrootPolicy(PreTrainedPolicy):
@@ -147,7 +148,7 @@ class GrootPolicy(PreTrainedPolicy):
actions = outputs.get("action_pred") 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] actions = actions[:, :, :original_action_dim]
return actions return actions
+16 -12
View File
@@ -51,7 +51,11 @@ from lerobot.processor.converters import (
) )
from lerobot.processor.core import EnvTransition, TransitionKey from lerobot.processor.core import EnvTransition, TransitionKey
from lerobot.utils.constants import ( from lerobot.utils.constants import (
ACTION,
HF_LEROBOT_HOME, HF_LEROBOT_HOME,
OBS_IMAGE,
OBS_IMAGES,
OBS_STATE,
POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_POSTPROCESSOR_DEFAULT_NAME,
POLICY_PREPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME,
) )
@@ -107,9 +111,9 @@ def make_groot_pre_post_processors(
# Define feature specs for optional normalization steps # Define feature specs for optional normalization steps
_features: dict[str, PolicyFeature] = { _features: dict[str, PolicyFeature] = {
# Observation features (only add those we may normalize) # 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 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) # 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) # Determine env action dimension from config (simple, object-like PolicyFeature)
try: try:
env_action_dim = int(config.output_features["action"].shape[0]) env_action_dim = int(config.output_features[ACTION].shape[0])
except Exception: except Exception:
env_action_dim = 0 env_action_dim = 0
@@ -268,9 +272,9 @@ class GrootPackInputsStep(ProcessorStep):
return torch.where(mask, mapped, torch.zeros_like(mapped)) return torch.where(mask, mapped, torch.zeros_like(mapped))
# 1) Video (B, T=1, V, H, W, C) uint8 # 1) Video (B, T=1, V, H, W, C) uint8
img_keys = sorted([k for k in obs if k.startswith("observation.images.")]) img_keys = sorted([k for k in obs if k.startswith(OBS_IMAGES)])
if not img_keys and "observation.image" in obs: if not img_keys and OBS_IMAGE in obs:
img_keys = ["observation.image"] img_keys = [OBS_IMAGE]
if img_keys: if img_keys:
cams = [_to_uint8_np_bhwc(obs[k]) for k in 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) video = np.stack(cams, axis=1) # (B, V, H, W, C)
@@ -294,14 +298,14 @@ class GrootPackInputsStep(ProcessorStep):
comp["language"] = lang comp["language"] = lang
# 3) State/state_mask -> (B, 1, max_state_dim) # 3) State/state_mask -> (B, 1, max_state_dim)
if "observation.state" in obs: if OBS_STATE in obs:
state = obs["observation.state"] # (B, D) state = obs[OBS_STATE] # (B, D)
if state.dim() != 2: if state.dim() != 2:
raise ValueError(f"state must be (B, D), got {tuple(state.shape)}") raise ValueError(f"state must be (B, D), got {tuple(state.shape)}")
bsz, d = state.shape bsz, d = state.shape
# Normalize BEFORE padding # Normalize BEFORE padding
if self.normalize_min_max: 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) state = state.unsqueeze(1) # (B, 1, D)
if d > self.max_state_dim: if d > self.max_state_dim:
state = state[:, :, : self.max_state_dim] state = state[:, :, : self.max_state_dim]
@@ -320,11 +324,11 @@ class GrootPackInputsStep(ProcessorStep):
# Normalize BEFORE temporal expansion/padding # Normalize BEFORE temporal expansion/padding
if self.normalize_min_max: if self.normalize_min_max:
if action.dim() == 2: if action.dim() == 2:
action = _min_max_norm(action, "action") action = _min_max_norm(action, ACTION)
elif action.dim() == 3: elif action.dim() == 3:
b, t, d = action.shape b, t, d = action.shape
flat = action.reshape(b * t, d) flat = action.reshape(b * t, d)
flat = _min_max_norm(flat, "action") flat = _min_max_norm(flat, ACTION)
action = flat.view(b, t, d) action = flat.view(b, t, d)
if action.dim() == 2: if action.dim() == 2:
action = action.unsqueeze(1).repeat(1, self.action_horizon, 1) 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 # 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 # inverse: x = (y+1)/2 * denom + min, and when denom==0 -> x = min
if self.normalize_min_max and self.stats is not None: 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] d = action.shape[-1]
min_v = torch.as_tensor( min_v = torch.as_tensor(
stats_k.get("min", torch.zeros(d)), dtype=action.dtype, device=action.device stats_k.get("min", torch.zeros(d)), dtype=action.dtype, device=action.device
@@ -21,7 +21,7 @@ from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
from lerobot.optim.optimizers import AdamWConfig from lerobot.optim.optimizers import AdamWConfig
from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig
from lerobot.policies.rtc.configuration_rtc import RTCConfig 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 DEFAULT_IMAGE_SIZE = 224
@@ -124,19 +124,19 @@ class PI0Config(PreTrainedConfig):
) )
self.input_features[key] = empty_camera 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( state_feature = PolicyFeature(
type=FeatureType.STATE, type=FeatureType.STATE,
shape=(self.max_state_dim,), # Padded to max_state_dim 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( action_feature = PolicyFeature(
type=FeatureType.ACTION, type=FeatureType.ACTION,
shape=(self.max_action_dim,), # Padded to max_action_dim 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: def get_optimizer_preset(self) -> AdamWConfig:
return AdamWConfig( return AdamWConfig(
@@ -21,6 +21,7 @@ from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
from lerobot.optim.optimizers import AdamWConfig from lerobot.optim.optimizers import AdamWConfig
from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig
from lerobot.policies.rtc.configuration_rtc import RTCConfig from lerobot.policies.rtc.configuration_rtc import RTCConfig
from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE
DEFAULT_IMAGE_SIZE = 224 DEFAULT_IMAGE_SIZE = 224
@@ -117,26 +118,26 @@ class PI05Config(PreTrainedConfig):
def validate_features(self) -> None: def validate_features(self) -> None:
"""Validate and set up input/output features.""" """Validate and set up input/output features."""
for i in range(self.empty_cameras): for i in range(self.empty_cameras):
key = f"observation.images.empty_camera_{i}" key = OBS_IMAGES + f".empty_camera_{i}"
empty_camera = PolicyFeature( empty_camera = PolicyFeature(
type=FeatureType.VISUAL, type=FeatureType.VISUAL,
shape=(3, *self.image_resolution), # Use configured image resolution shape=(3, *self.image_resolution), # Use configured image resolution
) )
self.input_features[key] = empty_camera 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( state_feature = PolicyFeature(
type=FeatureType.STATE, type=FeatureType.STATE,
shape=(self.max_state_dim,), # Padded to max_state_dim 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( action_feature = PolicyFeature(
type=FeatureType.ACTION, type=FeatureType.ACTION,
shape=(self.max_action_dim,), # Padded to max_action_dim 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: def get_optimizer_preset(self) -> AdamWConfig:
return AdamWConfig( return AdamWConfig(
@@ -21,6 +21,7 @@ from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
from lerobot.optim.optimizers import AdamWConfig from lerobot.optim.optimizers import AdamWConfig
from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig
from lerobot.policies.rtc.configuration_rtc import RTCConfig from lerobot.policies.rtc.configuration_rtc import RTCConfig
from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE
DEFAULT_IMAGE_SIZE = 224 DEFAULT_IMAGE_SIZE = 224
@@ -110,26 +111,26 @@ class PI0FastConfig(PreTrainedConfig):
def validate_features(self) -> None: def validate_features(self) -> None:
"""Validate and set up input/output features.""" """Validate and set up input/output features."""
for i in range(self.empty_cameras): for i in range(self.empty_cameras):
key = f"observation.images.empty_camera_{i}" key = OBS_IMAGES + f".empty_camera_{i}"
empty_camera = PolicyFeature( empty_camera = PolicyFeature(
type=FeatureType.VISUAL, type=FeatureType.VISUAL,
shape=(3, *self.image_resolution), # Use configured image resolution shape=(3, *self.image_resolution), # Use configured image resolution
) )
self.input_features[key] = empty_camera 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( state_feature = PolicyFeature(
type=FeatureType.STATE, type=FeatureType.STATE,
shape=(self.max_state_dim,), # Padded to max_state_dim 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( action_feature = PolicyFeature(
type=FeatureType.ACTION, type=FeatureType.ACTION,
shape=(self.max_action_dim,), # Padded to max_action_dim 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: def get_optimizer_preset(self) -> AdamWConfig:
return AdamWConfig( return AdamWConfig(
@@ -26,6 +26,7 @@ from lerobot.configs.policies import PreTrainedConfig
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
from lerobot.optim.optimizers import AdamWConfig from lerobot.optim.optimizers import AdamWConfig
from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig
from lerobot.utils.constants import OBS_IMAGES, OBS_STATE
@PreTrainedConfig.register_subclass("sarm") @PreTrainedConfig.register_subclass("sarm")
@@ -86,8 +87,8 @@ class SARMConfig(PreTrainedConfig):
pretrained_model_path: str | None = None pretrained_model_path: str | None = None
device: str | None = None device: str | None = None
image_key: str = "observation.images.top" # Key for image used from the dataset image_key: str = OBS_IMAGES + ".top" # Key for image used from the dataset
state_key: str = "observation.state" state_key: str = OBS_STATE
# Populated by the processor (video_features, state_features, text_features) # Populated by the processor (video_features, state_features, text_features)
input_features: dict = field(default_factory=lambda: {}) input_features: dict = field(default_factory=lambda: {})
+2 -1
View File
@@ -40,6 +40,7 @@ from lerobot.policies.sarm.sarm_utils import (
normalize_stage_tau, normalize_stage_tau,
pad_state_to_max_dim, pad_state_to_max_dim,
) )
from lerobot.utils.constants import OBS_STR
class StageTransformer(nn.Module): class StageTransformer(nn.Module):
@@ -721,7 +722,7 @@ class SARMRewardModel(PreTrainedPolicy):
Returns: Returns:
Tuple of (total_loss, output_dict with loss components) Tuple of (total_loss, output_dict with loss components)
""" """
observation = batch.get("observation", batch) observation = batch.get(OBS_STR, batch)
# Extract features # Extract features
video_features = observation["video_features"].to(self.device) video_features = observation["video_features"].to(self.device)
+1 -2
View File
@@ -16,7 +16,6 @@
import logging import logging
from collections import deque from collections import deque
from typing import Any
import numpy as np import numpy as np
import torch import torch
@@ -140,7 +139,7 @@ def prepare_observation_for_inference(
def build_inference_frame( def build_inference_frame(
observation: dict[str, Any], observation: RobotObservation,
device: torch.device, device: torch.device,
ds_features: dict[str, dict], ds_features: dict[str, dict],
task: str | None = None, task: str | None = None,
@@ -18,6 +18,7 @@ from lerobot.configs.policies import PreTrainedConfig
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
from lerobot.optim.optimizers import AdamWConfig from lerobot.optim.optimizers import AdamWConfig
from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig
from lerobot.utils.constants import ACTION, OBS_STATE
@PreTrainedConfig.register_subclass("wall_x") @PreTrainedConfig.register_subclass("wall_x")
@@ -105,14 +106,14 @@ class WallXConfig(PreTrainedConfig):
"No features of type FeatureType.VISUAL found in input_features." "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( state_feature = PolicyFeature(
type=FeatureType.STATE, type=FeatureType.STATE,
shape=(self.max_state_dim,), # Padded to max_state_dim 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: 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 state_dim = state_shape[0] if state_shape else 0
if state_dim > self.max_state_dim: if state_dim > self.max_state_dim:
raise ValueError( raise ValueError(
@@ -120,14 +121,14 @@ class WallXConfig(PreTrainedConfig):
f"Either reduce state dimension or increase max_state_dim in config." 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( action_feature = PolicyFeature(
type=FeatureType.ACTION, type=FeatureType.ACTION,
shape=(self.max_action_dim,), # Padded to max_action_dim shape=(self.max_action_dim,), # Padded to max_action_dim
) )
self.output_features["action"] = action_feature self.output_features[ACTION] = action_feature
else: 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 action_dim = action_shape[0] if action_shape else 0
if action_dim > self.max_action_dim: if action_dim > self.max_action_dim:
raise ValueError( raise ValueError(
@@ -1861,7 +1861,7 @@ class WallXPolicy(PreTrainedPolicy):
dim=-1, dim=-1,
) )
else: else:
action_dim = self.config.output_features["action"].shape[0] action_dim = self.config.output_features[ACTION].shape[0]
dof_mask = torch.cat( dof_mask = torch.cat(
[ [
torch.ones( torch.ones(
@@ -1977,7 +1977,7 @@ class WallXPolicy(PreTrainedPolicy):
elif self.config.prediction_mode == "fast": elif self.config.prediction_mode == "fast":
output = self.model( output = self.model(
**batch, **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, pred_horizon=self.config.chunk_size,
mode="predict", mode="predict",
predict_mode="fast", predict_mode="fast",
@@ -1989,7 +1989,7 @@ class WallXPolicy(PreTrainedPolicy):
actions = output["predict_action"] actions = output["predict_action"]
# Unpad actions to actual action dimension # 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] actions = actions[:, :, :action_dim]
return actions return actions
+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.processor.core import EnvTransition, TransitionKey
from lerobot.utils.constants import ( from lerobot.utils.constants import (
OBS_IMAGES, OBS_IMAGES,
OBS_PREFIX,
OBS_STATE, OBS_STATE,
POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_POSTPROCESSOR_DEFAULT_NAME,
POLICY_PREPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME,
@@ -137,8 +138,9 @@ class LiberoProcessorStep(ObservationProcessorStep):
processed_obs[key] = img processed_obs[key] = img
# Process robot_state into a flat state vector # Process robot_state into a flat state vector
if "observation.robot_state" in processed_obs: robot_state_str = OBS_PREFIX + "robot_state"
robot_state = processed_obs.pop("observation.robot_state") if robot_state_str in processed_obs:
robot_state = processed_obs.pop(robot_state_str)
# Extract components # Extract components
eef_pos = robot_state["eef"]["pos"] # (B, 3,) eef_pos = robot_state["eef"]["pos"] # (B, 3,)
@@ -174,8 +176,8 @@ class LiberoProcessorStep(ObservationProcessorStep):
state_feats = {} state_feats = {}
# add our new flattened state # add our new flattened state
state_feats["observation.state"] = PolicyFeature( state_feats[OBS_STATE] = PolicyFeature(
key="observation.state", key=OBS_STATE,
shape=(20,), shape=(20,),
dtype="float32", dtype="float32",
) )
@@ -247,7 +249,7 @@ class XVLAImageScaleProcessorStep(ProcessorStep):
keys_to_scale = self.image_keys keys_to_scale = self.image_keys
if keys_to_scale is None: if keys_to_scale is None:
# Auto-detect image keys # 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 # Scale each image
for key in keys_to_scale: for key in keys_to_scale:
@@ -303,7 +305,7 @@ class XVLAImageToFloatProcessorStep(ProcessorStep):
keys_to_convert = self.image_keys keys_to_convert = self.image_keys
if keys_to_convert is None: if keys_to_convert is None:
# Auto-detect image keys # 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 # Convert each image
for key in keys_to_convert: for key in keys_to_convert:
@@ -376,7 +378,7 @@ class XVLAImageNetNormalizeProcessorStep(ProcessorStep):
keys_to_normalize = self.image_keys keys_to_normalize = self.image_keys
if keys_to_normalize is None: if keys_to_normalize is None:
# Auto-detect image keys # 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 # Normalize each image
for key in keys_to_normalize: for key in keys_to_normalize:
-3
View File
@@ -49,7 +49,6 @@ from .hil_processor import (
RewardClassifierProcessorStep, RewardClassifierProcessorStep,
TimeLimitProcessorStep, TimeLimitProcessorStep,
) )
from .joint_observations_processor import JointVelocityProcessorStep, MotorCurrentProcessorStep
from .normalize_processor import NormalizerProcessorStep, UnnormalizerProcessorStep, hotswap_stats from .normalize_processor import NormalizerProcessorStep, UnnormalizerProcessorStep, hotswap_stats
from .observation_processor import VanillaObservationProcessorStep from .observation_processor import VanillaObservationProcessorStep
from .pipeline import ( from .pipeline import (
@@ -94,14 +93,12 @@ __all__ = [
"ImageCropResizeProcessorStep", "ImageCropResizeProcessorStep",
"InfoProcessorStep", "InfoProcessorStep",
"InterventionActionProcessorStep", "InterventionActionProcessorStep",
"JointVelocityProcessorStep",
"make_default_processors", "make_default_processors",
"make_default_teleop_action_processor", "make_default_teleop_action_processor",
"make_default_robot_action_processor", "make_default_robot_action_processor",
"make_default_robot_observation_processor", "make_default_robot_observation_processor",
"MapDeltaActionToRobotActionStep", "MapDeltaActionToRobotActionStep",
"MapTensorToDeltaActionDictStep", "MapTensorToDeltaActionDictStep",
"MotorCurrentProcessorStep",
"NormalizerProcessorStep", "NormalizerProcessorStep",
"Numpy2TorchActionProcessorStep", "Numpy2TorchActionProcessorStep",
"ObservationProcessorStep", "ObservationProcessorStep",
+3 -3
View File
@@ -23,7 +23,7 @@ from typing import Any
import numpy as np import numpy as np
import torch 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 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( def create_transition(
observation: dict[str, Any] | None = None, observation: RobotObservation | None = None,
action: PolicyAction | RobotAction | None = None, action: PolicyAction | RobotAction | None = None,
reward: float = 0.0, reward: float = 0.0,
done: bool = False, done: bool = False,
@@ -384,7 +384,7 @@ def transition_to_batch(transition: EnvTransition) -> dict[str, Any]:
REWARD: transition.get(TransitionKey.REWARD, 0.0), REWARD: transition.get(TransitionKey.REWARD, 0.0),
DONE: transition.get(TransitionKey.DONE, False), DONE: transition.get(TransitionKey.DONE, False),
TRUNCATED: transition.get(TransitionKey.TRUNCATED, False), TRUNCATED: transition.get(TransitionKey.TRUNCATED, False),
"info": transition.get(TransitionKey.INFO, {}), INFO: transition.get(TransitionKey.INFO, {}),
} }
# Add complementary data. # Add complementary data.
+1 -1
View File
@@ -45,7 +45,7 @@ RobotObservation: TypeAlias = dict[str, Any]
EnvTransition = TypedDict( EnvTransition = TypedDict(
"EnvTransition", "EnvTransition",
{ {
TransitionKey.OBSERVATION.value: dict[str, Any] | None, TransitionKey.OBSERVATION.value: RobotObservation | None,
TransitionKey.ACTION.value: PolicyAction | RobotAction | EnvAction | None, TransitionKey.ACTION.value: PolicyAction | RobotAction | EnvAction | None,
TransitionKey.REWARD.value: float | torch.Tensor | None, TransitionKey.REWARD.value: float | torch.Tensor | None,
TransitionKey.DONE.value: bool | torch.Tensor | None, TransitionKey.DONE.value: bool | torch.Tensor | None,
+6 -5
View File
@@ -18,7 +18,7 @@ from dataclasses import dataclass
import torch import torch
from lerobot.configs.types import PipelineFeatureType, PolicyFeature 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 from .pipeline import ObservationProcessorStep, ProcessorStepRegistry
@@ -60,8 +60,9 @@ class LiberoProcessorStep(ObservationProcessorStep):
processed_obs[key] = img processed_obs[key] = img
# Process robot_state into a flat state vector # Process robot_state into a flat state vector
if "observation.robot_state" in processed_obs: observation_robot_state_str = OBS_PREFIX + "robot_state"
robot_state = processed_obs.pop("observation.robot_state") if observation_robot_state_str in processed_obs:
robot_state = processed_obs.pop(observation_robot_state_str)
# Extract components # Extract components
eef_pos = robot_state["eef"]["pos"] # (B, 3,) eef_pos = robot_state["eef"]["pos"] # (B, 3,)
@@ -98,8 +99,8 @@ class LiberoProcessorStep(ObservationProcessorStep):
state_feats = {} state_feats = {}
# add our new flattened state # add our new flattened state
state_feats["observation.state"] = PolicyFeature( state_feats[OBS_STATE] = PolicyFeature(
key="observation.state", key=OBS_STATE,
shape=(8,), # [eef_pos(3), axis_angle(3), gripper(2)] shape=(8,), # [eef_pos(3), axis_angle(3), gripper(2)]
dtype="float32", dtype="float32",
description=("Concatenated end-effector position (3), axis-angle (3), and gripper qpos (2)."), description=("Concatenated end-effector position (3), axis-angle (3), and gripper qpos (2)."),
+2 -2
View File
@@ -30,7 +30,7 @@ from lerobot.utils.constants import ACTION
from .converters import from_tensor_to_numpy, to_tensor from .converters import from_tensor_to_numpy, to_tensor
from .core import EnvTransition, PolicyAction, TransitionKey from .core import EnvTransition, PolicyAction, TransitionKey
from .pipeline import PolicyProcessorPipeline, ProcessorStep, ProcessorStepRegistry from .pipeline import PolicyProcessorPipeline, ProcessorStep, ProcessorStepRegistry, RobotObservation
@dataclass @dataclass
@@ -239,7 +239,7 @@ class _NormalizationMixin:
config["normalize_observation_keys"] = sorted(self.normalize_observation_keys) config["normalize_observation_keys"] = sorted(self.normalize_observation_keys)
return config 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. Applies (un)normalization to all relevant features in an observation dictionary.
+3 -3
View File
@@ -49,7 +49,7 @@ from lerobot.configs.types import PipelineFeatureType, PolicyFeature
from lerobot.utils.hub import HubMixin from lerobot.utils.hub import HubMixin
from .converters import batch_to_transition, create_transition, transition_to_batch 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. # Generic type variables for pipeline input and output.
TInput = TypeVar("TInput") TInput = TypeVar("TInput")
@@ -1337,7 +1337,7 @@ class DataProcessorPipeline(HubMixin, Generic[TInput, TOutput]):
return features return features
# Convenience methods for processing individual parts of a transition. # 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. """Processes only the observation part of a transition through the pipeline.
Args: Args:
@@ -1440,7 +1440,7 @@ class ObservationProcessorStep(ProcessorStep, ABC):
"""An abstract `ProcessorStep` that specifically targets the observation in a transition.""" """An abstract `ProcessorStep` that specifically targets the observation in a transition."""
@abstractmethod @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. """Processes an observation dictionary. Subclasses must implement this method.
Args: Args:
+2 -2
View File
@@ -38,7 +38,7 @@ from lerobot.utils.constants import (
) )
from lerobot.utils.import_utils import _transformers_available 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 from .pipeline import ActionProcessorStep, ObservationProcessorStep, ProcessorStepRegistry
# Conditional import for type checking and lazy loading # Conditional import for type checking and lazy loading
@@ -139,7 +139,7 @@ class TokenizerProcessorStep(ObservationProcessorStep):
return None 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. Tokenizes the task description and adds it to the observation dictionary.
+6 -5
View File
@@ -38,13 +38,12 @@ from lerobot.processor import (
GripperPenaltyProcessorStep, GripperPenaltyProcessorStep,
ImageCropResizeProcessorStep, ImageCropResizeProcessorStep,
InterventionActionProcessorStep, InterventionActionProcessorStep,
JointVelocityProcessorStep,
MapDeltaActionToRobotActionStep, MapDeltaActionToRobotActionStep,
MapTensorToDeltaActionDictStep, MapTensorToDeltaActionDictStep,
MotorCurrentProcessorStep,
Numpy2TorchActionProcessorStep, Numpy2TorchActionProcessorStep,
RewardClassifierProcessorStep, RewardClassifierProcessorStep,
RobotActionToPolicyActionProcessorStep, RobotActionToPolicyActionProcessorStep,
RobotObservation,
TimeLimitProcessorStep, TimeLimitProcessorStep,
Torch2NumpyActionProcessorStep, Torch2NumpyActionProcessorStep,
TransitionKey, 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.robot_utils import precise_sleep
from lerobot.utils.utils import log_say from lerobot.utils.utils import log_say
from .joint_observations_processor import JointVelocityProcessorStep, MotorCurrentProcessorStep
logging.basicConfig(level=logging.INFO) logging.basicConfig(level=logging.INFO)
@@ -163,7 +164,7 @@ class RobotEnv(gym.Env):
self._setup_spaces() 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.""" """Get current robot observation including joint positions and camera images."""
obs_dict = self.robot.get_observation() obs_dict = self.robot.get_observation()
raw_joint_joint_position = {f"{name}.pos": obs_dict[f"{name}.pos"] for name in self._joint_names} 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( def reset(
self, *, seed: int | None = None, options: dict[str, Any] | None = None 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. """Reset environment to initial state.
Args: 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} self._raw_joint_positions = {f"{key}.pos": obs[f"{key}.pos"] for key in self._joint_names}
return obs, {TeleopEvents.IS_INTERVENTION: False} 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.""" """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())} joint_targets_dict = {f"{key}.pos": action[i] for i, key in enumerate(self.robot.bus.motors.keys())}
@@ -16,8 +16,8 @@
import logging import logging
from functools import cached_property from functools import cached_property
from typing import Any
from lerobot.processor import RobotAction, RobotObservation
from lerobot.robots.so_follower import SOFollower, SOFollowerRobotConfig from lerobot.robots.so_follower import SOFollower, SOFollowerRobotConfig
from ..robot import Robot from ..robot import Robot
@@ -116,7 +116,7 @@ class BiSOFollower(Robot):
self.left_arm.setup_motors() self.left_arm.setup_motors()
self.right_arm.setup_motors() self.right_arm.setup_motors()
def get_observation(self) -> dict[str, Any]: def get_observation(self) -> RobotObservation:
obs_dict = {} obs_dict = {}
# Add "left_" prefix # Add "left_" prefix
@@ -129,7 +129,7 @@ class BiSOFollower(Robot):
return obs_dict return obs_dict
def send_action(self, action: dict[str, Any]) -> dict[str, Any]: def send_action(self, action: RobotAction) -> RobotAction:
# Remove "left_" prefix # Remove "left_" prefix
left_action = { left_action = {
key.removeprefix("left_"): value for key, value in action.items() if key.startswith("left_") key.removeprefix("left_"): value for key, value in action.items() if key.startswith("left_")
@@ -18,12 +18,12 @@
import base64 import base64
import logging import logging
from functools import cached_property from functools import cached_property
from typing import Any
import cv2 import cv2
import numpy as np import numpy as np
import requests import requests
from lerobot.processor import RobotAction, RobotObservation
from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
from ..robot import Robot from ..robot import Robot
@@ -197,11 +197,11 @@ class EarthRoverMiniPlus(Robot):
ACTION_ANGULAR_VEL: float, ACTION_ANGULAR_VEL: float,
} }
def get_observation(self) -> dict[str, Any]: def get_observation(self) -> RobotObservation:
"""Get current robot observation from SDK. """Get current robot observation from SDK.
Returns: Returns:
dict: Observation containing: RobotObservation: Observation containing:
- front: Front camera image (480, 640, 3) in RGB format - front: Front camera image (480, 640, 3) in RGB format
- rear: Rear 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) - linear.vel: Current speed (0-1, SDK reports only positive speeds)
@@ -255,7 +255,7 @@ class EarthRoverMiniPlus(Robot):
return observation 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. """Send action to robot via SDK.
Args: Args:
@@ -264,8 +264,7 @@ class EarthRoverMiniPlus(Robot):
- angular.vel: Target angular velocity (-1 to 1) - angular.vel: Target angular velocity (-1 to 1)
Returns: Returns:
dict: The action that was sent (matches action_features keys) RobotAction: The action that was sent (matches action_features keys)
Raises: Raises:
DeviceNotConnectedError: If robot is not connected DeviceNotConnectedError: If robot is not connected
+3 -3
View File
@@ -17,7 +17,6 @@
import logging import logging
import time import time
from functools import cached_property from functools import cached_property
from typing import Any
from lerobot.cameras.utils import make_cameras_from_configs from lerobot.cameras.utils import make_cameras_from_configs
from lerobot.motors import Motor, MotorNormMode from lerobot.motors import Motor, MotorNormMode
@@ -25,6 +24,7 @@ from lerobot.motors.calibration_gui import RangeFinderGUI
from lerobot.motors.feetech import ( from lerobot.motors.feetech import (
FeetechMotorsBus, FeetechMotorsBus,
) )
from lerobot.processor import RobotAction, RobotObservation
from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
from ..robot import Robot from ..robot import Robot
@@ -128,7 +128,7 @@ class HopeJrArm(Robot):
self.bus.setup_motor(motor) self.bus.setup_motor(motor)
print(f"'{motor}' motor id set to {self.bus.motors[motor].id}") 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: if not self.is_connected:
raise DeviceNotConnectedError(f"{self} is not connected.") raise DeviceNotConnectedError(f"{self} is not connected.")
@@ -149,7 +149,7 @@ class HopeJrArm(Robot):
return obs_dict 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: if not self.is_connected:
raise DeviceNotConnectedError(f"{self} is not connected.") raise DeviceNotConnectedError(f"{self} is not connected.")
+3 -3
View File
@@ -17,7 +17,6 @@
import logging import logging
import time import time
from functools import cached_property from functools import cached_property
from typing import Any
from lerobot.cameras.utils import make_cameras_from_configs from lerobot.cameras.utils import make_cameras_from_configs
from lerobot.motors import Motor, MotorNormMode from lerobot.motors import Motor, MotorNormMode
@@ -25,6 +24,7 @@ from lerobot.motors.calibration_gui import RangeFinderGUI
from lerobot.motors.feetech import ( from lerobot.motors.feetech import (
FeetechMotorsBus, FeetechMotorsBus,
) )
from lerobot.processor import RobotAction, RobotObservation
from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
from ..robot import Robot from ..robot import Robot
@@ -159,7 +159,7 @@ class HopeJrHand(Robot):
self.bus.setup_motor(motor) self.bus.setup_motor(motor)
print(f"'{motor}' motor id set to {self.bus.motors[motor].id}") 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: if not self.is_connected:
raise DeviceNotConnectedError(f"{self} is not connected.") raise DeviceNotConnectedError(f"{self} is not connected.")
@@ -181,7 +181,7 @@ class HopeJrHand(Robot):
return obs_dict 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: if not self.is_connected:
raise DeviceNotConnectedError(f"{self} is not connected.") raise DeviceNotConnectedError(f"{self} is not connected.")
@@ -17,7 +17,6 @@
import logging import logging
import time import time
from functools import cached_property from functools import cached_property
from typing import Any
from lerobot.cameras.utils import make_cameras_from_configs from lerobot.cameras.utils import make_cameras_from_configs
from lerobot.motors import Motor, MotorCalibration, MotorNormMode from lerobot.motors import Motor, MotorCalibration, MotorNormMode
@@ -25,6 +24,7 @@ from lerobot.motors.dynamixel import (
DynamixelMotorsBus, DynamixelMotorsBus,
OperatingMode, OperatingMode,
) )
from lerobot.processor import RobotAction, RobotObservation
from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
from ..robot import Robot from ..robot import Robot
@@ -182,7 +182,7 @@ class KochFollower(Robot):
self.bus.setup_motor(motor) self.bus.setup_motor(motor)
print(f"'{motor}' motor id set to {self.bus.motors[motor].id}") 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: if not self.is_connected:
raise DeviceNotConnectedError(f"{self} is not connected.") raise DeviceNotConnectedError(f"{self} is not connected.")
@@ -202,7 +202,7 @@ class KochFollower(Robot):
return obs_dict 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. """Command arm to move to a target joint configuration.
The relative action magnitude may be clipped depending on the configuration parameter 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. Thus, this function always returns the action actually sent.
Args: Args:
action (dict[str, float]): The goal positions for the motors. action (RobotAction): The goal positions for the motors.
Returns: 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: if not self.is_connected:
raise DeviceNotConnectedError(f"{self} is not connected.") raise DeviceNotConnectedError(f"{self} is not connected.")
+4 -3
View File
@@ -28,6 +28,7 @@ from lerobot.motors.feetech import (
FeetechMotorsBus, FeetechMotorsBus,
OperatingMode, OperatingMode,
) )
from lerobot.processor import RobotAction, RobotObservation
from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
from ..robot import Robot from ..robot import Robot
@@ -338,7 +339,7 @@ class LeKiwi(Robot):
"theta.vel": theta, "theta.vel": theta,
} # m/s and deg/s } # m/s and deg/s
def get_observation(self) -> dict[str, Any]: def get_observation(self) -> RobotObservation:
if not self.is_connected: if not self.is_connected:
raise DeviceNotConnectedError(f"{self} is not connected.") raise DeviceNotConnectedError(f"{self} is not connected.")
@@ -369,7 +370,7 @@ class LeKiwi(Robot):
return obs_dict 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. """Command lekiwi to move to a target joint configuration.
The relative action magnitude may be clipped depending on the configuration parameter 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. RobotDeviceNotConnectedError: if robot is not connected.
Returns: 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: if not self.is_connected:
raise DeviceNotConnectedError(f"{self} is not connected.") raise DeviceNotConnectedError(f"{self} is not connected.")
+9 -10
View File
@@ -18,11 +18,11 @@ import base64
import json import json
import logging import logging
from functools import cached_property from functools import cached_property
from typing import Any
import cv2 import cv2
import numpy as np import numpy as np
from lerobot.processor import RobotAction, RobotObservation
from lerobot.utils.constants import ACTION, OBS_STATE from lerobot.utils.constants import ACTION, OBS_STATE
from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
@@ -172,7 +172,7 @@ class LeKiwiClient(Robot):
return last_msg 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.""" """Parses the JSON observation string."""
try: try:
return json.loads(obs_string) return json.loads(obs_string)
@@ -196,15 +196,15 @@ class LeKiwiClient(Robot):
return None return None
def _remote_state_from_obs( def _remote_state_from_obs(
self, observation: dict[str, Any] self, observation: RobotObservation
) -> tuple[dict[str, np.ndarray], dict[str, Any]]: ) -> tuple[dict[str, np.ndarray], RobotObservation]:
"""Extracts frames, and state from the parsed observation.""" """Extracts frames, and state from the parsed observation."""
flat_state = {key: observation.get(key, 0.0) for key in self._state_order} 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) 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 # Decode images
current_frames: dict[str, np.ndarray] = {} current_frames: dict[str, np.ndarray] = {}
@@ -217,7 +217,7 @@ class LeKiwiClient(Robot):
return current_frames, obs_dict 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. Polls the video socket for the latest observation data.
@@ -252,7 +252,7 @@ class LeKiwiClient(Robot):
return new_frames, new_state 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, Capture observations from the remote robot: current follower arm positions,
present wheel speeds (converted to body-frame velocities: x, y, theta), present wheel speeds (converted to body-frame velocities: x, y, theta),
@@ -307,12 +307,11 @@ class LeKiwiClient(Robot):
def configure(self): def configure(self):
pass 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 """Command lekiwi to move to a target joint configuration. Translates to motor space + sends over ZMQ
Args: Args:
action (np.ndarray): array containing the goal positions for the motors. action (RobotAction): array containing the goal positions for the motors.
Raises: Raises:
RobotDeviceNotConnectedError: if robot is not connected. RobotDeviceNotConnectedError: if robot is not connected.
@@ -17,7 +17,6 @@
import logging import logging
import time import time
from functools import cached_property from functools import cached_property
from typing import Any
from lerobot.cameras.utils import make_cameras_from_configs from lerobot.cameras.utils import make_cameras_from_configs
from lerobot.motors import Motor, MotorCalibration, MotorNormMode from lerobot.motors import Motor, MotorCalibration, MotorNormMode
@@ -26,6 +25,7 @@ from lerobot.motors.dynamixel import (
DynamixelMotorsBus, DynamixelMotorsBus,
OperatingMode, OperatingMode,
) )
from lerobot.processor import RobotAction, RobotObservation
from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
from ..robot import Robot from ..robot import Robot
@@ -165,7 +165,7 @@ class OmxFollower(Robot):
self.bus.setup_motor(motor) self.bus.setup_motor(motor)
print(f"'{motor}' motor id set to {self.bus.motors[motor].id}") 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: if not self.is_connected:
raise DeviceNotConnectedError(f"{self} is not connected.") raise DeviceNotConnectedError(f"{self} is not connected.")
@@ -185,7 +185,7 @@ class OmxFollower(Robot):
return obs_dict 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. """Command arm to move to a target joint configuration.
The relative action magnitude may be clipped depending on the configuration parameter 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. Thus, this function always returns the action actually sent.
Args: Args:
action (dict[str, float]): The goal positions for the motors. action (RobotAction): The goal positions for the motors.
Returns: 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: if not self.is_connected:
raise DeviceNotConnectedError(f"{self} is not connected.") raise DeviceNotConnectedError(f"{self} is not connected.")
@@ -30,6 +30,8 @@ class Reachy2RobotConfig(RobotConfig):
# IP address of the Reachy 2 robot # IP address of the Reachy 2 robot
ip_address: str | None = "localhost" 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. # If True, turn_off_smoothly() will be sent to the robot before disconnecting.
disable_torque_on_disconnect: bool = False disable_torque_on_disconnect: bool = False
@@ -51,11 +53,16 @@ class Reachy2RobotConfig(RobotConfig):
# Robot cameras # Robot cameras
# Set to True if you want to use the corresponding cameras in the observations. # Set to True if you want to use the corresponding cameras in the observations.
# By default, only the teleop cameras are used. # By default, no camera is used.
with_left_teleop_camera: bool = True with_left_teleop_camera: bool = False
with_right_teleop_camera: bool = True with_right_teleop_camera: bool = False
with_torso_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) cameras: dict[str, CameraConfig] = field(default_factory=dict)
def __post_init__(self) -> None: def __post_init__(self) -> None:
@@ -65,9 +72,10 @@ class Reachy2RobotConfig(RobotConfig):
name="teleop", name="teleop",
image_type="left", image_type="left",
ip_address=self.ip_address, ip_address=self.ip_address,
fps=15, port=self.port,
width=640, width=self.camera_width,
height=480, height=self.camera_height,
fps=30, # Not configurable for Reachy 2 cameras
color_mode=ColorMode.RGB, color_mode=ColorMode.RGB,
) )
if self.with_right_teleop_camera: if self.with_right_teleop_camera:
@@ -75,9 +83,10 @@ class Reachy2RobotConfig(RobotConfig):
name="teleop", name="teleop",
image_type="right", image_type="right",
ip_address=self.ip_address, ip_address=self.ip_address,
fps=15, port=self.port,
width=640, width=self.camera_width,
height=480, height=self.camera_height,
fps=30, # Not configurable for Reachy 2 cameras
color_mode=ColorMode.RGB, color_mode=ColorMode.RGB,
) )
if self.with_torso_camera: if self.with_torso_camera:
@@ -85,9 +94,10 @@ class Reachy2RobotConfig(RobotConfig):
name="depth", name="depth",
image_type="rgb", image_type="rgb",
ip_address=self.ip_address, ip_address=self.ip_address,
fps=15, port=self.port,
width=640, width=self.camera_width,
height=480, height=self.camera_height,
fps=30, # Not configurable for Reachy 2 cameras
color_mode=ColorMode.RGB, color_mode=ColorMode.RGB,
) )
+12 -7
View File
@@ -13,19 +13,24 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
from __future__ import annotations
import time import time
from typing import Any from typing import TYPE_CHECKING, Any
import numpy as np
from reachy2_sdk import ReachySDK
from lerobot.cameras.utils import make_cameras_from_configs from lerobot.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 ..robot import Robot
from ..utils import ensure_safe_goal_position from ..utils import ensure_safe_goal_position
from .configuration_reachy2 import Reachy2RobotConfig 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} # {lerobot_keys: reachy2_sdk_keys}
REACHY2_NECK_JOINTS = { REACHY2_NECK_JOINTS = {
"neck_yaw.pos": "head.neck.yaw", "neck_yaw.pos": "head.neck.yaw",
@@ -165,8 +170,8 @@ class Reachy2Robot(Robot):
else: else:
return {} return {}
def get_observation(self) -> dict[str, np.ndarray]: def get_observation(self) -> RobotObservation:
obs_dict: dict[str, Any] = {} obs_dict: RobotObservation = {}
# Read Reachy 2 state # Read Reachy 2 state
before_read_t = time.perf_counter() before_read_t = time.perf_counter()
@@ -179,7 +184,7 @@ class Reachy2Robot(Robot):
return obs_dict 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 self.reachy is not None:
if not self.is_connected: if not self.is_connected:
raise ConnectionError() raise ConnectionError()
+6 -6
View File
@@ -15,11 +15,11 @@
import abc import abc
import builtins import builtins
from pathlib import Path from pathlib import Path
from typing import Any
import draccus import draccus
from lerobot.motors import MotorCalibration from lerobot.motors import MotorCalibration
from lerobot.processor import RobotAction, RobotObservation
from lerobot.utils.constants import HF_LEROBOT_CALIBRATION, ROBOTS from lerobot.utils.constants import HF_LEROBOT_CALIBRATION, ROBOTS
from .config import RobotConfig from .config import RobotConfig
@@ -153,28 +153,28 @@ class Robot(abc.ABC):
pass pass
@abc.abstractmethod @abc.abstractmethod
def get_observation(self) -> dict[str, Any]: def get_observation(self) -> RobotObservation:
""" """
Retrieve the current observation from the robot. Retrieve the current observation from the robot.
Returns: 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`. should match :pymeth:`observation_features`.
""" """
pass pass
@abc.abstractmethod @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. Send an action command to the robot.
Args: 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`. :pymeth:`action_features`.
Returns: 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. safety limits on velocity.
""" """
pass pass
@@ -28,6 +28,7 @@ from lerobot.processor import (
ProcessorStepRegistry, ProcessorStepRegistry,
RobotAction, RobotAction,
RobotActionProcessorStep, RobotActionProcessorStep,
RobotObservation,
TransitionKey, TransitionKey,
) )
from lerobot.utils.rotation import Rotation from lerobot.utils.rotation import Rotation
@@ -438,7 +439,7 @@ class ForwardKinematicsJointsToEEObservation(ObservationProcessorStep):
kinematics: RobotKinematics kinematics: RobotKinematics
motor_names: list[str] 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) return compute_forward_kinematics_joints_to_ee(observation, self.kinematics, self.motor_names)
def transform_features( def transform_features(
@@ -17,7 +17,7 @@
import logging import logging
import time import time
from functools import cached_property 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.cameras.utils import make_cameras_from_configs
from lerobot.motors import Motor, MotorCalibration, MotorNormMode from lerobot.motors import Motor, MotorCalibration, MotorNormMode
@@ -25,6 +25,7 @@ from lerobot.motors.feetech import (
FeetechMotorsBus, FeetechMotorsBus,
OperatingMode, OperatingMode,
) )
from lerobot.processor import RobotAction, RobotObservation
from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
from ..robot import Robot from ..robot import Robot
@@ -175,7 +176,7 @@ class SOFollower(Robot):
self.bus.setup_motor(motor) self.bus.setup_motor(motor)
print(f"'{motor}' motor id set to {self.bus.motors[motor].id}") 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: if not self.is_connected:
raise DeviceNotConnectedError(f"{self} is not connected.") raise DeviceNotConnectedError(f"{self} is not connected.")
@@ -195,7 +196,7 @@ class SOFollower(Robot):
return obs_dict 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. """Command arm to move to a target joint configuration.
The relative action magnitude may be clipped depending on the configuration parameter 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. RobotDeviceNotConnectedError: if robot is not connected.
Returns: Returns:
the action sent to the motors, potentially clipped. RobotAction: the action sent to the motors, potentially clipped.
""" """
if not self.is_connected: if not self.is_connected:
raise DeviceNotConnectedError(f"{self} is not connected.") raise DeviceNotConnectedError(f"{self} is not connected.")
@@ -16,6 +16,8 @@
from dataclasses import dataclass, field from dataclasses import dataclass, field
from lerobot.cameras import CameraConfig
from ..config import RobotConfig from ..config import RobotConfig
_GAINS: dict[str, dict[str, list[float]]] = { _GAINS: dict[str, dict[str, list[float]]] = {
@@ -60,3 +62,6 @@ class UnitreeG1Config(RobotConfig):
# Socket config for ZMQ bridge # Socket config for ZMQ bridge
robot_ip: str = "192.168.123.164" # default G1 IP robot_ip: str = "192.168.123.164" # default G1 IP
# Cameras (ZMQ-based remote cameras)
cameras: dict[str, CameraConfig] = field(default_factory=dict)
+189 -88
View File
@@ -23,14 +23,10 @@ from functools import cached_property
from typing import Any from typing import Any
import numpy as np 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.envs.factory import make_env
from lerobot.processor import RobotAction, RobotObservation
from lerobot.robots.unitree_g1.g1_utils import G1_29_JointIndex from lerobot.robots.unitree_g1.g1_utils import G1_29_JointIndex
from ..robot import Robot from ..robot import Robot
@@ -43,8 +39,6 @@ logger = logging.getLogger(__name__)
kTopicLowCommand_Debug = "rt/lowcmd" kTopicLowCommand_Debug = "rt/lowcmd"
kTopicLowState = "rt/lowstate" kTopicLowState = "rt/lowstate"
G1_29_Num_Motors = 29
@dataclass @dataclass
class MotorState: class MotorState:
@@ -66,28 +60,12 @@ class IMUState:
# g1 observation class # g1 observation class
@dataclass @dataclass
class G1_29_LowState: # noqa: N801 class G1_29_LowState: # noqa: N801
motor_state: list[MotorState] = field( motor_state: list[MotorState] = field(default_factory=lambda: [MotorState() for _ in G1_29_JointIndex])
default_factory=lambda: [MotorState() for _ in range(G1_29_Num_Motors)]
)
imu_state: IMUState = field(default_factory=IMUState) imu_state: IMUState = field(default_factory=IMUState)
wireless_remote: Any = None # Raw wireless remote data wireless_remote: Any = None # Raw wireless remote data
mode_machine: int = 0 # Robot mode 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): class UnitreeG1(Robot):
config_class = UnitreeG1Config config_class = UnitreeG1Config
name = "unitree_g1" name = "unitree_g1"
@@ -117,9 +95,12 @@ class UnitreeG1(Robot):
logger.info("Initialize UnitreeG1...") logger.info("Initialize UnitreeG1...")
self.config = config self.config = config
self.control_dt = config.control_dt 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: if config.is_simulation:
from unitree_sdk2py.core.channel import ( from unitree_sdk2py.core.channel import (
ChannelFactoryInitialize, ChannelFactoryInitialize,
@@ -133,62 +114,33 @@ class UnitreeG1(Robot):
ChannelSubscriber, ChannelSubscriber,
) )
# connect robot # Store for use in connect()
self.ChannelFactoryInitialize = ChannelFactoryInitialize self._ChannelFactoryInitialize = ChannelFactoryInitialize
self.connect() self._ChannelPublisher = ChannelPublisher
self._ChannelSubscriber = ChannelSubscriber
# initialize direct motor control interface # Initialize state variables
self.lowcmd_publisher = ChannelPublisher(kTopicLowCommand_Debug, hg_LowCmd) self.sim_env = None
self.lowcmd_publisher.Init() self._env_wrapper = None
self.lowstate_subscriber = ChannelSubscriber(kTopicLowState, hg_LowState) self._lowstate = None
self.lowstate_subscriber.Init()
self.lowstate_buffer = DataBuffer()
# initialize subscribe thread to read robot state
self._shutdown_event = threading.Event() self._shutdown_event = threading.Event()
self.subscribe_thread = threading.Thread(target=self._subscribe_motor_state) self.subscribe_thread = None
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.remote_controller = self.RemoteController() self.remote_controller = self.RemoteController()
def _subscribe_motor_state(self): # polls robot state @ 250Hz def _subscribe_motor_state(self): # polls robot state @ 250Hz
while not self._shutdown_event.is_set(): while not self._shutdown_event.is_set():
start_time = time.time() 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() msg = self.lowstate_subscriber.Read()
if msg is not None: if msg is not None:
lowstate = G1_29_LowState() lowstate = G1_29_LowState()
# Capture motor states # Capture motor states using jointindex
for id in range(G1_29_Num_Motors): for id in G1_29_JointIndex:
lowstate.motor_state[id].q = msg.motor_state[id].q lowstate.motor_state[id].q = msg.motor_state[id].q
lowstate.motor_state[id].dq = msg.motor_state[id].dq lowstate.motor_state[id].dq = msg.motor_state[id].dq
lowstate.motor_state[id].tau_est = msg.motor_state[id].tau_est lowstate.motor_state[id].tau_est = msg.motor_state[id].tau_est
@@ -207,7 +159,7 @@ class UnitreeG1(Robot):
# Capture mode_machine # Capture mode_machine
lowstate.mode_machine = msg.mode_machine lowstate.mode_machine = msg.mode_machine
self.lowstate_buffer.set_data(lowstate) self._lowstate = lowstate
current_time = time.time() current_time = time.time()
all_t_elapsed = current_time - start_time all_t_elapsed = current_time - start_time
@@ -216,7 +168,7 @@ class UnitreeG1(Robot):
@cached_property @cached_property
def action_features(self) -> dict[str, type]: 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 def calibrate(self) -> None: # robot is already calibrated
pass pass
@@ -225,20 +177,153 @@ class UnitreeG1(Robot):
pass pass
def connect(self, calibrate: bool = True) -> None: # connect to DDS 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: if self.config.is_simulation:
self.ChannelFactoryInitialize(0, "lo") self._ChannelFactoryInitialize(0, "lo")
self.mujoco_env = make_env("lerobot/unitree-g1-mujoco", trust_remote_code=True) 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: 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): def disconnect(self):
# Signal thread to stop and unblock any waits
self._shutdown_event.set() 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]: # Wait for subscribe thread to finish
return self.lowstate_buffer.get_data() 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 @property
def is_calibrated(self) -> bool: def is_calibrated(self) -> bool:
@@ -246,11 +331,15 @@ class UnitreeG1(Robot):
@property @property
def is_connected(self) -> bool: def is_connected(self) -> bool:
return self.lowstate_buffer.get_data() is not None return self._lowstate is not None
@property @property
def _motors_ft(self) -> dict[str, type]: 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 @property
def _cameras_ft(self) -> dict[str, tuple]: def _cameras_ft(self) -> dict[str, tuple]:
@@ -262,7 +351,7 @@ class UnitreeG1(Robot):
def observation_features(self) -> dict[str, type | tuple]: def observation_features(self) -> dict[str, type | tuple]:
return {**self._motors_ft, **self._cameras_ft} 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: for motor in G1_29_JointIndex:
key = f"{motor.name}.q" key = f"{motor.name}.q"
if key in action: if key in action:
@@ -293,22 +382,34 @@ class UnitreeG1(Robot):
self, self,
control_dt: float | None = None, control_dt: float | None = None,
default_positions: list[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: if control_dt is None:
control_dt = self.config.control_dt control_dt = self.config.control_dt
if default_positions is None: if default_positions is None:
default_positions = np.array(self.config.default_positions, dtype=np.float32) 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 total_time = 3.0
num_steps = int(total_time / control_dt) num_steps = int(total_time / control_dt)
# get current state # get current state
robot_state = self.get_observation() obs = self.get_observation()
# record current positions # record current positions
init_dof_pos = np.zeros(29, dtype=np.float32) init_dof_pos = np.zeros(29, dtype=np.float32)
for i in range(29): for motor in G1_29_JointIndex:
init_dof_pos[i] = robot_state.motor_state[i].q init_dof_pos[motor.value] = obs[f"{motor.name}.q"]
# Interpolate to default position # Interpolate to default position
for step in range(num_steps): for step in range(num_steps):
+2 -2
View File
@@ -177,9 +177,9 @@ def rollout(
action = policy.select_action(observation) action = policy.select_action(observation)
action = postprocessor(action) action = postprocessor(action)
action_transition = {"action": action} action_transition = {ACTION: action}
action_transition = env_postprocessor(action_transition) action_transition = env_postprocessor(action_transition)
action = action_transition["action"] action = action_transition[ACTION]
# Convert to CPU / numpy. # Convert to CPU / numpy.
action_numpy: np.ndarray = action.to("cpu").numpy() action_numpy: np.ndarray = action.to("cpu").numpy()
+10
View File
@@ -73,7 +73,9 @@ from lerobot.cameras import ( # noqa: F401
CameraConfig, # noqa: F401 CameraConfig, # noqa: F401
) )
from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig # 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.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 import parser
from lerobot.configs.policies import PreTrainedConfig from lerobot.configs.policies import PreTrainedConfig
from lerobot.datasets.image_writer import safe_stop_image_writer from lerobot.datasets.image_writer import safe_stop_image_writer
@@ -102,7 +104,9 @@ from lerobot.robots import ( # noqa: F401
koch_follower, koch_follower,
make_robot_from_config, make_robot_from_config,
omx_follower, omx_follower,
reachy2,
so_follower, so_follower,
unitree_g1,
) )
from lerobot.teleoperators import ( # noqa: F401 from lerobot.teleoperators import ( # noqa: F401
Teleoperator, Teleoperator,
@@ -112,6 +116,7 @@ from lerobot.teleoperators import ( # noqa: F401
koch_leader, koch_leader,
make_teleoperator_from_config, make_teleoperator_from_config,
omx_leader, omx_leader,
reachy2_teleoperator,
so_leader, so_leader,
) )
from lerobot.teleoperators.keyboard.teleop_keyboard import KeyboardTeleop 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"] (recorded_episodes < cfg.dataset.num_episodes - 1) or events["rerecord_episode"]
): ):
log_say("Reset the environment", cfg.play_sounds) log_say("Reset the environment", cfg.play_sounds)
# reset g1 robot
if robot.name == "unitree_g1":
robot.reset()
record_loop( record_loop(
robot=robot, robot=robot,
events=events, events=events,
+2
View File
@@ -59,7 +59,9 @@ from lerobot.robots import ( # noqa: F401
koch_follower, koch_follower,
make_robot_from_config, make_robot_from_config,
omx_follower, omx_follower,
reachy2,
so_follower, so_follower,
unitree_g1,
) )
from lerobot.utils.constants import ACTION from lerobot.utils.constants import ACTION
from lerobot.utils.import_utils import register_third_party_plugins from lerobot.utils.import_utils import register_third_party_plugins
+3 -1
View File
@@ -76,6 +76,7 @@ from lerobot.robots import ( # noqa: F401
koch_follower, koch_follower,
make_robot_from_config, make_robot_from_config,
omx_follower, omx_follower,
reachy2,
so_follower, so_follower,
) )
from lerobot.teleoperators import ( # noqa: F401 from lerobot.teleoperators import ( # noqa: F401
@@ -88,6 +89,7 @@ from lerobot.teleoperators import ( # noqa: F401
koch_leader, koch_leader,
make_teleoperator_from_config, make_teleoperator_from_config,
omx_leader, omx_leader,
reachy2_teleoperator,
so_leader, so_leader,
) )
from lerobot.utils.import_utils import register_third_party_plugins from lerobot.utils.import_utils import register_third_party_plugins
@@ -163,7 +165,7 @@ def teleop_loop(
# Process action for robot through pipeline # Process action for robot through pipeline
robot_action_to_send = robot_action_processor((teleop_action, obs)) 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) _ = robot.send_action(robot_action_to_send)
if display_data: if display_data:
+2 -2
View File
@@ -253,6 +253,8 @@ def train(cfg: TrainPipelineConfig, accelerator: Accelerator | None = None):
cfg: A `TrainPipelineConfig` object containing all training configurations. cfg: A `TrainPipelineConfig` object containing all training configurations.
accelerator: Optional Accelerator instance. If None, one will be created automatically. accelerator: Optional Accelerator instance. If None, one will be created automatically.
""" """
cfg.validate()
# Create Accelerator if not provided # Create Accelerator if not provided
# It will automatically detect if running in distributed mode or single-process mode # 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 # 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 # When using accelerate, only the main process should log to avoid duplicate outputs
is_main_process = accelerator.is_main_process is_main_process = accelerator.is_main_process
cfg.validate()
# Only log on main process # Only log on main process
if is_main_process: if is_main_process:
logging.info(pformat(cfg.to_dict())) logging.info(pformat(cfg.to_dict()))
+110 -78
View File
@@ -14,11 +14,14 @@
"""Train FAST tokenizer for action encoding. """Train FAST tokenizer for action encoding.
This script: This script:
1. Loads action chunks from LeRobotDataset (with sampling) 1. Loads action chunks from LeRobotDataset (with episode sampling)
2. Applies delta transforms and per-timestamp normalization 2. Optionally applies delta transforms (relative vs absolute actions)
3. Trains FAST tokenizer on specified action dimensions 3. Extracts specified action dimensions for encoding
4. Saves tokenizer to assets directory 4. Applies normalization (MEAN_STD, MIN_MAX, QUANTILES, or other modes)
5. Reports compression statistics 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: Example:
@@ -42,16 +45,63 @@ lerobot-train-tokenizer \
""" """
import json import json
from dataclasses import dataclass
from pathlib import Path from pathlib import Path
from typing import TYPE_CHECKING
import numpy as np import numpy as np
import torch import torch
import tyro
from huggingface_hub import HfApi 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.configs.types import NormalizationMode
from lerobot.datasets.lerobot_dataset import LeRobotDataset 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: 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: else:
# if no state key, use zeros (no delta transform) # if no state key, use zeros (no delta transform)
state = np.zeros_like( 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 = ( action = 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"])
)
states.append(state) states.append(state)
actions.append(action) actions.append(action)
@@ -327,88 +375,57 @@ def compute_compression_stats(tokenizer, action_chunks: np.ndarray):
return stats return stats
def main( @parser.wrap()
repo_id: str, def train_tokenizer(cfg: TokenizerTrainingConfig):
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,
):
""" """
Train FAST tokenizer for action encoding. Train FAST tokenizer for action encoding.
Args: Args:
repo_id: LeRobot dataset repository ID cfg: TokenizerTrainingConfig dataclass with all configuration parameters
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
""" """
# load dataset # load dataset
print(f"Loading dataset: {repo_id}") print(f"Loading dataset: {cfg.repo_id}")
dataset = LeRobotDataset(repo_id=repo_id, root=root) dataset = LeRobotDataset(repo_id=cfg.repo_id, root=cfg.root)
print(f"Dataset loaded: {dataset.num_episodes} episodes, {dataset.num_frames} frames") print(f"Dataset loaded: {dataset.num_episodes} episodes, {dataset.num_frames} frames")
# parse normalization mode # parse normalization mode
try: try:
norm_mode = NormalizationMode(normalization_mode) norm_mode = NormalizationMode(cfg.normalization_mode)
except ValueError as err: except ValueError as err:
raise ValueError( 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])}" f"Must be one of: {', '.join([m.value for m in NormalizationMode])}"
) from err ) from err
print(f"Normalization mode: {norm_mode.value}") print(f"Normalization mode: {norm_mode.value}")
# parse encoded dimensions # parse encoded dimensions
encoded_dim_ranges = [] 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(":")) start, end = map(int, range_str.strip().split(":"))
encoded_dim_ranges.append((start, end)) encoded_dim_ranges.append((start, end))
total_encoded_dims = sum(end - start for start, end in encoded_dim_ranges) 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 # parse delta dimensions
delta_dim_list = None delta_dim_list = None
if delta_dims is not None and delta_dims.strip(): if cfg.delta_dims is not None and cfg.delta_dims.strip():
delta_dim_list = [int(d.strip()) for d in delta_dims.split(",")] delta_dim_list = [int(d.strip()) for d in cfg.delta_dims.split(",")]
print(f"Delta dimensions: {delta_dim_list}") print(f"Delta dimensions: {delta_dim_list}")
else: else:
print("No delta dimensions specified") print("No delta dimensions specified")
print(f"Use delta transform: {use_delta_transform}") print(f"Use delta transform: {cfg.use_delta_transform}")
if use_delta_transform and (delta_dim_list is None or len(delta_dim_list) == 0): 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("Warning: use_delta_transform=True but no delta_dims specified. No delta will be applied.")
print(f"Action horizon: {action_horizon}") print(f"Action horizon: {cfg.action_horizon}")
print(f"State key: {state_key}") print(f"State key: {cfg.state_key}")
# determine episodes to process # determine episodes to process
num_episodes = dataset.num_episodes num_episodes = dataset.num_episodes
if max_episodes is not None: if cfg.max_episodes is not None:
num_episodes = min(max_episodes, num_episodes) num_episodes = min(cfg.max_episodes, num_episodes)
print(f"Processing {num_episodes} episodes...") print(f"Processing {num_episodes} episodes...")
@@ -419,7 +436,15 @@ def main(
print(f" Processing episode {ep_idx}/{num_episodes}...") print(f" Processing episode {ep_idx}/{num_episodes}...")
chunks = process_episode( 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: if chunks is not None:
all_chunks.append(chunks) all_chunks.append(chunks)
@@ -442,8 +467,8 @@ def main(
# get normalization stats from dataset # get normalization stats from dataset
norm_stats = dataset.meta.stats norm_stats = dataset.meta.stats
if norm_stats is not None and "action" in norm_stats: if norm_stats is not None and ACTION in norm_stats:
action_stats = norm_stats["action"] action_stats = norm_stats[ACTION]
# build encoded dimension indices # build encoded dimension indices
encoded_dim_indices = [] encoded_dim_indices = []
@@ -495,16 +520,17 @@ def main(
# train FAST tokenizer # train FAST tokenizer
tokenizer = train_fast_tokenizer( tokenizer = train_fast_tokenizer(
encoded_chunks, encoded_chunks,
vocab_size=vocab_size, vocab_size=cfg.vocab_size,
scale=scale, scale=cfg.scale,
) )
# compute compression statistics # compute compression statistics
compression_stats = compute_compression_stats(tokenizer, encoded_chunks) compression_stats = compute_compression_stats(tokenizer, encoded_chunks)
# save tokenizer # save tokenizer
output_dir = cfg.output_dir
if output_dir is None: 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 = Path(output_dir)
output_path.mkdir(parents=True, exist_ok=True) output_path.mkdir(parents=True, exist_ok=True)
@@ -512,18 +538,18 @@ def main(
# save metadata # save metadata
metadata = { metadata = {
"repo_id": repo_id, "repo_id": cfg.repo_id,
"vocab_size": vocab_size, "vocab_size": cfg.vocab_size,
"scale": scale, "scale": cfg.scale,
"encoded_dims": encoded_dims, "encoded_dims": cfg.encoded_dims,
"encoded_dim_ranges": encoded_dim_ranges, "encoded_dim_ranges": encoded_dim_ranges,
"total_encoded_dims": total_encoded_dims, "total_encoded_dims": total_encoded_dims,
"delta_dims": delta_dims, "delta_dims": cfg.delta_dims,
"delta_dim_list": delta_dim_list, "delta_dim_list": delta_dim_list,
"use_delta_transform": use_delta_transform, "use_delta_transform": cfg.use_delta_transform,
"state_key": state_key, "state_key": cfg.state_key,
"normalization_mode": norm_mode.value, "normalization_mode": norm_mode.value,
"action_horizon": action_horizon, "action_horizon": cfg.action_horizon,
"num_training_chunks": len(encoded_chunks), "num_training_chunks": len(encoded_chunks),
"compression_stats": compression_stats, "compression_stats": compression_stats,
} }
@@ -535,21 +561,22 @@ def main(
print(f"Metadata: {json.dumps(metadata, indent=2)}") print(f"Metadata: {json.dumps(metadata, indent=2)}")
# push to Hugging Face Hub if requested # push to Hugging Face Hub if requested
if push_to_hub: if cfg.push_to_hub:
# determine the hub repository ID # determine the hub repository ID
hub_repo_id = cfg.hub_repo_id
if hub_repo_id is None: if hub_repo_id is None:
hub_repo_id = output_path.name hub_repo_id = output_path.name
print(f"\nNo hub_repo_id provided, using: {hub_repo_id}") print(f"\nNo hub_repo_id provided, using: {hub_repo_id}")
print(f"\nPushing tokenizer to Hugging Face Hub: {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: try:
# use the tokenizer's push_to_hub method # use the tokenizer's push_to_hub method
tokenizer.push_to_hub( tokenizer.push_to_hub(
repo_id=hub_repo_id, repo_id=hub_repo_id,
private=hub_private, private=cfg.hub_private,
commit_message=f"Upload FAST tokenizer trained on {repo_id}", commit_message=f"Upload FAST tokenizer trained on {cfg.repo_id}",
) )
# also upload the metadata.json file separately # also upload the metadata.json file separately
@@ -568,5 +595,10 @@ def main(
print(" Make sure you're logged in with `huggingface-cli login`") 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__": if __name__ == "__main__":
tyro.cli(main) main()
@@ -20,6 +20,8 @@ from typing import Any
import numpy as np import numpy as np
from lerobot.processor import RobotAction
from ..teleoperator import Teleoperator from ..teleoperator import Teleoperator
from ..utils import TeleopEvents from ..utils import TeleopEvents
from .configuration_gamepad import GamepadTeleopConfig from .configuration_gamepad import GamepadTeleopConfig
@@ -83,7 +85,7 @@ class GamepadTeleop(Teleoperator):
self.gamepad = Gamepad() self.gamepad = Gamepad()
self.gamepad.start() self.gamepad.start()
def get_action(self) -> dict[str, Any]: def get_action(self) -> RobotAction:
# Update the controller to get fresh inputs # Update the controller to get fresh inputs
self.gamepad.update() self.gamepad.update()
@@ -21,6 +21,7 @@ import time
from queue import Queue from queue import Queue
from typing import Any from typing import Any
from lerobot.processor import RobotAction
from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
from ..teleoperator import Teleoperator from ..teleoperator import Teleoperator
@@ -124,7 +125,7 @@ class KeyboardTeleop(Teleoperator):
def configure(self): def configure(self):
pass pass
def get_action(self) -> dict[str, Any]: def get_action(self) -> RobotAction:
before_read_t = time.perf_counter() before_read_t = time.perf_counter()
if not self.is_connected: if not self.is_connected:
@@ -181,7 +182,7 @@ class KeyboardEndEffectorTeleop(KeyboardTeleop):
"names": {"delta_x": 0, "delta_y": 1, "delta_z": 2}, "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: if not self.is_connected:
raise DeviceNotConnectedError( raise DeviceNotConnectedError(
"KeyboardTeleop is not connected. You need to run `connect()` before `get_action()`." "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 # Only remove key if it's being released
self.current_pressed.pop(key_char, None) 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. Get the current action based on pressed keys.
Returns: Returns:
dict with 'linear.vel' and 'angular.vel' keys RobotAction with 'linear.vel' and 'angular.vel' keys
""" """
before_read_t = time.perf_counter() before_read_t = time.perf_counter()
@@ -13,11 +13,20 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
from __future__ import annotations
import logging import logging
import time 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 from reachy2_sdk import ReachySDK
else:
ReachySDK = None
from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
from ..teleoperator import Teleoperator from ..teleoperator import Teleoperator
from .config_reachy2_teleoperator import Reachy2TeleoperatorConfig from .config_reachy2_teleoperator import Reachy2TeleoperatorConfig
@@ -75,6 +84,7 @@ class Reachy2Teleoperator(Teleoperator):
def __init__(self, config: Reachy2TeleoperatorConfig): def __init__(self, config: Reachy2TeleoperatorConfig):
super().__init__(config) super().__init__(config)
self.config = config self.config = config
self.reachy: None | ReachySDK = None 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 return self.reachy.is_connected() if self.reachy is not None else False
def connect(self, calibrate: bool = True) -> None: def connect(self, calibrate: bool = True) -> None:
if self.is_connected:
raise DeviceAlreadyConnectedError(f"{self} already connected")
self.reachy = ReachySDK(self.config.ip_address) self.reachy = ReachySDK(self.config.ip_address)
if not self.is_connected: if not self.is_connected:
raise ConnectionError() raise DeviceNotConnectedError()
logger.info(f"{self} connected.") logger.info(f"{self} connected.")
@property @property
@@ -135,19 +149,20 @@ class Reachy2Teleoperator(Teleoperator):
def get_action(self) -> dict[str, float]: def get_action(self) -> dict[str, float]:
start = time.perf_counter() 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: if self.config.use_present_position:
joint_action = { joint_action = {k: self.reachy.joints[v].present_position for k, v in self.joints_dict.items()}
k: self.reachy.joints[v].present_position for k, v in self.joints_dict.items()
}
else: else:
joint_action = {k: self.reachy.joints[v].goal_position for k, v in self.joints_dict.items()} joint_action = {k: self.reachy.joints[v].goal_position for k, v in self.joints_dict.items()}
if not self.config.with_mobile_base: if not self.config.with_mobile_base:
dt_ms = (time.perf_counter() - start) * 1e3 dt_ms = (time.perf_counter() - start) * 1e3
logger.debug(f"{self} read action: {dt_ms:.1f}ms") logger.debug(f"{self} read action: {dt_ms:.1f}ms")
return joint_action return joint_action
if self.config.use_present_position: if self.config.use_present_position:
vel_action = {k: self.reachy.mobile_base.odometry[v] for k, v in REACHY2_VEL.items()} vel_action = {k: self.reachy.mobile_base.odometry[v] for k, v in REACHY2_VEL.items()}
else: else:
@@ -160,5 +175,5 @@ class Reachy2Teleoperator(Teleoperator):
raise NotImplementedError raise NotImplementedError
def disconnect(self) -> None: def disconnect(self) -> None:
if self.reachy and self.is_connected: if self.is_connected:
self.reachy.disconnect() self.reachy.disconnect()
+3 -2
View File
@@ -20,6 +20,7 @@ from typing import Any
import draccus import draccus
from lerobot.motors.motors_bus import MotorCalibration from lerobot.motors.motors_bus import MotorCalibration
from lerobot.processor import RobotAction
from lerobot.utils.constants import HF_LEROBOT_CALIBRATION, TELEOPERATORS from lerobot.utils.constants import HF_LEROBOT_CALIBRATION, TELEOPERATORS
from .config import TeleoperatorConfig from .config import TeleoperatorConfig
@@ -150,12 +151,12 @@ class Teleoperator(abc.ABC):
pass pass
@abc.abstractmethod @abc.abstractmethod
def get_action(self) -> dict[str, Any]: def get_action(self) -> RobotAction:
""" """
Retrieve the current action from the teleoperator. Retrieve the current action from the teleoperator.
Returns: 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`. structure should match :pymeth:`observation_features`.
""" """
pass pass
+2
View File
@@ -28,11 +28,13 @@ OBS_LANGUAGE_TOKENS = OBS_LANGUAGE + ".tokens"
OBS_LANGUAGE_ATTENTION_MASK = OBS_LANGUAGE + ".attention_mask" OBS_LANGUAGE_ATTENTION_MASK = OBS_LANGUAGE + ".attention_mask"
ACTION = "action" ACTION = "action"
ACTION_PREFIX = ACTION + "."
ACTION_TOKENS = ACTION + ".tokens" ACTION_TOKENS = ACTION + ".tokens"
ACTION_TOKEN_MASK = ACTION + ".token_mask" ACTION_TOKEN_MASK = ACTION + ".token_mask"
REWARD = "next.reward" REWARD = "next.reward"
TRUNCATED = "next.truncated" TRUNCATED = "next.truncated"
DONE = "next.done" DONE = "next.done"
INFO = "info"
ROBOTS = "robots" ROBOTS = "robots"
TELEOPERATORS = "teleoperators" TELEOPERATORS = "teleoperators"
+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") _transformers_available = is_package_available("transformers")
_peft_available = is_package_available("peft") _peft_available = is_package_available("peft")
_scipy_available = is_package_available("scipy") _scipy_available = is_package_available("scipy")
_reachy2_sdk_available = is_package_available("reachy2_sdk")
def make_device_from_device_class(config: ChoiceRegistry) -> Any: def make_device_from_device_class(config: ChoiceRegistry) -> Any:
+6 -5
View File
@@ -14,12 +14,13 @@
import numbers import numbers
import os import os
from typing import Any
import numpy as np import numpy as np
import rerun as rr 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( def init_rerun(
@@ -50,8 +51,8 @@ def _is_scalar(x):
def log_rerun_data( def log_rerun_data(
observation: dict[str, Any] | None = None, observation: RobotObservation | None = None,
action: dict[str, Any] | None = None, action: RobotAction | None = None,
compress_images: bool = False, compress_images: bool = False,
) -> None: ) -> None:
""" """
@@ -96,7 +97,7 @@ def log_rerun_data(
for k, v in action.items(): for k, v in action.items():
if v is None: if v is None:
continue 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): if _is_scalar(v):
rr.log(key, rr.Scalars(float(v))) rr.log(key, rr.Scalars(float(v)))
+2 -12
View File
@@ -20,6 +20,8 @@ from unittest.mock import MagicMock, patch
import numpy as np import numpy as np
import pytest import pytest
pytest.importorskip("reachy2_sdk")
from lerobot.cameras.reachy2_camera import Reachy2Camera, Reachy2CameraConfig from lerobot.cameras.reachy2_camera import Reachy2Camera, Reachy2CameraConfig
from lerobot.utils.errors import DeviceNotConnectedError from lerobot.utils.errors import DeviceNotConnectedError
@@ -127,24 +129,12 @@ def test_async_read(camera):
try: try:
img = camera.async_read() img = camera.async_read()
assert camera.thread is not None
assert camera.thread.is_alive()
assert isinstance(img, np.ndarray) assert isinstance(img, np.ndarray)
finally: finally:
if camera.is_connected: if camera.is_connected:
camera.disconnect() 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): def test_read_before_connect(camera):
with pytest.raises(DeviceNotConnectedError): with pytest.raises(DeviceNotConnectedError):
_ = camera.read() _ = camera.read()
-1
View File
@@ -28,7 +28,6 @@ pytest_plugins = [
"tests.fixtures.files", "tests.fixtures.files",
"tests.fixtures.hub", "tests.fixtures.hub",
"tests.fixtures.optimizers", "tests.fixtures.optimizers",
"tests.plugins.reachy2_sdk",
] ]
+3 -3
View File
@@ -17,10 +17,10 @@
import random import random
from dataclasses import dataclass, field from dataclasses import dataclass, field
from functools import cached_property from functools import cached_property
from typing import Any
from lerobot.cameras import CameraConfig, make_cameras_from_configs from lerobot.cameras import CameraConfig, make_cameras_from_configs
from lerobot.motors.motors_bus import Motor, MotorNormMode from lerobot.motors.motors_bus import Motor, MotorNormMode
from lerobot.processor import RobotAction, RobotObservation
from lerobot.robots import Robot, RobotConfig from lerobot.robots import Robot, RobotConfig
from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
from tests.mocks.mock_motors_bus import MockMotorsBus from tests.mocks.mock_motors_bus import MockMotorsBus
@@ -119,7 +119,7 @@ class MockRobot(Robot):
def configure(self) -> None: def configure(self) -> None:
pass pass
def get_observation(self) -> dict[str, Any]: def get_observation(self) -> RobotObservation:
if not self.is_connected: if not self.is_connected:
raise DeviceNotConnectedError(f"{self} is not 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) 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: if not self.is_connected:
raise DeviceNotConnectedError(f"{self} is not connected.") raise DeviceNotConnectedError(f"{self} is not connected.")
+2 -1
View File
@@ -19,6 +19,7 @@ from dataclasses import dataclass
from functools import cached_property from functools import cached_property
from typing import Any from typing import Any
from lerobot.processor import RobotAction
from lerobot.teleoperators import Teleoperator, TeleoperatorConfig from lerobot.teleoperators import Teleoperator, TeleoperatorConfig
from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
@@ -88,7 +89,7 @@ class MockTeleop(Teleoperator):
def configure(self) -> None: def configure(self) -> None:
pass pass
def get_action(self) -> dict[str, Any]: def get_action(self) -> RobotAction:
if not self.is_connected: if not self.is_connected:
raise DeviceNotConnectedError(f"{self} is not connected.") raise DeviceNotConnectedError(f"{self} is not connected.")
-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()
+2
View File
@@ -19,6 +19,8 @@ from unittest.mock import MagicMock, patch
import numpy as np import numpy as np
import pytest import pytest
pytest.importorskip("reachy2_sdk")
from lerobot.robots.reachy2 import ( from lerobot.robots.reachy2 import (
REACHY2_ANTENNAS_JOINTS, REACHY2_ANTENNAS_JOINTS,
REACHY2_L_ARM_JOINTS, REACHY2_L_ARM_JOINTS,