Compare commits

..

8 Commits

Author SHA1 Message Date
Nikodem Bartnik ee3e796270 fix link 2026-07-29 13:21:16 +02:00
Nikodem Bartnik e548375e4c fix emoji and formatting 2026-07-29 13:11:36 +02:00
Nikodem Bartnik 17d5675738 improved explanation 2026-07-29 13:09:20 +02:00
Nikodem Bartnik ffd0e26fbe text fix 2026-07-29 09:43:25 +02:00
Nikodem Bartnik 9cade7db18 fix logo size 2026-07-28 17:16:34 +02:00
Nikodem Bartnik 775bbc19f4 fix formatting 2026-07-28 17:13:44 +02:00
Nikodem Bartnik 2b78c7e582 bring community section up + bigger logo 2026-07-28 16:28:38 +02:00
Nikodem Bartnik b101ff0e56 extend index.mdx 2026-07-28 16:10:15 +02:00
57 changed files with 457 additions and 1642 deletions
+6 -10
View File
@@ -61,20 +61,16 @@ Full details in [`docs/source/so101.mdx`](./docs/source/so101.mdx) and [`docs/so
**4.1 Install**
```bash
# uv (recommended — see AGENTS.md and CLAUDE.md)
uv sync --locked --extra feetech # SO-100/SO-101 motor stack
# uv sync --locked --extra all # everything
# uv sync --locked --extra smolvla # add SmolVLA deps
# pip (alternative, e.g. when not working from source)
# pip install 'lerobot[feetech]'
# pip install 'lerobot[all]'
# pip install 'lerobot[smolvla]'
pip install 'lerobot[feetech]' # SO-100/SO-101 motor stack
# pip install 'lerobot[all]' # everything
# pip install 'lerobot[aloha,pusht]' # specific features
# pip install 'lerobot[smolvla]' # add SmolVLA deps
git lfs install && git lfs pull
hf auth login # required to push datasets/policies
```
Contributors can alternatively use `uv sync --locked --extra feetech` (see `AGENTS.md`).
**4.2 Find USB ports** — run once per arm, unplug when prompted.
```bash
+5 -4
View File
@@ -68,16 +68,17 @@ ENV HOME=/home/user_lerobot \
# issues with MuJoCo and OpenGL drivers.
RUN uv venv --python python${PYTHON_VERSION}
# Install third-party dependencies separately for layer caching
# Install Python dependencies for caching
COPY --chown=user_lerobot:user_lerobot setup.py pyproject.toml uv.lock README.md MANIFEST.in ./
RUN uv sync --locked --extra all --no-install-project --no-cache
COPY --chown=user_lerobot:user_lerobot src/ src/
RUN uv sync --locked --extra all --no-cache
RUN chmod +x /lerobot/.venv/lib/python${PYTHON_VERSION}/site-packages/triton/backends/nvidia/bin/ptxas
# Copy the application source code and install the local project
# Copy the rest of the application source code
# Make sure to have the git-LFS files for testing
COPY --chown=user_lerobot:user_lerobot . .
RUN uv sync --locked --extra all --no-cache
# Set the default command
CMD ["/bin/bash"]
+5 -4
View File
@@ -60,14 +60,15 @@ ENV HOME=/home/user_lerobot \
# run other Python projects in the same container without dependency conflicts.
RUN uv venv
# Install third-party dependencies separately for layer caching
# Install Python dependencies for caching
COPY --chown=user_lerobot:user_lerobot setup.py pyproject.toml uv.lock README.md MANIFEST.in ./
RUN uv sync --locked --extra all --no-install-project --no-cache
COPY --chown=user_lerobot:user_lerobot src/ src/
# Copy the application code and install the local project
RUN uv sync --locked --extra all --no-cache
# Copy the rest of the application code
# Make sure to have the git-LFS files for testing
COPY --chown=user_lerobot:user_lerobot . .
RUN uv sync --locked --extra all --no-cache
# Set the default command
CMD ["/bin/bash"]
+3 -3
View File
@@ -58,7 +58,7 @@ final_action = postprocessor(action)
## Hardware API redesign
PR [#777](https://github.com/huggingface/lerobot/pull/777) improves the LeRobot calibration but is **not backward-compatible**. Below is an overview of what changed and how you can continue to work with datasets created before this pull request.
PR [#777](https://github.com/huggingface/lerobot/pull/777) improves the LeRobot calibration but is **not backward-compatible**. Below is a overview of what changed and how you can continue to work with datasets created before this pull request.
### What changed?
@@ -129,8 +129,8 @@ python examples/backward_compatibility/replay.py \
Policies output actions in the same format as the datasets (`torch.Tensors`). Therefore, the same transformations should be applied.
To find these transformations, we recommend first replaying an episode of the dataset your policy was trained on using the section above.
Then, add these same transformations to your inference script (shown here in the `record.py` script):
To find these transformations, we recommend to first try and and replay an episode of the dataset your policy was trained on using the section above.
Then, add these same transformations on your inference script (shown here in the `record.py` script):
```diff
action_values = predict_action(
+15 -2
View File
@@ -88,6 +88,20 @@ policy_preprocessor = NormalizerProcessorStep(stats=dataset_stats)
The same policy can work with different environment processors, and the same environment processor can work with different policies:
````python
# Use SmolVLA policy with LIBERO environment
# Use SmolVLA policy with LIBERO environment
libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
env_cfg=libero_cfg,
policy_cfg=smolvla_cfg,
)
smolvla_preprocessor, smolvla_postprocessor = make_pre_post_processors(smolvla_cfg)
# Or use ACT policy with the same LIBERO environment
libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
env_cfg=libero_cfg,
policy_cfg=act_cfg,
)
act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg)
```python
# Use SmolVLA policy with LIBERO environment
libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
@@ -102,7 +116,6 @@ libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
policy_cfg=act_cfg,
)
act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg)
```
### 3. **Easier Experimentation**
@@ -132,7 +145,7 @@ class LiberoVelocityProcessorStep(ObservationProcessorStep):
state = torch.cat([eef_pos, eef_axisangle, eef_vel,
gripper_pos, gripper_vel], dim=-1) # 14D
return state
```
````
### 4. **Cleaner Environment Code**
+4 -4
View File
@@ -40,10 +40,10 @@ This tutorial guides you through updating the firmware of Feetech motors using t
For each motor you want to update:
1. **Select the motor** from the list by clicking on it
2. **Click the Upgrade tab**:
3. **Click the Online button**:
- If a potential firmware update is found, it will be displayed in the box
4. **Click the Upgrade button**:
2. **Click on Upgrade tab**:
3. **Click on Online button**:
- If an potential firmware update is found, it will be displayed in the box
4. **Click on Upgrade button**:
- The update progress will be displayed
## Step 6: Verify Update
+1 -1
View File
@@ -211,7 +211,7 @@ Record, Replay and Train with Hope-JR is still experimental.
### Record
This step records the dataset, which can be seen as an example [here](https://huggingface.co/datasets/nepyope/hand_record_test_with_video_data).
This step records the dataset, which can be seen as an example [here](https://huggingface.co/datasets/nepyope/hand_record_test_with_video_data/settings).
```bash
lerobot-record \
+163 -9
View File
@@ -1,23 +1,177 @@
# LeRobot
<div class="flex justify-center">
<a target="_blank" href="https://huggingface.co/lerobot">
<img
alt="HuggingFace Expert Acceleration Program"
alt="LeRobot, Hugging Face Robotics Library"
src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/lerobot/lerobot-logo-thumbnail.png"
style="width: 100%"
></img>
</a>
</div>
# LeRobot
**State-of-the-art machine learning for real-world robotics**
🤗 LeRobot aims to provide models, datasets, and tools for real-world robotics in PyTorch. The goal is to lower the barrier for entry to robotics so that everyone can contribute and benefit from sharing datasets and pretrained models.
🤗 LeRobot provides a hardware-agnostic, Python-native interface for controlling real robots - from affordable arms like the SO-ARM101 to full humanoids. Plus the tools to record, store, and share the datasets they generate. Every dataset uses the standardized **LeRobotDataset** format (synchronized video + action/state data) and can be streamed directly from the [Hugging Face Hub](https://huggingface.co/lerobot).
🤗 LeRobot contains state-of-the-art approaches that have been shown to transfer to the real-world with a focus on imitation learning and reinforcement learning.
🤗 On top of that data, LeRobot implements state-of-the-art policies - from lightweight imitation-learning models like ACT to large vision-language-action models like π₀ and SmolVLA - all trainable, shareable, and deployable with the same handful of CLI commands.
🤗 LeRobot already provides a set of pretrained models, datasets with human collected demonstrations, and simulated environments so that everyone can get started.
The goal: lower the barrier to entry for robotics, so that everyone can contribute to, and benefit from, shared datasets and pretrained models.
🤗 LeRobot hosts pretrained models and datasets on the LeRobot HuggingFace page.
<div align="center" style="display: flex; justify-content: center; gap: 8px; flex-wrap: wrap; margin: 20px 0;">
<a href="https://discord.gg/s3KuuzsPFb" target="_blank">
<img alt="Discord" src="https://img.shields.io/badge/Discord-Join_the_Community-5865F2?style=flat&logo=discord&logoColor=white">
</a>
<a href="https://x.com/LeRobotHF" target="_blank">
<img alt="X (Twitter)" src="https://img.shields.io/badge/X-Follow_%40LeRobotHF-black?style=flat&logo=x&logoColor=white">
</a>
<a href="https://huggingface.co/lerobot" target="_blank">
<img alt="Hugging Face Hub" src="https://img.shields.io/badge/HF_Hub-Models_%26_Datasets-FFD21E?style=flat">
</a>
</div>
Join the LeRobot community on [Discord](https://discord.gg/s3KuuzsPFb)
<div align="center">
<img src="../../media/readme/robots_control_video.webp" width="640px" alt="Reachy 2 Demo">
</div>
## How It Works
**Teleoperate → Record → Train → Deploy**
1. **Teleoperate** - control the robot yourself (with a leader arm, keyboard, or phone) so it can learn from your movements.
2. **Record** - each demonstration is saved as a dataset: synchronized camera video plus the actions you took.
3. **Train** - a policy (the neural network that will control the robot) learns to imitate your demonstrations.
4. **Deploy** - run the trained policy on the robot and watch it complete the task on its own.
## Get Started
New here? [Install LeRobot](./installation), then pick your path:
<div class="grid grid-cols-1 md:grid-cols-3 gap-4 my-6">
<div class="border dark:border-gray-700 rounded-lg p-4 shadow">
<div class="text-lg font-semibold mb-2">🔧 I have a robot</div>
<p class="text-gray-700 dark:text-gray-300 text-sm">
LeRobot supports a wide range of arms and mobile robots. Popular picks:
</p>
<ul class="text-gray-700 dark:text-gray-300 text-sm list-disc pl-5 mb-2">
<li>
<a href="./so101">SO-101</a> - our flagship, low-cost arm
</li>
<li>
<a href="./lekiwi">LeKiwi</a> - a mobile base with an arm on top
</li>
<li>
<a href="./koch">Koch v1.1</a> - a long-time community favorite
</li>
<li>
or find yours under <strong>Robots</strong> in the sidebar
</li>
</ul>
<p class="text-gray-700 dark:text-gray-300 text-sm">
Once it's assembled and calibrated, record a dataset and train your first
policy with the <a href="./il_robots">imitation learning tutorial</a> - or
skip the CLI entirely with <a href="./lelab">LeLab</a>, a browser GUI for
the same workflow.
</p>
</div>
<div class="border dark:border-gray-700 rounded-lg p-4 shadow">
<div class="text-lg font-semibold mb-2">💻 No hardware yet</div>
<p class="text-gray-700 dark:text-gray-300 text-sm">
You can still train and evaluate policies without owning a robot:
</p>
<ul class="text-gray-700 dark:text-gray-300 text-sm list-disc pl-5 mb-2">
<li>
train on an existing
<a href="https://huggingface.co/datasets?other=LeRobot">
LeRobot dataset
</a>
from the Hub
</li>
<li>
evaluate in <a href="./envhub">simulation</a>, against benchmarks like
LIBERO or Meta-World
</li>
<li>
try the free <a href="./notebooks">Colab notebooks</a> - nothing to
install
</li>
</ul>
</div>
<div class="border dark:border-gray-700 rounded-lg p-4 shadow">
<div class="text-lg font-semibold mb-2">🤝 I want to contribute</div>
<p class="text-gray-700 dark:text-gray-300 text-sm">
Start with the <a href="./contributing">Contributing guide</a>, then
<a href="./bring_your_own_policies">add a new policy</a> or
<a href="./integrate_hardware">bring your own hardware</a>.
</p>
</div>
</div>
## Explore the Docs
<div class="grid grid-cols-1 md:grid-cols-3 gap-4 my-6">
<a
class="!no-underline border dark:border-gray-700 rounded-lg p-4 shadow hover:shadow-lg"
href="./cheat-sheet"
>
<div class="font-semibold mb-1">📋 Cheat Sheet</div>
<p class="text-gray-700 dark:text-gray-300 text-sm">
Every LeRobot CLI command, copy-paste ready.
</p>
</a>
<a
class="!no-underline border dark:border-gray-700 rounded-lg p-4 shadow hover:shadow-lg"
href="./hardware_guide"
>
<div class="font-semibold mb-1">🖥️ Compute & Hardware Guide</div>
<p class="text-gray-700 dark:text-gray-300 text-sm">
Which policy fits your GPU, and how long training takes.
</p>
</a>
<a
class="!no-underline border dark:border-gray-700 rounded-lg p-4 shadow hover:shadow-lg"
href="./lerobot-dataset-v3"
>
<div class="font-semibold mb-1">🗂️ LeRobotDataset</div>
<p class="text-gray-700 dark:text-gray-300 text-sm">
Load, stream, and visualize robot datasets from the Hub.
</p>
</a>
<a
class="!no-underline border dark:border-gray-700 rounded-lg p-4 shadow hover:shadow-lg"
href="./lelab"
>
<div class="font-semibold mb-1">🖼 LeLab</div>
<p class="text-gray-700 dark:text-gray-300 text-sm">
A browser GUI for calibrating, recording, and training - no CLI required.
</p>
</a>
<a
class="!no-underline border dark:border-gray-700 rounded-lg p-4 shadow hover:shadow-lg"
href="./act"
>
<div class="font-semibold mb-1">🧠 Policies</div>
<p class="text-gray-700 dark:text-gray-300 text-sm">
Start with ACT, our recommended first policy - or browse SmolVLA, π₀, and
more in the sidebar.
</p>
</a>
<a
class="!no-underline border dark:border-gray-700 rounded-lg p-4 shadow hover:shadow-lg"
href="./envhub"
>
<div class="font-semibold mb-1">🎮 Simulation & Benchmarks</div>
<p class="text-gray-700 dark:text-gray-300 text-sm">
Train and evaluate in simulated environments before touching real
hardware.
</p>
</a>
</div>
## Common Problems
Running into issues? A few of the most frequent ones:
- **Blurry or unusable camera footage** - lighting matters more than resolution. See the [Cameras](./cameras) guide.
- **Build or install errors** (`cmake`, `ffmpeg`, CUDA) - see the Troubleshooting section of the [Installation guide](./installation#troubleshooting).
- **Not sure which policy fits your GPU** - check the [Compute & Hardware Guide](./hardware_guide).
- **Still stuck?** Ask on [Discord](https://discord.gg/s3KuuzsPFb) - the community (and the LeRobot team) is there to help.
+1 -1
View File
@@ -18,7 +18,7 @@ If you're using Feetech or Dynamixel motors, LeRobot provides built-in bus inter
- [`DynamixelMotorsBus`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/motors/dynamixel/dynamixel.py) for controlling Dynamixel servos
Please refer to the [`MotorsBus`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/motors/motors_bus.py) abstract class to learn about its API.
For a good example of how it can be used, you can have a look at our own [SO101 follower implementation](https://github.com/huggingface/lerobot/blob/main/src/lerobot/robots/so_follower/so_follower.py)
For a good example of how it can be used, you can have a look at our own [SO101 follower implementation](https://github.com/huggingface/lerobot/blob/main/src/lerobot/robots/so_follower/so101_follower/so101_follower.py)
Use these if compatible. Otherwise, you'll need to find or write a Python interface (not covered in this tutorial):
+1 -1
View File
@@ -51,7 +51,7 @@ In addition to these instructions, you need to install the Feetech SDK & ZeroMQ
pip install -e ".[lekiwi]"
```
Great 🤗! You are now done installing LeRobot, and we can begin assembling the SO100/SO101 arms and the mobile base 🤖.
Great :hugs:! You are now done installing LeRobot, and we can begin assembling the SO100/SO101 arms and the mobile base :robot:.
Every time you now want to use LeRobot, you can go to the `~/lerobot` folder where we installed LeRobot and run one of the commands.
# Step-by-Step Assembly Instructions
+1 -1
View File
@@ -174,7 +174,7 @@ The model takes images, text instructions, and robot state as input, and outputs
## Reproducing π₀Fast results
We reproduce the results of π₀Fast on the LIBERO benchmark using the LeRobot implementation. We take the LeRobot PiFast base model [lerobot/pi0fast-base](https://huggingface.co/lerobot/pi0fast-base) and finetune for an additional 40k steps in bfloat16, with batch size of 256 on 8 H100 GPUs using the [HuggingFace LIBERO dataset](https://huggingface.co/datasets/HuggingFaceVLA/libero).
We reproduce the results of π₀Fast on the LIBERO benchmark using the LeRobot implementation. We take the LeRobot PiFast base model [lerobot/pi0fast-base](https://huggingface.co/lerobot/pi0fast-base) and finetune for an additional 40kk steps in bfloat16, with batch size of 256 on 8 H100 GPUs using the [HuggingFace LIBERO dataset](https://huggingface.co/datasets/HuggingFaceVLA/libero).
The finetuned model can be found here:
+3 -3
View File
@@ -22,7 +22,7 @@ With processors, you choose the learning features you want to use for your polic
## Three pipelines
We often compose three pipelines. Depending on your setup, some can be empty if action and observation spaces already match.
Each of these pipelines handles different conversions between different action and observation spaces. Below is a quick explanation of each pipeline.
Each of these pipelines handle different conversions between different action and observation spaces. Below is a quick explanation of each pipeline.
1. Pipeline 1: Teleop action space → dataset action space (phone pose → EE targets)
2. Pipeline 2: Dataset action space → robot command space (EE targets → joints)
@@ -74,7 +74,7 @@ In the phone to SO-100 follower examples we use the following adapters:
- `robot_action_to_transition`: transforms the teleop action dict to a pipeline transition.
- `transition_to_robot_action`: transforms the pipeline transition to a robot action dict.
- `observation_to_transition`: transforms the robot observation dict to a pipeline transition.
- `transition_to_observation`: transforms the pipeline transition to an observation dict.
- `transition_to_observation`: transforms the pipeline transition to a observation dict.
Checkout [src/lerobot/processor/converters.py](https://github.com/huggingface/lerobot/blob/main/src/lerobot/processor/converters.py) for more details.
@@ -82,7 +82,7 @@ Check out [src/lerobot/processor/converters.py](https://github.com/huggingface/l
Dataset features are determined by the keys saved in the dataset. Each step can declare what features it modifies in a contract called `transform_features(...)`. Once you build a processor, the processor can then aggregate all of these features with `aggregate_pipeline_dataset_features()` and merge multiple feature dicts with `combine_feature_dicts(...)`.
Below is an example of how we declare features with the `transform_features` method in the phone to SO-100 follower examples:
Below is and example of how we declare features with the `transform_features` method in the phone to SO-100 follower examples:
```python
def transform_features(
+2 -2
View File
@@ -57,7 +57,7 @@ policy_cfg.rtc_config = RTCConfig(
policy = PI0Policy.from_pretrained("lerobot/pi0_base", policy_cfg=policy_cfg, device="cuda")
# Now use predict_action_chunk with RTC parameters
inference_delay = 4 # How many steps of inference latency, this value should be calculated based on the inference latency of the policy
inference_delay = 4 # How many steps of inference latency, this values should be calculated based on the inference latency of the policy
# Initialize the action queue
action_queue = ActionQueue(policy_cfg.rtc_config)
@@ -100,7 +100,7 @@ Typical values: 8-12 steps
RTCConfig(execution_horizon=10)
```
**`max_guidance_weight`**: How strongly to enforce consistency with the previous chunk. This is a hyperparameter that can be tuned to balance the smoothness of the transitions and the reactivity of the policy. For 10 steps flow matching (SmolVLA, Pi0, Pi0.5), a value of 10.0 is an optimal value.
**`max_guidance_weight`**: How strongly to enforce consistency with the previous chunk. This is a hyperparameter that can be tuned to balance the smoothness of the transitions and the reactivity of the policy. For 10 steps flow matching (SmolVLA, Pi0, Pi0.5), a value of 10.0 is a optimal value.
**`prefix_attention_schedule`**: How to weight consistency across the overlap region.
+1 -1
View File
@@ -93,7 +93,7 @@ lerobot-train --help
## Evaluate the finetuned model and run it in real-time
Similarly for when recording an episode, it is recommended that you are logged in to the HuggingFace Hub. You can follow the corresponding steps: [Record a dataset](./il_robots#record-a-dataset).
Similarly for when recording an episode, it is recommended that you are logged in to the HuggingFace Hub. You can follow the corresponding steps: [Record a dataset](./il_robots).
Once you are logged in, you can run inference in your setup by doing:
```bash
+2 -2
View File
@@ -50,11 +50,11 @@ lerobot-edit-dataset \
Divide a dataset into multiple subsets.
```bash
# Split by fractions (e.g. 60% train, 20% val, 20% test)
# Split by fractions (e.g. 80% train, 20% test, 20% val)
lerobot-edit-dataset \
--repo_id lerobot/pusht \
--operation.type split \
--operation.splits '{"train": 0.6, "val": 0.2, "test": 0.2}'
--operation.splits '{"train": 0.8, "test": 0.2, "val": 0.2}'
# Split by specific episode indices
lerobot-edit-dataset \
-13
View File
@@ -494,19 +494,6 @@ ignore_errors = true
module = "lerobot.envs.*"
ignore_errors = false
[[tool.mypy.overrides]]
module = "lerobot.annotations.*"
ignore_errors = false
disallow_untyped_defs = true
disallow_incomplete_defs = true
check_untyped_defs = true
[[tool.mypy.overrides]]
module = "lerobot.transforms.*"
ignore_errors = false
disallow_untyped_defs = true
disallow_incomplete_defs = true
check_untyped_defs = true
# [[tool.mypy.overrides]]
# module = "lerobot.utils.*"
+40 -96
View File
@@ -19,7 +19,6 @@ import copy
import logging
import shutil
from pathlib import Path
from typing import Any, NotRequired, TypedDict
import datasets
import pandas as pd
@@ -50,32 +49,8 @@ from .utils import (
)
from .video_utils import concatenate_video_files, get_video_duration_in_s
logger = logging.getLogger(__name__)
type FeatureDict = dict[str, dict[str, Any]]
type ChunkFile = tuple[int, int]
class IndexState(TypedDict):
chunk: int
file: int
src_to_dst: NotRequired[dict[ChunkFile, ChunkFile]]
class VideoIndex(TypedDict):
chunk: int
file: int
latest_duration: float
episode_duration: float
src_to_offset: NotRequired[dict[ChunkFile, float]]
src_to_dst: NotRequired[dict[ChunkFile, ChunkFile]]
dst_file_durations: NotRequired[dict[ChunkFile, float]]
type VideoIndexState = dict[str, VideoIndex]
def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMetadata]) -> FeatureDict:
def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMetadata]) -> dict[str, dict]:
"""Create a merged video feature info dictionary for aggregation. The video encoder info is merged field-by-field: each key is kept only when every source agrees; otherwise that key is set to ``null`` (or ``{}`` for ``video.extra_options``) and a warning is logged.
Args:
@@ -84,14 +59,14 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
Returns:
dict: A dictionary of merged video feature info.
"""
merged_info: FeatureDict = copy.deepcopy(all_metadata[0].features)
merged_info = copy.deepcopy(all_metadata[0].features)
video_keys = [k for k in merged_info if merged_info[k].get("dtype") == "video"]
for vk in video_keys:
video_infos = [m.features.get(vk, {}).get("info") or {} for m in all_metadata]
base_video_info = video_infos[0]
merged_encoder_info: dict[str, Any] = {}
merged_encoder_info: dict = {}
fallback_keys: list[str] = []
for info_key in VIDEO_ENCODER_INFO_KEYS:
values = [info.get(info_key, None) for info in video_infos]
@@ -105,7 +80,7 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
merged_encoder_info[info_key] = {} if info_key == "video.extra_options" else None
if fallback_keys:
logger.warning(
logging.warning(
f"Merging heterogeneous or incomplete video encoder metadata for feature {vk}. "
f"Setting these keys to null: {fallback_keys}.",
)
@@ -117,7 +92,7 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
return merged_info
def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]) -> tuple[int, str | None, FeatureDict]:
def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]):
"""Validates that all dataset metadata have consistent properties.
Ensures all datasets have the same fps, robot_type, and features to guarantee
@@ -154,9 +129,7 @@ def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]) -> tuple[i
return fps, robot_type, features
def update_data_df(
df: pd.DataFrame, src_meta: LeRobotDatasetMetadata, dst_meta: LeRobotDatasetMetadata
) -> pd.DataFrame:
def update_data_df(df, src_meta, dst_meta):
"""Updates a data DataFrame with new indices and task mappings for aggregation.
Adjusts episode indices, frame indices, and task indices to account for
@@ -181,12 +154,12 @@ def update_data_df(
def update_meta_data(
df: pd.DataFrame,
dst_meta: LeRobotDatasetMetadata,
meta_idx: IndexState,
data_idx: IndexState,
videos_idx: VideoIndexState,
) -> pd.DataFrame:
df,
dst_meta,
meta_idx,
data_idx,
videos_idx,
):
"""Updates metadata DataFrame with new chunk, file, and timestamp indices.
Adjusts all indices and timestamps to account for previously aggregated
@@ -316,7 +289,7 @@ def aggregate_datasets(
chunk_size: int | None = None,
concatenate_videos: bool = True,
concatenate_data: bool = True,
) -> None:
):
"""Aggregates multiple LeRobot datasets into a single unified dataset.
This is the main function that orchestrates the aggregation process by:
@@ -336,7 +309,7 @@ def aggregate_datasets(
concatenate_videos: When False, keep one mp4 per source file instead of packing into shards.
concatenate_data: When False, keep one parquet per source file instead of packing into shards.
"""
logger.info("Start aggregate_datasets")
logging.info("Start aggregate_datasets")
if data_files_size_in_mb is None:
data_files_size_in_mb = DEFAULT_DATA_FILE_SIZE_IN_MB
@@ -368,15 +341,15 @@ def aggregate_datasets(
video_files_size_in_mb=video_files_size_in_mb,
)
logger.info("Find all tasks")
logging.info("Find all tasks")
unique_tasks = pd.concat([m.tasks for m in all_metadata]).index.unique()
dst_meta.tasks = pd.DataFrame(
{"task_index": range(len(unique_tasks))}, index=pd.Index(unique_tasks, name="task")
)
meta_idx: IndexState = {"chunk": 0, "file": 0}
data_idx: IndexState = {"chunk": 0, "file": 0}
videos_idx: VideoIndexState = {
meta_idx = {"chunk": 0, "file": 0}
data_idx = {"chunk": 0, "file": 0}
videos_idx = {
key: {"chunk": 0, "file": 0, "latest_duration": 0, "episode_duration": 0} for key in video_keys
}
@@ -400,17 +373,12 @@ def aggregate_datasets(
dst_meta.info.total_frames += src_meta.total_frames
finalize_aggregation(dst_meta, all_metadata)
logger.info("Aggregation complete.")
logging.info("Aggregation complete.")
def aggregate_videos(
src_meta: LeRobotDatasetMetadata,
dst_meta: LeRobotDatasetMetadata,
videos_idx: VideoIndexState,
video_files_size_in_mb: float,
chunk_size: int,
concatenate_videos: bool = True,
) -> VideoIndexState:
src_meta, dst_meta, videos_idx, video_files_size_in_mb, chunk_size, concatenate_videos=True
):
"""Aggregates video chunks from a source dataset into the destination dataset.
Handles video file concatenation and rotation based on file size limits.
@@ -438,8 +406,7 @@ def aggregate_videos(
videos_idx[key]["dst_file_durations"] = {}
for key, video_idx in videos_idx.items():
unique_chunk_file_pairs: list[ChunkFile] = sorted(
{
unique_chunk_file_pairs = {
(chunk, file)
for chunk, file in zip(
src_meta.episodes[f"videos/{key}/chunk_index"],
@@ -447,7 +414,7 @@ def aggregate_videos(
strict=False,
)
}
)
unique_chunk_file_pairs = sorted(unique_chunk_file_pairs)
chunk_idx = video_idx["chunk"]
file_idx = video_idx["file"]
@@ -522,14 +489,7 @@ def aggregate_videos(
return videos_idx
def aggregate_data(
src_meta: LeRobotDatasetMetadata,
dst_meta: LeRobotDatasetMetadata,
data_idx: IndexState,
data_files_size_in_mb: float,
chunk_size: int,
concatenate_data: bool = True,
) -> IndexState:
def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_size, concatenate_data=True):
"""Aggregates data chunks from a source dataset into the destination dataset.
Reads source data files, updates indices to match the aggregated dataset,
@@ -550,16 +510,14 @@ def aggregate_data(
Returns:
dict: Updated data_idx with current chunk and file indices.
"""
unique_chunk_file_ids: list[ChunkFile] = sorted(
{
unique_chunk_file_ids = {
(c, f)
for c, f in zip(
src_meta.episodes["data/chunk_index"],
src_meta.episodes["data/file_index"],
strict=False,
src_meta.episodes["data/chunk_index"], src_meta.episodes["data/file_index"], strict=False
)
}
)
unique_chunk_file_ids = sorted(unique_chunk_file_ids)
contains_images = len(dst_meta.image_keys) > 0
# retrieve features schema for proper image typing in parquet
@@ -567,7 +525,7 @@ def aggregate_data(
# Track source to destination file mapping for metadata update
# This is critical for handling datasets that are already results of a merge
src_to_dst: dict[ChunkFile, ChunkFile] = {}
src_to_dst: dict[tuple[int, int], tuple[int, int]] = {}
for src_chunk_idx, src_file_idx in unique_chunk_file_ids:
src_path = src_meta.root / DEFAULT_DATA_PATH.format(
@@ -606,13 +564,7 @@ def aggregate_data(
return data_idx
def aggregate_metadata(
src_meta: LeRobotDatasetMetadata,
dst_meta: LeRobotDatasetMetadata,
meta_idx: IndexState,
data_idx: IndexState,
videos_idx: VideoIndexState,
) -> IndexState:
def aggregate_metadata(src_meta, dst_meta, meta_idx, data_idx, videos_idx):
"""Aggregates metadata from a source dataset into the destination dataset.
Reads source metadata files, updates all indices and timestamps,
@@ -628,8 +580,7 @@ def aggregate_metadata(
Returns:
dict: Updated meta_idx with current chunk and file indices.
"""
chunk_file_ids: list[ChunkFile] = sorted(
{
chunk_file_ids = {
(c, f)
for c, f in zip(
src_meta.episodes["meta/episodes/chunk_index"],
@@ -637,7 +588,8 @@ def aggregate_metadata(
strict=False,
)
}
)
chunk_file_ids = sorted(chunk_file_ids)
for chunk_idx, file_idx in chunk_file_ids:
src_path = src_meta.root / DEFAULT_EPISODES_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
df = pd.read_parquet(src_path)
@@ -670,16 +622,16 @@ def aggregate_metadata(
def append_or_create_parquet_file(
df: pd.DataFrame,
src_path: Path,
idx: IndexState,
idx: dict[str, int],
max_mb: float,
chunk_size: int,
default_path: str,
contains_images: bool = False,
aggr_root: Path | None = None,
aggr_root: Path = None,
hf_features: datasets.Features | None = None,
concatenate: bool = True,
one_row_group_per_episode: bool = False,
) -> tuple[IndexState, ChunkFile]:
) -> tuple[dict[str, int], tuple[int, int]]:
"""Appends data to an existing parquet file or creates a new one based on size constraints.
Manages file rotation when size limits are exceeded to prevent individual files
@@ -702,13 +654,7 @@ def append_or_create_parquet_file(
Returns:
tuple: (updated_idx, (dst_chunk, dst_file)) where updated_idx is the index dict
and (dst_chunk, dst_file) is the actual destination file the data was written to.
Raises:
ValueError: If aggr_root is not provided.
"""
if aggr_root is None:
raise ValueError("aggr_root must be provided.")
dst_chunk, dst_file = idx["chunk"], idx["file"]
dst_path = aggr_root / default_path.format(chunk_index=dst_chunk, file_index=dst_file)
@@ -752,9 +698,7 @@ def append_or_create_parquet_file(
return idx, (dst_chunk, dst_file)
def finalize_aggregation(
aggr_meta: LeRobotDatasetMetadata, all_metadata: list[LeRobotDatasetMetadata]
) -> None:
def finalize_aggregation(aggr_meta, all_metadata):
"""Finalizes the dataset aggregation by writing summary files and statistics.
Writes the tasks file, info file with total counts and splits, and
@@ -764,16 +708,16 @@ def finalize_aggregation(
aggr_meta: Aggregated dataset metadata.
all_metadata: List of all source dataset metadata objects.
"""
logger.info("write tasks")
logging.info("write tasks")
write_tasks(aggr_meta.tasks, aggr_meta.root)
logger.info("write info")
logging.info("write info")
aggr_meta.info.total_tasks = len(aggr_meta.tasks)
aggr_meta.info.total_episodes = sum(m.total_episodes for m in all_metadata)
aggr_meta.info.total_frames = sum(m.total_frames for m in all_metadata)
aggr_meta.info.splits = {"train": f"0:{sum(m.total_episodes for m in all_metadata)}"}
write_info(aggr_meta.info, aggr_meta.root)
logger.info("write stats")
logging.info("write stats")
aggr_meta.stats = aggregate_stats([m.stats for m in all_metadata])
write_stats(aggr_meta.stats, aggr_meta.root)
+2 -2
View File
@@ -188,8 +188,8 @@ class LeRobotDatasetMetadata:
def _load_metadata(self):
self.info = load_info(self.root)
check_version_compatibility(self.repo_id, self._version, CODEBASE_VERSION)
self.tasks = load_tasks(self.root) if self.total_tasks > 0 else None
self.episodes = load_episodes(self.root) if self.total_episodes > 0 else None
self.tasks = load_tasks(self.root)
self.episodes = load_episodes(self.root)
self.stats = load_stats(self.root)
def ensure_readable(self) -> None:
-5
View File
@@ -384,12 +384,7 @@ class LiberoEnv(gym.Env):
def close(self):
if self._env is not None:
try:
self._env.close()
finally:
# LIBERO deletes its inner env on close, so this wrapper must
# be recreated before the next reset.
self._env = None
def _make_env_fns(
+1 -3
View File
@@ -384,9 +384,7 @@ class RoboTwinEnv(gym.Env):
self._env: Any | None = None # deferred — created on first reset() inside worker
self._step_count: int = 0
self._black_frame: np.ndarray = np.zeros(
(self.observation_height, self.observation_width, 3), dtype=np.uint8
)
self._black_frame = np.zeros((self.observation_height, self.observation_width, 3), dtype=np.uint8)
image_spaces = {
cam: spaces.Box(
+1 -1
View File
@@ -373,7 +373,7 @@ class VLABenchEnv(gym.Env):
if action.shape[0] != 7:
# Unknown layout — fall back to zero-pad so the sim doesn't crash.
padded: np.ndarray = np.zeros(ctrl_dim, dtype=np.float64)
padded = np.zeros(ctrl_dim, dtype=np.float64)
padded[: min(action.shape[0], ctrl_dim)] = action[:ctrl_dim]
return padded
+5 -35
View File
@@ -302,33 +302,6 @@ def _pad_evo1_stats(
return padded_stats
def _refresh_evo1_normalization_steps(
config: Evo1Config,
preprocessor: PolicyProcessorPipeline,
postprocessor: PolicyProcessorPipeline,
) -> None:
"""Re-pad checkpoint-loaded (un)normalizer stats/features to EVO1's fixed widths.
Loading a checkpoint injects the raw dataset stats (unpadded to max_state_dim/max_action_dim)
into the (un)normalizer via the generic override path in make_pre_post_processors. Those stats
and their declared features must be re-padded/reshaped to EVO1's fixed widths, otherwise
normalization fails against the padded state/action tensors (e.g. state padded to 24 vs. 8-dim
LIBERO stats). Padding is a no-op when stats are already at the target width.
"""
normalization_features = _evo1_normalization_features(config)
action_features = _evo1_action_features(config)
for step in preprocessor.steps:
if isinstance(step, NormalizerProcessorStep):
step.features = normalization_features
step.stats = _pad_evo1_stats(config, step.stats)
step.to(device=step.device, dtype=step.dtype)
for step in postprocessor.steps:
if isinstance(step, UnnormalizerProcessorStep):
step.features = action_features
step.stats = _pad_evo1_stats(config, step.stats)
step.to(device=step.device, dtype=step.dtype)
def reconcile_evo1_processors(
config: Evo1Config,
preprocessor: PolicyProcessorPipeline,
@@ -336,19 +309,16 @@ def reconcile_evo1_processors(
) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]:
"""Reconcile checkpoint-loaded pipelines with the current EVO1 config.
Three things cannot be restored from a serialized pipeline alone: the EVO1 batch converter
(converters are plain functions and are never serialized), eval-time CLI overrides of the
action postprocessing flags (`postprocess_action_dim`, `binarize_gripper`, `gripper_*`), and the
(un)normalizer stats/features when the generic override path injects raw, unpadded dataset
stats. This restores the converter, re-pads the normalization stats to EVO1's fixed widths, and
rebuilds the action step from the current config so those overrides take effect.
Two things cannot be restored from a serialized pipeline alone: the EVO1 batch converter
(converters are plain functions and are never serialized), and eval-time CLI overrides of the
action postprocessing flags (`postprocess_action_dim`, `binarize_gripper`, `gripper_*`). This
restores the converter and rebuilds the action step from the current config so those overrides
take effect.
"""
# Pipelines reloaded from a checkpoint come back with the default batch converter, which drops
# non-observation extras (embodiment_id, state_mask, custom task fields) needed by EVO1.
preprocessor.to_transition = evo1_batch_to_transition
_refresh_evo1_normalization_steps(config, preprocessor, postprocessor)
action_step = Evo1ActionProcessorStep(
action_dim=_evo1_action_dim(config),
binarize_gripper=config.binarize_gripper,
+2 -2
View File
@@ -18,7 +18,7 @@ import functools
import threading
from collections.abc import Callable, Sequence
from contextlib import suppress
from typing import NotRequired, TypedDict
from typing import TypedDict
import torch
import torch.nn.functional as F # noqa: N812
@@ -36,7 +36,7 @@ class BatchTransition(TypedDict):
next_state: dict[str, torch.Tensor]
done: torch.Tensor
truncated: torch.Tensor
complementary_info: NotRequired[dict[str, torch.Tensor | float | int] | None]
complementary_info: dict[str, torch.Tensor | float | int] | None = None
def random_crop_vectorized(images: torch.Tensor, output_size: tuple) -> torch.Tensor:
@@ -46,12 +46,6 @@ class SOFollowerConfig:
position_i_coefficient: int = 0
position_d_coefficient: int = 32
# Number of extra attempts when a `sync_read` of the motors fails. Feetech buses can occasionally
# return a corrupted status packet ("Incorrect status packet!"), especially when several joints move
# at once, which otherwise aborts the control loop. Retries are immediate (no sleep) and only happen on
# failure, so the steady-state read cost is unchanged.
num_read_retries: int = 2
@RobotConfig.register_subclass("so101_follower")
@RobotConfig.register_subclass("so100_follower")
@@ -510,10 +510,10 @@ class ForwardKinematicsJointsToEEAction(RobotActionProcessorStep):
# We only use the ee pose in the dataset, so we don't need the joint positions
for n in self.motor_names:
features[PipelineFeatureType.ACTION].pop(f"{n}.pos", None)
# Store end-effector features as actions in the dataset schema
# We specify the dataset features of this step that we want to be stored in the dataset
for k in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]:
features[PipelineFeatureType.ACTION][f"ee.{k}"] = PolicyFeature(
type=FeatureType.ACTION, shape=(1,)
type=FeatureType.STATE, shape=(1,)
)
return features
@@ -180,7 +180,7 @@ class SOFollower(Robot):
def get_observation(self) -> RobotObservation:
# Read arm position
start = time.perf_counter()
obs_dict = self.bus.sync_read("Present_Position", num_retry=self.config.num_read_retries)
obs_dict = self.bus.sync_read("Present_Position")
obs_dict = {f"{motor}.pos": val for motor, val in obs_dict.items()}
dt_ms = (time.perf_counter() - start) * 1e3
logger.debug(f"{self} read state: {dt_ms:.1f}ms")
@@ -221,7 +221,7 @@ class SOFollower(Robot):
# Cap goal position when too far away from present position.
# /!\ Slower fps expected due to reading from the follower.
if self.config.max_relative_target is not None:
present_pos = self.bus.sync_read("Present_Position", num_retry=self.config.num_read_retries)
present_pos = self.bus.sync_read("Present_Position")
goal_present_pos = {key: (g_pos, present_pos[key]) for key, g_pos in goal_pos.items()}
goal_pos = ensure_safe_goal_position(goal_present_pos, self.config.max_relative_target)
@@ -68,10 +68,6 @@ class UnitreeG1Config(RobotConfig):
# Compensates for gravity on the unitree's arms using the arm ik solver
gravity_compensation: bool = False
# Locomotion controller class name, e.g. "GrootLocomotionController",
# "HolosomaLocomotionController", or "SonicWholeBodyController". None disables it.
# Selecting "SonicWholeBodyController" implicitly switches the robot to the 64-D
# latent-token action/observation interface (``motion_token.{i}.pos`` action and a
# ``motion_token_state.{i}.pos`` state echo) so ``lerobot-rollout`` can drive a
# policy trained on SONIC motion tokens (e.g. nepyope/sonic_walk).
# Lower-body controller class name, e.g. "GrootLocomotionController" or
# "HolosomaLocomotionController". None disables it.
controller: str | None = None
@@ -1,27 +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.
"""Unitree G1 locomotion controllers (Groot, Holosoma, SONIC)."""
from .gr00t_locomotion import GrootLocomotionController
from .holosoma_locomotion import HolosomaLocomotionController
from .sonic_whole_body import SonicWholeBodyController
__all__ = [
"GrootLocomotionController",
"HolosomaLocomotionController",
"SonicWholeBodyController",
]
@@ -1,378 +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.
"""SONIC decoder whole-body controller for the Unitree G1 (token-only).
Pure-Python/ONNX re-implementation of the *decode* half of NVIDIA's SONIC deploy stack.
The encoder is intentionally absent: a token-output VLA (e.g. ``nepyope/sonic_walk``)
supplies the 64-D latent ``motion_token`` directly each tick, and the SONIC **decoder**
maps ``token + recent proprioception history`` to a residual action that is scaled and
added onto the standing pose (``default_angles``) to produce 50 Hz joint-position targets
for the robot's PD controller.
Index spaces: joints exist in two orderings — **IsaacLab** (policy/training order) and
**MuJoCo** (deploy order). ``ISAACLAB_TO_MUJOCO`` / ``MUJOCO_TO_ISAACLAB`` (in g1_utils)
convert between them. Quaternions are scalar-first ``(w, x, y, z)``.
"""
from __future__ import annotations
import json
import logging
import numpy as np
import onnx
import onnxruntime as ort
from huggingface_hub import hf_hub_download
from ..g1_utils import (
ISAACLAB_TO_MUJOCO,
MUJOCO_TO_ISAACLAB,
G1_29_JointIndex,
get_gravity_orientation,
)
from ..unitree_g1 import lowstate_to_obs
logger = logging.getLogger(__name__)
# ── Constants (hardware-validated; see the NVIDIA SONIC deploy reference) ──────
CONTROL_DT = 0.02 # 50 Hz control period (s)
TOKEN_DIM = 64 # decoder latent size
# SONIC decoder checkpoint: NVIDIA's decoder ONNX re-packaged with its deploy constants
# (kp/kd PD gains, the standing pose default_angles, and the residual action_scale) embedded
# in the ONNX metadata; see upload_sonic_decoder.py for provisioning. The runtime loads the
# model *and* all of these straight from the checkpoint (the Holosoma convention), so no
# motor-physics math happens at deploy time.
DEFAULT_SONIC_REPO_ID = "lerobot/sonic_decoder"
DECODER_FILENAME = "model_decoder.onnx"
DECODER_INPUT_DIM = 994 # token(64) + 10-frame proprio history + gravity
def load_sonic_decoder(repo_id: str = DEFAULT_SONIC_REPO_ID):
"""Load the SONIC decoder ONNX and its baked-in deploy constants from the checkpoint.
Returns ``(decoder_session, kp, kd, default_angles, action_scale, neutral_token)``. The
gains/pose/scale are (29,) float32 in IsaacLab joint order and ``neutral_token`` is the
(64,) float32 idle latent -- all read from the ONNX ``metadata_props`` rather than
recomputed/hardcoded at deploy time (mirrors ``holosoma_locomotion.load_policy``).
"""
decoder_path = hf_hub_download(repo_id=repo_id, filename=DECODER_FILENAME)
so = ort.SessionOptions()
so.log_severity_level = 3 # quiet ORT logs
session = ort.InferenceSession(decoder_path, sess_options=so)
dec_dim = int(session.get_inputs()[0].shape[1])
if dec_dim != DECODER_INPUT_DIM:
raise RuntimeError(f"Unexpected decoder input dim {dec_dim} (expected {DECODER_INPUT_DIM})")
meta = {p.key: p.value for p in onnx.load(decoder_path, load_external_data=False).metadata_props}
required = ("kp", "kd", "default_angles", "action_scale", "neutral_token")
missing = [k for k in required if k not in meta]
if missing:
raise ValueError(
f"SONIC decoder ONNX at {repo_id} is missing metadata {missing}; "
"re-run upload_sonic_decoder.py to (re)provision the checkpoint."
)
arr = {k: np.array(json.loads(meta[k]), dtype=np.float32) for k in required}
logger.info("Loaded SONIC deploy constants from %s (%d joints)", repo_id, len(arr["kp"]))
return session, arr["kp"], arr["kd"], arr["default_angles"], arr["action_scale"], arr["neutral_token"]
def _to_mujoco(a):
"""Apply the ``MUJOCO_TO_ISAACLAB`` gather to a 29-vector (deploy-order reorder).
NOTE: this returns ``a[MUJOCO_TO_ISAACLAB]``. The ``_mj`` suffixes and the exact
permutation direction are a fixed convention validated against the deployed SONIC ONNX
policy (the decoder consumes vectors in this order). Do not "correct" the table or
rename toward the opposite direction without re-validating on hardware.
"""
return a[MUJOCO_TO_ISAACLAB]
# Action-feature prefix for the latent-token interface (see _extract_token_from_action).
TOKEN_ACTION_PREFIX = "motion_token" # nosec B105 - feature-key prefix, not a secret
# Proprio-state prefix for the token interface: the robot echoes the last commanded token
# here so ``lerobot-rollout`` aggregates it into a 64-D ``observation.state``.
TOKEN_STATE_PREFIX = "motion_token_state" # nosec B105 - feature-key prefix, not a secret
def token_action_key(i: int) -> str:
"""Action-dict key for the i-th component of the 64-D SONIC latent token.
The ``.pos`` suffix is required so the value flows through ``lerobot-rollout``, which
only routes ``.pos`` scalar features onto the policy action vector.
"""
return f"{TOKEN_ACTION_PREFIX}.{i}.pos"
def token_state_key(i: int) -> str:
"""Observation key for the i-th component of the 64-D SONIC latent token state."""
return f"{TOKEN_STATE_PREFIX}.{i}.pos"
# Startup blend duration: over the first control ticks, linearly interpolate every joint
# from the robot's initial measured pose into the policy's commanded target, so control
# eases in without a snap on the first command.
INIT_RAMP_S = 3.0
def _extract_token_from_action(action: dict | None) -> np.ndarray | None:
"""Reassemble a dense (64,) latent token from ``motion_token.{i}`` keys, or None.
The token-only interface: the caller supplies the 64-D encoder latent directly (e.g. a
token-output VLA's action), which the decoder consumes with the encoder bypassed.
Requires the full dense token; a partial one is ignored (returns None).
"""
if not action:
return None
keys = [token_action_key(i) for i in range(TOKEN_DIM)]
if any(key not in action for key in keys):
return None
return np.fromiter((float(action[key]) for key in keys), dtype=np.float32, count=TOKEN_DIM)
class SonicDecoder:
"""Runs the SONIC decoder ONNX model and owns the proprioception history.
Each tick it appends the latest robot state to 10-frame history buffers, then maps the
supplied 64-D ``token`` + that history to a residual action added onto ``default_angles``.
The encoder is bypassed entirely (token supplied by the policy). ``default_angles`` and
``action_scale`` are (29,) float32 in IsaacLab order, loaded from the checkpoint.
"""
def __init__(self, decoder, default_angles, action_scale):
self.decoder = decoder
self.decoder_input = decoder.get_inputs()[0].name
self.default_angles = np.asarray(default_angles, np.float32)
self.action_scale = np.asarray(action_scale, np.float32)
self.default_angles_mj = _to_mujoco(self.default_angles)
self.token = np.zeros(TOKEN_DIM, np.float32)
self.last_action_mj = np.zeros(29, np.float32)
self.h_q_mj = [np.zeros(29, np.float32)] * 10
self.h_dq_mj = [np.zeros(29, np.float32)] * 10
self.h_ang = [np.zeros(3, np.float32)] * 10
self.h_act_mj = [np.zeros(29, np.float32)] * 10
self.h_quat = [np.array([1, 0, 0, 0], np.float32)] * 10
def reset(self):
"""Clear the token and 10-frame proprioception history.
``UnitreeG1.reset()`` relies on this so the first decoder outputs of a new episode
are not contaminated by the previous episode's state.
"""
self.token = np.zeros(TOKEN_DIM, np.float32)
self.last_action_mj = np.zeros(29, np.float32)
self.h_q_mj = [np.zeros(29, np.float32)] * 10
self.h_dq_mj = [np.zeros(29, np.float32)] * 10
self.h_ang = [np.zeros(3, np.float32)] * 10
self.h_act_mj = [np.zeros(29, np.float32)] * 10
self.h_quat = [np.array([1, 0, 0, 0], np.float32)] * 10
def update_history(self, q, dq, ang, quat):
"""Push the latest proprioception (pos/vel/gyro/orientation) into the 10-frame buffers."""
quat = quat / (np.linalg.norm(quat) + 1e-8)
q_mj = _to_mujoco(q)
dq_mj = _to_mujoco(dq)
self.h_q_mj = [q_mj - self.default_angles_mj] + self.h_q_mj[:-1]
self.h_dq_mj = [dq_mj] + self.h_dq_mj[:-1]
self.h_ang = [ang.copy()] + self.h_ang[:-1]
self.h_act_mj = [self.last_action_mj.copy()] + self.h_act_mj[:-1]
self.h_quat = [quat.copy()] + self.h_quat[:-1]
def build_decoder_obs(self):
"""Assemble the 994-D decoder input: token + 10-frame proprioception history + gravity."""
obs = np.zeros(994, np.float32)
off = 0
obs[off : off + 64] = self.token
off += 64
for h, sz in [
(list(reversed(self.h_ang)), 3),
(list(reversed(self.h_q_mj)), 29),
(list(reversed(self.h_dq_mj)), 29),
(list(reversed(self.h_act_mj)), 29),
]:
for f in range(10):
obs[off : off + sz] = h[f]
off += sz
for q in reversed(self.h_quat):
obs[off : off + 3] = get_gravity_orientation(q)
off += 3
assert off == 994, f"Decoder obs mismatch: {off}"
return obs
def step(self, robot_obs, token, debug=False):
"""One control tick: read robot obs, decode the supplied token -> joint targets.
Args:
robot_obs: dict with ``<joint>.q``/``.dq`` and ``imu.*`` fields.
token: 64-D latent supplied by the policy (encoder bypassed).
debug: log action/delta norms.
Returns:
dict of ``<joint>.q`` target positions (rad) in IsaacLab joint order.
"""
self.token = np.asarray(token, np.float32)
jnames = [m.name for m in G1_29_JointIndex]
q = np.array(
[
robot_obs.get(f"{n}.q", self.default_angles[m.value])
for m, n in zip(G1_29_JointIndex, jnames, strict=False)
],
np.float32,
)
dq = np.array([robot_obs.get(f"{n}.dq", 0.0) for n in jnames], np.float32)
quat = np.array(
[
robot_obs.get("imu.quat.w", 1),
robot_obs.get("imu.quat.x", 0),
robot_obs.get("imu.quat.y", 0),
robot_obs.get("imu.quat.z", 0),
],
np.float32,
)
ang = np.array([robot_obs.get(f"imu.gyro.{a}", 0) for a in "xyz"], np.float32)
self.update_history(q, dq, ang, quat)
action_mj = (
self.decoder.run(None, {self.decoder_input: self.build_decoder_obs().reshape(1, -1)})[0]
.squeeze()
.astype(np.float32)
)
self.last_action_mj = action_mj.copy()
target = self.default_angles + action_mj[ISAACLAB_TO_MUJOCO] * self.action_scale
if debug:
delta = target - q
logger.debug(
"token_norm=%.4f action_norm=%.4f delta_max=%.4f delta_rms=%.4f",
np.linalg.norm(self.token),
np.linalg.norm(action_mj),
np.max(np.abs(delta)),
np.sqrt(np.mean(delta**2)),
)
return {f"{m.name}.q": float(target[m.value]) for m in G1_29_JointIndex}
class SonicRuntime:
"""Loads the SONIC decoder ONNX model and owns the decode controller.
Token-only deploy: the encoder is bypassed; each tick the decoder consumes a 64-D
latent token supplied directly by the policy.
"""
def __init__(self):
decoder_sess, self.kp, self.kd, default_angles, action_scale, neutral_token = load_sonic_decoder()
self.default_angles = default_angles
self.neutral_token = neutral_token
self.controller = SonicDecoder(decoder_sess, default_angles, action_scale)
@property
def pipeline(self):
return self.controller
def reset(self):
self.controller.reset()
def shutdown(self):
pass
class SonicWholeBodyController:
"""Full-body SONIC controller for UnitreeG1's background controller thread."""
control_dt = CONTROL_DT
full_body = True
def __init__(self):
logger.info("Loading SONIC whole-body controller...")
self._runtime = SonicRuntime()
self.kp = self._runtime.kp
self.kd = self._runtime.kd
self.controller = self._runtime.controller
self._default_angles = self._runtime.default_angles
self._neutral_token = self._runtime.neutral_token
# Startup blend: ease from the robot's initial pose into the first commanded policy
# targets over INIT_RAMP_S (captured on the first control tick).
self._init_ramp_steps = max(1, round(INIT_RAMP_S / CONTROL_DT))
self._init_step = 0
self._start_pose: dict[str, float] = {}
# Token-interface state. ``token_mode`` is set True by the robot whenever a SONIC
# whole-body controller is selected (token-driven deploy): the controller then holds a
# stable *neutral* token until the first real token arrives, and afterwards holds the
# *last* token received between ticks (the async controller runs ~50 Hz while a token
# VLA streams ~30 Hz). This lives here (not in the entry-point script) so it applies
# uniformly to run_g1_server, lerobot-rollout and the sim replays.
self.token_mode = False
self._last_token: np.ndarray | None = None
logger.info("SONIC ready (decoder, 64-D token command path)")
def _startup_blend(self, obs: dict, out: dict) -> dict:
"""Ease into policy control at startup: for the first ``INIT_RAMP_S`` seconds,
interpolate between the robot's pose captured on the first tick and the policy's
live commanded target, so the handoff has no snap.
``out`` is the policy's ``<joint>.q`` target dict for this tick; the blend ratio
climbs 0->1 over the ramp, after which the raw policy target passes through.
"""
if self._init_step >= self._init_ramp_steps or not out:
return out
if self._init_step == 0:
# Capture the robot's actual pose as the interpolation start point.
self._start_pose = {
f"{m.name}.q": float(obs.get(f"{m.name}.q", self._default_angles[m.value]))
for m in G1_29_JointIndex
}
self._init_step += 1
ratio = min(1.0, self._init_step / self._init_ramp_steps)
blended = {
k: self._start_pose.get(k, float(tgt)) * (1.0 - ratio) + float(tgt) * ratio
for k, tgt in out.items()
}
if self._init_step >= self._init_ramp_steps:
logger.info("SONIC startup blend complete -> full policy control")
return blended
def run_step(self, action: dict, lowstate) -> dict:
if lowstate is None:
return {}
obs = lowstate_to_obs(lowstate)
# Token-only interface (token-output VLA): a dense 64-D ``motion_token.{i}`` command
# is decoded directly, encoder bypassed.
token = _extract_token_from_action(action)
if token is not None:
self._last_token = token
elif self._last_token is None and self.token_mode:
# Token-driven deploy, but no token has arrived yet: hold the checkpoint's neutral
# token, which the decoder maps to a stable, natural standing pose.
self._last_token = self._neutral_token.copy()
if self._last_token is None:
# No token yet and not in token_mode: hold (keep last target).
return {}
# Either a fresh token this tick or the last one received (held between the ~30 Hz
# token stream and the ~50 Hz control loop).
return self._startup_blend(obs, self.controller.step(obs, self._last_token))
def reset(self):
self._runtime.reset()
self._init_step = 0 # re-run the startup blend after a reset
self._start_pose = {}
# Drop the held token so token_mode re-seeds the neutral token after a reset.
self._last_token = None
def shutdown(self):
self._runtime.shutdown()
+2 -44
View File
@@ -23,47 +23,6 @@ import numpy as np
NUM_MOTORS = 29
# Joint-order permutations between the two 29-DoF layouts used across the G1 stack:
# IsaacLab (policy/training order) and MuJoCo (deploy order). ``a[ISAACLAB_TO_MUJOCO]``
# reorders an IsaacLab-ordered vector into MuJoCo order, and vice-versa.
ISAACLAB_TO_MUJOCO = np.array(
[
0,
3,
6,
9,
13,
17,
1,
4,
7,
10,
14,
18,
2,
5,
8,
11,
15,
19,
21,
23,
25,
27,
12,
16,
20,
22,
24,
26,
28,
],
dtype=np.int32,
)
# The two orderings are inverses of each other, so derive one from the other (argsort) to
# guarantee they can never drift out of sync.
MUJOCO_TO_ISAACLAB = np.argsort(ISAACLAB_TO_MUJOCO).astype(np.int32)
REMOTE_AXES = ("remote.lx", "remote.ly", "remote.rx", "remote.ry")
REMOTE_BUTTONS = tuple(f"remote.button.{i}" for i in range(16))
REMOTE_KEYS = REMOTE_AXES + REMOTE_BUTTONS
@@ -109,9 +68,8 @@ def make_locomotion_controller(name: str | None):
if name is None:
return None
controllers = {
"GrootLocomotionController": "lerobot.robots.unitree_g1.controllers.gr00t_locomotion",
"HolosomaLocomotionController": "lerobot.robots.unitree_g1.controllers.holosoma_locomotion",
"SonicWholeBodyController": "lerobot.robots.unitree_g1.controllers.sonic_whole_body",
"GrootLocomotionController": "lerobot.robots.unitree_g1.gr00t_locomotion",
"HolosomaLocomotionController": "lerobot.robots.unitree_g1.holosoma_locomotion",
}
module_path = controllers.get(name)
if module_path is None:
@@ -21,7 +21,7 @@ import numpy as np
import onnxruntime as ort
from huggingface_hub import hf_hub_download
from ..g1_utils import (
from .g1_utils import (
REMOTE_AXES,
REMOTE_BUTTONS,
G1_29_JointIndex,
@@ -22,7 +22,7 @@ import onnx
import onnxruntime as ort
from huggingface_hub import hf_hub_download
from ..g1_utils import (
from .g1_utils import (
REMOTE_AXES,
G1_29_JointArmIndex,
G1_29_JointIndex,
+9 -171
View File
@@ -34,6 +34,7 @@ from .config_unitree_g1 import UnitreeG1Config
from .g1_kinematics import G1_29_ArmIK
from .g1_utils import (
REMOTE_AXES,
REMOTE_KEYS,
G1_29_JointArmIndex,
G1_29_JointIndex,
default_remote_input,
@@ -105,47 +106,6 @@ class G1_29_LowState: # noqa: N801
mode_machine: int = 0 # Robot mode
def lowstate_to_obs(lowstate) -> dict:
"""Build a robot observation dict from a Unitree lowstate.
Shared by ``UnitreeG1.get_observation`` and the SONIC pipeline so the
lowstate -> obs mapping lives in exactly one place. Keys match the
``<joint>.q``/``imu.*`` schema consumed across the controllers.
"""
obs: dict = {}
for motor in G1_29_JointIndex:
idx = motor.value
obs[f"{motor.name}.q"] = lowstate.motor_state[idx].q
obs[f"{motor.name}.dq"] = lowstate.motor_state[idx].dq
obs[f"{motor.name}.tau"] = lowstate.motor_state[idx].tau_est
imu = lowstate.imu_state
if imu.gyroscope:
obs["imu.gyro.x"] = imu.gyroscope[0]
obs["imu.gyro.y"] = imu.gyroscope[1]
obs["imu.gyro.z"] = imu.gyroscope[2]
if imu.accelerometer:
obs["imu.accel.x"] = imu.accelerometer[0]
obs["imu.accel.y"] = imu.accelerometer[1]
obs["imu.accel.z"] = imu.accelerometer[2]
if imu.quaternion:
obs["imu.quat.w"] = imu.quaternion[0]
obs["imu.quat.x"] = imu.quaternion[1]
obs["imu.quat.y"] = imu.quaternion[2]
obs["imu.quat.z"] = imu.quaternion[3]
if imu.rpy:
obs["imu.rpy.roll"] = imu.rpy[0]
obs["imu.rpy.pitch"] = imu.rpy[1]
obs["imu.rpy.yaw"] = imu.rpy[2]
wr = getattr(lowstate, "wireless_remote", None)
if wr:
obs["wireless_remote"] = bytes(wr) if not isinstance(wr, (bytes, bytearray)) else wr
return obs
class UnitreeG1(Robot):
config_class = UnitreeG1Config
name = "unitree_g1"
@@ -188,60 +148,22 @@ class UnitreeG1(Robot):
self.arm_ik = G1_29_ArmIK() if config.gravity_compensation else None
# Lower-body / whole-body controller loaded dynamically
# Lower-body controller loaded dynamically
self.controller: LocomotionController | None = make_locomotion_controller(config.controller)
# A SONIC whole-body controller always runs in token mode: it holds a neutral
# token until the first real one arrives, then holds the last token between ticks.
if self.controller is not None and hasattr(self.controller, "token_mode"):
self.controller.token_mode = True
# Controller thread state
self._controller_thread = None
# When set, the controller loop stops publishing low commands so reset() can
# drive the joints directly without two publishers fighting (single-publisher).
self._controller_paused = threading.Event()
self._controller_action_lock = threading.Lock()
self.controller_input = default_remote_input()
self.controller_output = {}
# Token-mode state: last 64-D SONIC latent token commanded by the policy,
# echoed back as ``observation.state`` so a token-output VLA closes the loop
# on its own previous token. Implicit whenever the SONIC whole-body controller
# is active. Seeded to zeros; the controller's startup blend eases joints in.
self._last_token: np.ndarray | None = None
if self._sonic_token:
from .controllers.sonic_whole_body import TOKEN_DIM
self._last_token = np.zeros(TOKEN_DIM, dtype=np.float32)
@property
def _sonic_token(self) -> bool:
"""Whether the SONIC whole-body decoder is active.
A SONIC controller consumes a 64-D latent motion token as its action and echoes
the last commanded token as ``observation.state``. Keyed purely off the selected
controller so the token interface is implicit -- no separate config flag.
"""
return self.config.controller == "SonicWholeBodyController"
def _subscribe_lowstate(self): # polls robot state @ 250Hz
while not self._shutdown_event.is_set():
start_time = time.time()
# Step simulation if in simulation mode
if self.config.is_simulation and self.sim_env is not None:
try:
self.sim_env.step()
except ValueError as e:
# Startup race: the sim thread can step once before reset() has
# written a valid base pose, giving a zero-norm pelvis quaternion
# (scipy>=1.11 raises instead of normalizing). Skip and retry so
# the thread survives instead of dying and freezing the sim.
if "zero norm" not in str(e).lower():
raise
time.sleep(self.control_dt)
continue
msg = self.lowstate_subscriber.Read()
if msg is not None:
@@ -309,38 +231,15 @@ class UnitreeG1(Robot):
features[f"{cam}_depth"] = (cfg.height, cfg.width, 1)
return features
@property
def _token_state_ft(self) -> dict[str, type]:
"""64-D SONIC latent-token proprio state (``motion_token_state.{i}.pos``).
Exposed only when a SONIC whole-body controller is active; aggregated by the
rollout into a 64-D ``observation.state`` (the last token the policy commanded).
"""
if not self._sonic_token:
return {}
from .controllers.sonic_whole_body import TOKEN_DIM, token_state_key
return {token_state_key(i): float for i in range(TOKEN_DIM)}
@cached_property
def observation_features(self) -> dict[str, type | tuple]:
return {**self._motors_ft, **self._token_state_ft, **self._cameras_ft}
return {**self._motors_ft, **self._cameras_ft}
@cached_property
def action_features(self) -> dict[str, type]:
# No controller configured at all: raw 29-DoF joint teleop.
if self.controller is None:
return {f"{G1_29_JointIndex(motor).name}.q": float for motor in G1_29_JointIndex}
# Token-output VLA (SONIC decoder): advertise a 64-D latent-token action space
# (``motion_token.{i}.pos``) so ``lerobot-rollout`` maps a 64-D policy output
# straight onto the decoder, bypassing the encoder.
if self._sonic_token:
from .controllers.sonic_whole_body import TOKEN_DIM, token_action_key
return {token_action_key(i): float for i in range(TOKEN_DIM)}
# Locomotion controllers (GR00T / Holosoma): arm joint targets + joystick axes.
arm_features = {f"{G1_29_JointArmIndex(motor).name}.q": float for motor in G1_29_JointArmIndex}
remote_features = dict.fromkeys(REMOTE_AXES, float)
return {**arm_features, **remote_features}
@@ -356,11 +255,6 @@ class UnitreeG1(Robot):
while not self._shutdown_event.is_set():
start_time = time.time()
# Paused during reset() so the reset routine is the sole low-cmd publisher.
if self._controller_paused.is_set():
time.sleep(control_dt)
continue
with self._lowstate_lock:
lowstate = self._lowstate
@@ -449,9 +343,6 @@ class UnitreeG1(Robot):
self.kp = np.array(self.config.kp, dtype=np.float32)
self.kd = np.array(self.config.kd, dtype=np.float32)
if self.controller is not None and hasattr(self.controller, "kp"):
self.kp = np.array(self.controller.kp, dtype=np.float32)
self.kd = np.array(self.controller.kd, dtype=np.float32)
for joint in G1_29_JointIndex:
self.msg.motor_cmd[joint].mode = 1
@@ -500,10 +391,6 @@ class UnitreeG1(Robot):
if self._controller_thread.is_alive():
logger.warning("Controller thread did not stop cleanly")
# Release controller resources (e.g. SONIC decoder sessions).
if self.controller is not None and hasattr(self.controller, "shutdown"):
self.controller.shutdown()
# Close simulation environment
if self.config.is_simulation and self.sim_env is not None:
try:
@@ -574,15 +461,6 @@ class UnitreeG1(Robot):
if lowstate.wireless_remote:
obs["wireless_remote"] = lowstate.wireless_remote
# Token mode: echo the last commanded latent token as observation.state so a
# token-output VLA closes the loop on its own previous token.
if self._sonic_token:
from .controllers.sonic_whole_body import token_state_key
token = self._last_token if self._last_token is not None else []
for i, v in enumerate(token):
obs[token_state_key(i)] = float(v)
# Cameras - read images from ZMQ cameras
for cam_name, cam in self._cameras.items():
if getattr(cam, "use_rgb", True):
@@ -595,22 +473,9 @@ class UnitreeG1(Robot):
def send_action(self, action: RobotAction) -> RobotAction:
action_to_publish = action
if self.controller is not None:
# SONIC decoder: pull the 64-D latent token out of the action and remember it
# for the observation.state echo. The controller thread reads it back from
# controller_input (populated below) and decodes it into a 29-DoF command.
if self._sonic_token:
from .controllers.sonic_whole_body import _extract_token_from_action
token = _extract_token_from_action(action)
if token is not None:
self._last_token = token
self._update_controller_action(action)
# Full-body controllers (SONIC) own the whole 29-DoF command; nothing to
# publish here (the controller thread is the sole publisher).
if getattr(self.controller, "full_body", False):
return action
# Controller thread owns legs/waist. Here we only update joystick inputs
# and publish arm targets from the teleoperator.
self._update_controller_action(action)
arm_prefixes = tuple(j.name for j in G1_29_JointArmIndex)
action_to_publish = {
key: value
@@ -638,17 +503,11 @@ class UnitreeG1(Robot):
return action
def _update_controller_action(self, action: RobotAction) -> None:
"""Update controller input state from an incoming teleop action.
Controller-agnostic: every value-carrying key (locomotion ``remote.*`` axes or
SONIC ``motion_token.*`` values) is forwarded verbatim into ``controller_input``
and each controller extracts only the keys it understands. The robot deliberately
does not enumerate any controller's key schema here.
"""
"""Update controller input state from incoming teleop action."""
with self._controller_action_lock:
for key, value in action.items():
if isinstance(key, str) and value is not None:
self.controller_input[key] = value
for key in REMOTE_KEYS:
if key in action:
self.controller_input[key] = action[key]
@property
def is_calibrated(self) -> bool:
@@ -678,18 +537,6 @@ class UnitreeG1(Robot):
if default_positions is None:
default_positions = np.array(self.config.default_positions, dtype=np.float32)
# Full-body controllers (SONIC) own the whole 29-DoF command and ignore
# ``<joint>.q`` in send_action(), so reset() must publish the default pose
# directly. Pause the background controller first so the two aren't both writing
# low commands while the robot moves to the default pose.
full_body = getattr(self.controller, "full_body", False)
paused = False
if full_body and self._controller_thread is not None:
self._controller_paused.set()
paused = True
time.sleep(control_dt) # let any in-flight controller tick settle
try:
if self.config.is_simulation and self.sim_env is not None:
self.sim_env.reset()
self.publish_lowcmd(
@@ -718,11 +565,6 @@ class UnitreeG1(Robot):
interp_pos = init_dof_pos[motor.value] * (1 - alpha) + target_pos * alpha
action_dict[f"{motor.name}.q"] = float(interp_pos)
# Full-body controllers no-op in send_action(); publish the pose
# directly (arm-only controllers keep the send_action() path).
if full_body:
self.publish_lowcmd(action_dict)
else:
self.send_action(action_dict)
# Maintain constant control rate
@@ -730,12 +572,8 @@ class UnitreeG1(Robot):
sleep_time = max(0, control_dt - elapsed)
time.sleep(sleep_time)
# Reset controller internal state (gait phase, obs history, etc.) before
# resuming so its buffers reflect the post-reset pose.
# Reset controller internal state (gait phase, obs history, etc.)
if self.controller is not None and hasattr(self.controller, "reset"):
self.controller.reset()
finally:
if paused:
self._controller_paused.clear()
logger.info("Reached default position")
@@ -36,7 +36,6 @@ python src/lerobot/scripts/augment_dataset_quantile_stats.py \
import argparse
import concurrent.futures
import logging
import os
from pathlib import Path
import numpy as np
@@ -53,7 +52,6 @@ from lerobot.datasets import (
get_feature_stats,
write_stats,
)
from lerobot.datasets.compute_stats import sample_indices
from lerobot.utils.utils import init_logging
@@ -79,14 +77,12 @@ def has_quantile_stats(stats: dict[str, dict] | None, quantile_list_keys: list[s
return False
def process_single_episode(dataset: LeRobotDataset, episode_idx: int, use_sampling: bool = True) -> dict:
def process_single_episode(dataset: LeRobotDataset, episode_idx: int) -> dict:
"""Process a single episode and return its statistics.
Args:
dataset: The LeRobot dataset
episode_idx: Index of the episode to process
use_sampling: If True, sub-sample image/video frames per episode to bound
memory. If False, use every frame (exact, higher memory).
Returns:
Dictionary containing episode statistics
@@ -96,31 +92,16 @@ def process_single_episode(dataset: LeRobotDataset, episode_idx: int, use_sampli
start_idx = dataset.meta.episodes[episode_idx]["dataset_from_index"]
end_idx = dataset.meta.episodes[episode_idx]["dataset_to_index"]
episode_len = end_idx - start_idx
# Images/video are the memory hog, so sub-sample those frames per episode;
# numeric columns are cheap, so read them in full (exact).
image_keys = [k for k in dataset.features if dataset.features[k]["dtype"] in ("image", "video")]
numeric_keys = [
k for k in dataset.features if dataset.features[k]["dtype"] not in ("image", "video", "string")
]
collected_data: dict[str, list] = {}
for idx in range(start_idx, end_idx):
item = dataset[idx]
for key, value in item.items():
if key not in dataset.features:
continue
# Numeric features: every frame, read directly from the underlying table.
if numeric_keys:
numeric_cols = dataset.hf_dataset.select_columns(numeric_keys)[start_idx:end_idx]
for key in numeric_keys:
collected_data[key] = [torch.as_tensor(v) for v in numeric_cols[key]]
# Image/video features: decode only a sampled subset of frames.
if image_keys:
sampled_offsets = sample_indices(episode_len) if use_sampling else list(range(episode_len))
for offset in sampled_offsets:
item = dataset[start_idx + offset]
for key in image_keys:
if key in item:
collected_data.setdefault(key, []).append(item[key])
if key not in collected_data:
collected_data[key] = []
collected_data[key].append(value)
ep_stats = {}
for key, data_list in collected_data.items():
@@ -150,13 +131,11 @@ def process_single_episode(dataset: LeRobotDataset, episode_idx: int, use_sampli
return ep_stats
def compute_quantile_stats_for_dataset(dataset: LeRobotDataset, use_sampling: bool = True) -> dict[str, dict]:
def compute_quantile_stats_for_dataset(dataset: LeRobotDataset) -> dict[str, dict]:
"""Compute quantile statistics for all episodes in the dataset.
Args:
dataset: The LeRobot dataset to compute statistics for
use_sampling: If True, sub-sample image/video frames per episode to bound
memory. If False, use every frame (exact, higher memory).
Returns:
Dictionary containing aggregated statistics with quantiles
@@ -174,15 +153,15 @@ def compute_quantile_stats_for_dataset(dataset: LeRobotDataset, use_sampling: bo
if has_videos:
logging.info("Dataset contains video keys - using sequential processing for thread safety")
for episode_idx in tqdm(range(dataset.num_episodes), desc="Processing episodes"):
ep_stats = process_single_episode(dataset, episode_idx, use_sampling)
ep_stats = process_single_episode(dataset, episode_idx)
episode_stats_list.append(ep_stats)
else:
logging.info("Dataset has no video keys - using parallel processing for better performance")
max_workers = min(dataset.num_episodes, int(os.environ.get("LEROBOT_STATS_MAX_WORKERS", 16)))
max_workers = min(dataset.num_episodes, 16)
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
future_to_episode = {
executor.submit(process_single_episode, dataset, episode_idx, use_sampling): episode_idx
executor.submit(process_single_episode, dataset, episode_idx): episode_idx
for episode_idx in range(dataset.num_episodes)
}
@@ -209,7 +188,6 @@ def augment_dataset_with_quantile_stats(
repo_id: str,
root: str | Path | None = None,
overwrite: bool = False,
use_sampling: bool = True,
) -> None:
"""Augment a dataset with quantile statistics if they are missing.
@@ -217,8 +195,6 @@ def augment_dataset_with_quantile_stats(
repo_id: Repository ID of the dataset
root: Local root directory for the dataset
overwrite: Overwrite existing quantile statistics if they already exist
use_sampling: If True, sub-sample image/video frames per episode to bound
memory. If False, use every frame (exact, higher memory).
"""
logging.info(f"Loading dataset: {repo_id}")
dataset = LeRobotDataset(
@@ -232,7 +208,7 @@ def augment_dataset_with_quantile_stats(
logging.info("Dataset does not contain quantile statistics. Computing them now...")
new_stats = compute_quantile_stats_for_dataset(dataset, use_sampling=use_sampling)
new_stats = compute_quantile_stats_for_dataset(dataset)
logging.info("Updating dataset metadata with new quantile statistics")
dataset.meta.stats = new_stats
@@ -272,14 +248,6 @@ def main():
action="store_true",
help="Overwrite existing quantile statistics if they already exist",
)
parser.add_argument(
"--no-sampling",
action="store_true",
help=(
"Compute stats over every frame (exact, higher memory). By default, "
"image/video frames are sub-sampled per episode to bound memory."
),
)
args = parser.parse_args()
root = Path(args.root) if args.root else None
@@ -290,7 +258,6 @@ def main():
repo_id=args.repo_id,
root=root,
overwrite=args.overwrite,
use_sampling=not args.no_sampling,
)
+1 -1
View File
@@ -564,7 +564,7 @@ def eval_policy(
if seeds:
all_seeds.extend(seeds)
else:
all_seeds.extend([None] * env.num_envs)
all_seeds.append(None)
# FIXME: episode_data is either None or it doesn't exist
if return_episode_data:
+39 -28
View File
@@ -28,6 +28,7 @@ lerobot-find-cameras
# NOTE(Steven): macOS cameras sometimes report different FPS at init time, not an issue here as we don't specify FPS when opening the cameras, but the information displayed might not be truthful.
import argparse
import concurrent.futures
import logging
import time
from pathlib import Path
@@ -132,7 +133,7 @@ def save_image(
camera_identifier: str | int,
images_dir: Path,
camera_type: str,
) -> None:
):
"""
Saves a single image to disk using Pillow. Handles color conversion if necessary.
"""
@@ -151,7 +152,7 @@ def save_image(
logger.error(f"Failed to save image for camera {camera_identifier} (type {camera_type}): {e}")
def create_camera_instance(cam_meta: dict[str, Any], *, warmup_s: int = 1) -> dict[str, Any] | None:
def create_camera_instance(cam_meta: dict[str, Any]) -> dict[str, Any] | None:
"""Create and connect to a camera instance based on metadata."""
cam_type = cam_meta.get("type")
cam_id = cam_meta.get("id")
@@ -164,14 +165,12 @@ def create_camera_instance(cam_meta: dict[str, Any], *, warmup_s: int = 1) -> di
cv_config = OpenCVCameraConfig(
index_or_path=cam_id,
color_mode=ColorMode.RGB,
warmup_s=warmup_s,
)
instance = OpenCVCamera(cv_config)
elif cam_type == "RealSense":
rs_config = RealSenseCameraConfig(
serial_number_or_name=cam_id,
color_mode=ColorMode.RGB,
warmup_s=warmup_s,
)
instance = RealSenseCamera(rs_config)
else:
@@ -189,7 +188,9 @@ def create_camera_instance(cam_meta: dict[str, Any], *, warmup_s: int = 1) -> di
return None
def process_camera_image(cam_dict: dict[str, Any], output_dir: Path, current_time: float) -> None:
def process_camera_image(
cam_dict: dict[str, Any], output_dir: Path, current_time: float
) -> concurrent.futures.Future | None:
"""Capture and process an image from a single camera."""
cam = cam_dict["instance"]
meta = cam_dict["meta"]
@@ -199,7 +200,7 @@ def process_camera_image(cam_dict: dict[str, Any], output_dir: Path, current_tim
try:
image_data = cam.read()
save_image(
return save_image(
image_data,
cam_id_str,
output_dir,
@@ -214,9 +215,10 @@ def process_camera_image(cam_dict: dict[str, Any], output_dir: Path, current_tim
return None
def cleanup_camera(cam_dict: dict[str, Any]) -> None:
def cleanup_cameras(cameras_to_use: list[dict[str, Any]]):
"""Disconnect all cameras."""
logger.info(f"Disconnecting camera with ID {cam_dict['meta'].get('id')}...")
logger.info(f"Disconnecting {len(cameras_to_use)} cameras...")
for cam_dict in cameras_to_use:
try:
if cam_dict["instance"] and cam_dict["instance"].is_connected:
cam_dict["instance"].disconnect()
@@ -228,7 +230,6 @@ def save_images_from_all_cameras(
output_dir: Path,
record_time_s: float = 2.0,
camera_type: str | None = None,
warmup_s: int = 1,
):
"""
Connects to detected cameras (optionally filtered by type) and saves images from each.
@@ -239,7 +240,6 @@ def save_images_from_all_cameras(
record_time_s: Duration in seconds to record images.
camera_type: Optional string to filter cameras ("realsense" or "opencv").
If None, uses all detected cameras.
warmup_s: Duration in seconds to warmup camera before recording images.
"""
output_dir.mkdir(parents=True, exist_ok=True)
logger.info(f"Saving images to {output_dir}")
@@ -249,23 +249,39 @@ def save_images_from_all_cameras(
logger.warning("No cameras detected matching the criteria. Cannot save images.")
return
logger.info(
f"Starting image capture for {record_time_s} seconds from {len(all_camera_metadata)} cameras."
)
try:
cameras_to_use = []
for cam_meta in all_camera_metadata:
cam_dict = create_camera_instance(cam_meta, warmup_s=warmup_s)
if cam_dict is None:
continue
camera_instance = create_camera_instance(cam_meta)
if camera_instance:
cameras_to_use.append(camera_instance)
if not cameras_to_use:
logger.warning("No cameras could be connected. Aborting image save.")
return
logger.info(f"Starting image capture for {record_time_s} seconds from {len(cameras_to_use)} cameras.")
start_time = time.perf_counter()
with concurrent.futures.ThreadPoolExecutor(max_workers=len(cameras_to_use) * 2) as executor:
try:
while time.perf_counter() - start_time < record_time_s:
futures = []
current_capture_time = time.perf_counter()
process_camera_image(cam_dict, output_dir, current_capture_time)
cleanup_camera(cam_dict)
for cam_dict in cameras_to_use:
future = process_camera_image(cam_dict, output_dir, current_capture_time)
if future:
futures.append(future)
if futures:
concurrent.futures.wait(futures)
except KeyboardInterrupt:
logger.info("Capture interrupted by user.")
finally:
print("\nFinalizing image saving...")
executor.shutdown(wait=True)
cleanup_cameras(cameras_to_use)
print(f"Image capture finished. Images saved to {output_dir}")
@@ -275,6 +291,7 @@ def main():
parser = argparse.ArgumentParser(
description="Unified camera utility script for listing cameras and capturing images."
)
parser.add_argument(
"camera_type",
type=str,
@@ -292,14 +309,8 @@ def main():
parser.add_argument(
"--record-time-s",
type=float,
default=2.0,
help="Time duration to attempt capturing frames. Default: 2 seconds.",
)
parser.add_argument(
"--warmup-s",
type=int,
default=1,
help="Time duration to warmup camera before attempting to capture frames. Default: 1 second.",
default=6.0,
help="Time duration to attempt capturing frames. Default: 6 seconds.",
)
args = parser.parse_args()
save_images_from_all_cameras(**vars(args))
+13 -17
View File
@@ -22,8 +22,7 @@ import dataclasses
import logging
import sys
import time
from collections.abc import Iterator
from contextlib import contextmanager, nullcontext
from contextlib import nullcontext
from pprint import pformat
from typing import TYPE_CHECKING, Any
@@ -77,20 +76,6 @@ else:
from .lerobot_eval import eval_policy_all
@contextmanager
def _make_eval_envs(cfg: TrainPipelineConfig) -> Iterator[dict[str, dict[int, Any]]]:
"""Create evaluation environments for one run and always dispose of them."""
envs = make_env(
cfg.env,
n_envs=cfg.eval.batch_size,
use_async_envs=cfg.eval.use_async_envs,
)
try:
yield envs
finally:
close_envs(envs)
def _dataloader_worker_kwargs(cfg: TrainPipelineConfig) -> dict[str, Any]:
"""Return worker-only DataLoader options, disabling them for single-process loading."""
workers_enabled = cfg.num_workers > 0
@@ -295,6 +280,14 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if not is_main_process:
dataset, eval_dataset = make_train_eval_datasets(cfg)
# Create environment used for evaluating checkpoints during training on simulation data.
# On real-world data, no need to create an environment as evaluations are done outside train.py,
# using the eval.py instead, with gym_dora environment and dora-rs.
eval_env = None
if cfg.env_eval_freq > 0 and cfg.env is not None and is_main_process:
logging.info("Creating env")
eval_env = make_env(cfg.env, n_envs=cfg.eval.batch_size, use_async_envs=cfg.eval.use_async_envs)
if cfg.is_reward_model_training:
if is_main_process:
logging.info("Creating reward model")
@@ -702,7 +695,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if is_main_process:
step_id = get_step_identifier(step, cfg.steps)
logging.info(f"Eval policy at step {step}")
with _make_eval_envs(cfg) as eval_env, torch.no_grad(), accelerator.autocast():
with torch.no_grad(), accelerator.autocast():
eval_info = eval_policy_all(
envs=eval_env, # dict[suite][task_id] -> vec_env
policy=accelerator.unwrap_model(policy),
@@ -750,6 +743,9 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if is_main_process:
progbar.close()
if eval_env:
close_envs(eval_env)
is_fsdp = accelerator.distributed_type == DistributedType.FSDP
model_state_dict = accelerator.get_state_dict(policy) if is_fsdp else None
if is_main_process:
@@ -171,13 +171,7 @@ class IOSPhone(BasePhone, Teleoperator):
# HEBI provides orientation in w, x, y, z format.
# Scipy's Rotation expects x, y, z, w.
quat_xyzw = np.concatenate((ar_quat[1:], [ar_quat[0]])) # wxyz to xyzw
# ARKit can emit zero/NaN quaternions before tracking is ready or on a
# dropped packet. Rotation.from_quat now rejects those; degrade the same
# way as a missing pose so teleop stays alive mid-session.
try:
rot = Rotation.from_quat(quat_xyzw)
except ValueError:
return False, None, None, None
pos = ar_pos - rot.apply(self.config.camera_offset)
return True, pos, rot, pose
@@ -29,12 +29,6 @@ class SOLeaderConfig:
# Whether to use degrees for angles
use_degrees: bool = True
# Number of extra attempts when a `sync_read` of the motors fails. Feetech buses can occasionally
# return a corrupted status packet ("Incorrect status packet!"), especially when several joints move
# at once, which otherwise aborts the teleoperation loop. Retries are immediate (no sleep) and only
# happen on failure, so the steady-state read cost is unchanged.
num_read_retries: int = 2
@TeleoperatorConfig.register_subclass("so101_leader")
@TeleoperatorConfig.register_subclass("so100_leader")
@@ -145,7 +145,7 @@ class SOLeader(Teleoperator):
@check_if_not_connected
def get_action(self) -> dict[str, float]:
start = time.perf_counter()
action = self.bus.sync_read("Present_Position", num_retry=self.config.num_read_retries)
action = self.bus.sync_read("Present_Position")
action = {f"{motor}.pos": val for motor, val in action.items()}
dt_ms = (time.perf_counter() - start) * 1e3
logger.debug(f"{self} read action: {dt_ms:.1f}ms")
+7 -7
View File
@@ -41,7 +41,7 @@ class RandomSubsetApply(Transform):
def __init__(
self,
transforms: Sequence[Callable[..., Any]],
transforms: Sequence[Callable],
p: list[float] | None = None,
n_subset: int | None = None,
random_order: bool = False,
@@ -50,7 +50,7 @@ class RandomSubsetApply(Transform):
if not isinstance(transforms, Sequence):
raise TypeError("Argument transforms should be a sequence of callables")
if p is None:
p = [1.0] * len(transforms)
p = [1] * len(transforms)
elif len(p) != len(transforms):
raise ValueError(
f"Length of p doesn't match the number of transforms: {len(p)} != {len(transforms)}"
@@ -69,7 +69,7 @@ class RandomSubsetApply(Transform):
self.n_subset = n_subset
self.random_order = random_order
self.selected_transforms: list[Callable[..., Any]] = []
self.selected_transforms = None
def forward(self, *inputs: Any) -> Any:
needs_unpacking = len(inputs) > 1
@@ -119,7 +119,7 @@ class SharpnessJitter(Transform):
super().__init__()
self.sharpness = self._check_input(sharpness)
def _check_input(self, sharpness: float | Sequence[float]) -> tuple[float, float]:
def _check_input(self, sharpness):
if isinstance(sharpness, (int | float)):
if sharpness < 0:
raise ValueError("If sharpness is a single number, it must be non negative.")
@@ -215,7 +215,7 @@ class ImageTransformsConfig:
)
def make_transform_from_config(cfg: ImageTransformConfig) -> Transform:
def make_transform_from_config(cfg: ImageTransformConfig):
if cfg.type == "SharpnessJitter":
return SharpnessJitter(**cfg.kwargs)
@@ -236,8 +236,8 @@ class ImageTransforms(Transform):
super().__init__()
self._cfg = cfg
self.weights: list[float] = []
self.transforms: dict[str, Transform] = {}
self.weights = []
self.transforms = {}
for tf_name, tf_cfg in cfg.tfs.items():
if tf_cfg.weight <= 0.0:
continue
+4 -13
View File
@@ -37,25 +37,16 @@ def auto_select_torch_device() -> torch.device:
# TODO(Steven): Remove log. log shouldn't be an argument, this should be handled by the logger level
def get_safe_torch_device(try_device: str, log: bool = False) -> torch.device:
"""Given a string, return a torch.device with checks on whether the device is available.
Raises:
ValueError: If the requested device family is known but not available on
this machine (``AssertionError`` was previously used and is easy to
mistake for a programmer bug under ``python -O`` where asserts vanish).
"""
"""Given a string, return a torch.device with checks on whether the device is available."""
try_device = str(try_device)
if try_device.startswith("cuda"):
if not torch.cuda.is_available():
raise ValueError(f"Requested device {try_device!r} but CUDA is not available.")
assert torch.cuda.is_available()
device = torch.device(try_device)
elif try_device == "mps":
if not torch.backends.mps.is_available():
raise ValueError("Requested device 'mps' but MPS is not available.")
assert torch.backends.mps.is_available()
device = torch.device("mps")
elif try_device == "xpu":
if not torch.xpu.is_available():
raise ValueError("Requested device 'xpu' but XPU is not available.")
assert torch.xpu.is_available()
device = torch.device("xpu")
elif try_device == "cpu":
device = torch.device("cpu")
+5 -5
View File
@@ -32,21 +32,21 @@ def load_json(fpath: Path) -> Any:
Returns:
Any: The data loaded from the JSON file.
"""
with open(fpath, encoding="utf-8") as f:
with open(fpath) as f:
return json.load(f)
def write_json(data: JsonLike, fpath: Path) -> None:
"""Write JSON-serializable data to a file.
def write_json(data: dict, fpath: Path) -> None:
"""Write data to a JSON file.
Creates parent directories if they don't exist.
Args:
data: JSON-serializable data to write.
data (dict): The dictionary to write.
fpath (Path): The path to the output JSON file.
"""
fpath.parent.mkdir(exist_ok=True, parents=True)
with open(fpath, "w", encoding="utf-8") as f:
with open(fpath, "w") as f:
json.dump(data, f, indent=4, ensure_ascii=False)
-4
View File
@@ -30,10 +30,6 @@ def precise_sleep(seconds: float, spin_threshold: float = 0.010, sleep_margin: f
"""
if seconds <= 0:
return
if spin_threshold < 0:
raise ValueError(f"spin_threshold must be >= 0, got {spin_threshold}")
if sleep_margin < 0:
raise ValueError(f"sleep_margin must be >= 0, got {sleep_margin}")
system = platform.system()
# On macOS and Windows the scheduler / sleep granularity can make
+2 -5
View File
@@ -29,12 +29,9 @@ class Rotation:
def __init__(self, quat: np.ndarray) -> None:
"""Initialize rotation from quaternion [x, y, z, w]."""
self._quat = np.asarray(quat, dtype=float)
if self._quat.shape != (4,):
raise ValueError(f"Quaternion must have shape (4,), got {self._quat.shape}")
# Normalize quaternion. Reject the zero vector — it has no orientation.
# Normalize quaternion
norm = np.linalg.norm(self._quat)
if norm <= 0.0 or not np.isfinite(norm):
raise ValueError(f"Quaternion must be a non-zero finite vector; got {self._quat} (norm={norm})")
if norm > 0:
self._quat = self._quat / norm
@classmethod
+2 -2
View File
@@ -14,7 +14,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import NotRequired, TypedDict
from typing import TypedDict
import torch
@@ -28,7 +28,7 @@ class Transition(TypedDict):
next_state: dict[str, torch.Tensor]
done: bool
truncated: bool
complementary_info: NotRequired[dict[str, torch.Tensor | float | int] | None]
complementary_info: dict[str, torch.Tensor | float | int] | None = None
def move_transition_to_device(transition: Transition, device: str = "cpu") -> Transition:
+9 -13
View File
@@ -24,6 +24,7 @@ import sys
import time
from collections.abc import Iterator
from copy import copy, deepcopy
from datetime import datetime
from pathlib import Path
from statistics import mean
from typing import TYPE_CHECKING, Any
@@ -60,16 +61,14 @@ def init_logging(
accelerator: Optional Accelerator instance (for multi-GPU detection)
"""
class LeRobotFormatter(logging.Formatter):
def format(self, record: logging.LogRecord) -> str:
record.lerobot_location = f"{record.pathname}:{record.lineno}"[-15:]
record.lerobot_pid = f"[PID: {os.getpid()}] " if display_pid else ""
return super().format(record)
def custom_format(record: logging.LogRecord) -> str:
dt = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
fnameline = f"{record.pathname}:{record.lineno}"
pid_str = f"[PID: {os.getpid()}] " if display_pid else ""
return f"{record.levelname} {pid_str}{dt} {fnameline[-15:]:>15} {record.getMessage()}"
formatter = LeRobotFormatter(
"%(levelname)s %(lerobot_pid)s%(asctime)s %(lerobot_location)15s %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
)
formatter = logging.Formatter()
formatter.format = custom_format
logger = logging.getLogger()
logger.setLevel(logging.NOTSET)
@@ -134,13 +133,10 @@ def say(text: str, blocking: bool = False):
else:
raise RuntimeError("Unsupported operating system for text-to-speech.")
try:
if blocking:
subprocess.run(cmd, check=True, timeout=5)
subprocess.run(cmd, check=True)
else:
subprocess.Popen(cmd, creationflags=subprocess.CREATE_NO_WINDOW if system == "Windows" else 0)
except (FileNotFoundError, subprocess.TimeoutExpired) as e:
logging.warning("Text-to-speech command failed: %s | Error: %s", cmd, e)
def log_say(text: str, play_sounds: bool = True, blocking: bool = False):
@@ -1,104 +0,0 @@
# 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 numpy as np
import pytest
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
from lerobot.scripts.augment_dataset_quantile_stats import (
compute_quantile_stats_for_dataset,
has_quantile_stats,
)
def _numeric_keys(dataset):
return [k for k, v in dataset.features.items() if v["dtype"] not in ("image", "video", "string")]
def _image_keys(dataset):
return [k for k, v in dataset.features.items() if v["dtype"] in ("image", "video")]
def test_numeric_stats_are_unaffected_by_sampling(tmp_path, lerobot_dataset_factory):
"""Sampling only touches image/video frames; numeric features are read in
full either way, so their stats must be identical with and without sampling."""
dataset = lerobot_dataset_factory(
root=tmp_path / "ds", total_episodes=2, total_frames=400, use_videos=False
)
exact = compute_quantile_stats_for_dataset(dataset, use_sampling=False)
sampled = compute_quantile_stats_for_dataset(dataset, use_sampling=True)
numeric_keys = _numeric_keys(dataset)
assert numeric_keys, "fixture should expose numeric features"
for key in numeric_keys:
if key not in exact:
continue
for stat in ("mean", "std", "q01", "q50", "q99"):
if stat in exact[key]:
np.testing.assert_allclose(
sampled[key][stat],
exact[key][stat],
rtol=1e-6,
atol=1e-6,
err_msg=f"numeric feature '{key}' stat '{stat}' changed under sampling",
)
def test_image_sampling_reduces_data_but_keeps_stats_close(tmp_path, lerobot_dataset_factory):
"""For images, sampling should reduce the number of samples considered while
keeping the resulting statistics close to the exact ones."""
dataset = lerobot_dataset_factory(
root=tmp_path / "ds", total_episodes=2, total_frames=400, use_videos=False
)
exact = compute_quantile_stats_for_dataset(dataset, use_sampling=False)
sampled = compute_quantile_stats_for_dataset(dataset, use_sampling=True)
image_keys = _image_keys(dataset)
assert image_keys, "fixture should expose at least one image feature"
for key in image_keys:
# sampling actually looked at fewer pixels
assert sampled[key]["count"][0] < exact[key]["count"][0]
# but per-channel mean stays close
np.testing.assert_allclose(
sampled[key]["mean"],
exact[key]["mean"],
rtol=0.15,
err_msg=f"image feature '{key}' mean drifted too far under sampling",
)
def test_short_episodes_use_all_frames(tmp_path, lerobot_dataset_factory):
"""With episodes shorter than the sampling floor, sampling is a no-op and
must produce exactly the same stats as the exact path."""
dataset = lerobot_dataset_factory(
root=tmp_path / "ds", total_episodes=2, total_frames=40, use_videos=False
)
exact = compute_quantile_stats_for_dataset(dataset, use_sampling=False)
sampled = compute_quantile_stats_for_dataset(dataset, use_sampling=True)
for key in _image_keys(dataset):
assert sampled[key]["count"][0] == exact[key]["count"][0]
def test_quantile_stats_present_after_compute(tmp_path, lerobot_dataset_factory):
"""The computed stats should contain quantile keys for the dataset."""
dataset = lerobot_dataset_factory(
root=tmp_path / "ds", total_episodes=2, total_frames=200, use_videos=False
)
stats = compute_quantile_stats_for_dataset(dataset, use_sampling=True)
assert has_quantile_stats(stats)
-14
View File
@@ -482,20 +482,6 @@ def test_add_frame_works_in_write_mode(tmp_path):
# ── Resume mode ──────────────────────────────────────────────────────
def test_resume_freshly_created_empty_dataset(tmp_path):
"""resume() accepts a local dataset created before any episode was recorded."""
root = tmp_path / "resume_empty_ds"
LeRobotDataset.create(repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root)
resumed = LeRobotDataset.resume(repo_id=DUMMY_REPO_ID, root=root)
assert isinstance(resumed.writer, DatasetWriter)
assert resumed.meta.total_episodes == 0
assert resumed.meta.total_frames == 0
assert resumed.meta.tasks is None
assert resumed.meta.episodes is None
def test_resume_creates_writer(tmp_path):
"""After resume(), writer is a DatasetWriter."""
root = tmp_path / "resume_ds"
-13
View File
@@ -294,19 +294,6 @@ def test__sync_read(addr, length, ids_values, mock_motors, dummy_motors):
assert read_values == ids_values
def test__sync_read_retries_after_transient_failure(mock_motors, dummy_motors):
addr, length, ids_values = (10, 4, {1: 1337})
stub = mock_motors.build_sync_read_stub(addr, length, ids_values, num_invalid_try=1)
bus = FeetechMotorsBus(port=mock_motors.port, motors=dummy_motors)
bus.connect(handshake=False)
read_values, read_comm = bus._sync_read(addr, length, list(ids_values), num_retry=1)
assert read_comm == scs.COMM_SUCCESS
assert read_values == ids_values
assert mock_motors.stubs[stub].calls == 2
@pytest.mark.parametrize("raise_on_error", (True, False))
def test__sync_read_comm(raise_on_error, mock_motors, dummy_motors):
addr, length, ids_values = (10, 4, {1: 1337})
-54
View File
@@ -496,60 +496,6 @@ def test_evo1_processor_save_load_round_trip_applies_config_overrides(tmp_path):
assert "embodiment_id" in processed
def test_reconcile_evo1_processors_repads_overridden_stats(tmp_path):
"""Loading a checkpoint and injecting raw (unpadded) dataset stats must be re-padded.
Regression test: lerobot-train passes the raw dataset stats as normalizer/unnormalizer
overrides when resuming from a checkpoint (e.g. stage2 from a stage1 checkpoint). Those stats
are at the dataset dims (e.g. LIBERO state=8/action=7), but EVO1 pads state/action to
max_state_dim/max_action_dim before normalization, so reconcile_evo1_processors must re-pad the
stats or normalization crashes with a shape mismatch.
"""
config = make_config()
preprocessor, postprocessor = make_evo1_pre_post_processors(config, dataset_stats=make_stats())
preprocessor.save_pretrained(tmp_path)
postprocessor.save_pretrained(tmp_path)
# Reload with the generic override path injecting raw, unpadded dataset stats.
raw_stats = make_stats()
loaded_pre = PolicyProcessorPipeline.from_pretrained(
tmp_path,
config_filename=f"{POLICY_PREPROCESSOR_DEFAULT_NAME}.json",
overrides={"normalizer_processor": {"stats": raw_stats}},
to_transition=batch_to_transition,
to_output=transition_to_batch,
)
loaded_post = PolicyProcessorPipeline.from_pretrained(
tmp_path,
config_filename=f"{POLICY_POSTPROCESSOR_DEFAULT_NAME}.json",
overrides={"unnormalizer_processor": {"stats": raw_stats}},
to_transition=policy_action_to_transition,
to_output=transition_to_policy_action,
)
# Sanity: the override really injected unpadded stats before reconciliation.
normalizer = next(step for step in loaded_pre.steps if isinstance(step, NormalizerProcessorStep))
assert normalizer._tensor_stats[OBS_STATE]["min"].shape == (STATE_DIM,)
loaded_pre, loaded_post = reconcile_evo1_processors(config, loaded_pre, loaded_post)
normalizer = next(step for step in loaded_pre.steps if isinstance(step, NormalizerProcessorStep))
unnormalizer = next(step for step in loaded_post.steps if isinstance(step, UnnormalizerProcessorStep))
assert normalizer._tensor_stats[OBS_STATE]["min"].shape == (MAX_STATE_DIM,)
assert normalizer._tensor_stats[ACTION]["min"].shape == (MAX_ACTION_DIM,)
assert unnormalizer._tensor_stats[ACTION]["min"].shape == (MAX_ACTION_DIM,)
# Normalizing a padded state must not raise (this is the exact runtime path that crashed).
processed = loaded_pre(
{
"task": "pick the block",
OBS_STATE: torch.zeros(STATE_DIM),
f"{OBS_IMAGES}.front": torch.rand(3, 16, 16),
}
)
assert processed[OBS_STATE].shape == (1, MAX_STATE_DIM)
def test_evo1_policy_forward_and_inference_use_batched_embedding(monkeypatch):
monkeypatch.setattr(modeling_evo1, "Evo1Model", DummyEvo1Model)
policy = modeling_evo1.Evo1Policy(make_config())
+8 -12
View File
@@ -12,6 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import sys
from types import SimpleNamespace
from unittest.mock import MagicMock
@@ -46,23 +47,18 @@ def test_make_policy_keeps_peft_adapter_and_base_revisions_separate(monkeypatch)
peft_config_from_pretrained = MagicMock(return_value=peft_config)
adapted_policy = torch.nn.Linear(1, 1)
peft_model_from_pretrained = MagicMock(return_value=adapted_policy)
require_package = MagicMock()
monkeypatch.setattr(policy_factory, "require_package", require_package)
monkeypatch.setattr(
policy_factory,
"PeftConfig",
SimpleNamespace(from_pretrained=peft_config_from_pretrained),
)
monkeypatch.setattr(
policy_factory,
"PeftModel",
SimpleNamespace(from_pretrained=peft_model_from_pretrained),
monkeypatch.setitem(
sys.modules,
"peft",
SimpleNamespace(
PeftConfig=SimpleNamespace(from_pretrained=peft_config_from_pretrained),
PeftModel=SimpleNamespace(from_pretrained=peft_model_from_pretrained),
),
)
policy = policy_factory.make_policy(cfg, ds_meta=dataset_meta)
assert policy is adapted_policy
require_package.assert_called_once_with("peft", extra="peft")
peft_config_from_pretrained.assert_called_once_with(
"user/adapter",
revision="adapter-sha",
@@ -1,45 +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 pytest
from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
from lerobot.robots.so_follower.robot_kinematic_processor import (
ForwardKinematicsJointsToEEAction,
ForwardKinematicsJointsToEEObservation,
)
MOTOR_NAMES = ["shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll", "gripper"]
EE_KEYS = {f"ee.{k}" for k in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]}
def _joint_bucket(feature_type: FeatureType) -> dict[str, PolicyFeature]:
return {f"{n}.pos": PolicyFeature(type=feature_type, shape=(1,)) for n in MOTOR_NAMES}
@pytest.mark.parametrize(
("step_cls", "bucket", "feature_type"),
[
(ForwardKinematicsJointsToEEAction, PipelineFeatureType.ACTION, FeatureType.ACTION),
(ForwardKinematicsJointsToEEObservation, PipelineFeatureType.OBSERVATION, FeatureType.STATE),
],
)
def test_fk_feature_schema(step_cls, bucket, feature_type):
features = {PipelineFeatureType.ACTION: {}, PipelineFeatureType.OBSERVATION: {}}
features[bucket] = _joint_bucket(feature_type)
out = step_cls(kinematics=None, motor_names=MOTOR_NAMES).transform_features(features)[bucket]
assert set(out) == EE_KEYS
assert {feature.type for feature in out.values()} == {feature_type}
+2 -23
View File
@@ -49,7 +49,7 @@ def _make_bus_mock() -> MagicMock:
@pytest.fixture
def follower(tmp_path):
def follower():
bus_mock = _make_bus_mock()
def _bus_side_effect(*_args, **kwargs):
@@ -71,7 +71,7 @@ def follower(tmp_path):
),
patch.object(SO100Follower, "configure", lambda self: None),
):
cfg = SO100FollowerConfig(port="/dev/null", calibration_dir=tmp_path)
cfg = SO100FollowerConfig(port="/dev/null")
robot = SO100Follower(cfg)
yield robot
if robot.is_connected:
@@ -99,27 +99,6 @@ def test_get_observation(follower):
assert obs[f"{motor}.pos"] == idx
def test_get_observation_uses_read_retries(follower):
# Feetech buses can intermittently fail a sync_read; the follower should forward the configured
# retry count so transient failures don't abort the control loop (see #3131).
follower.config.num_read_retries = 7
follower.connect()
follower.get_observation()
follower.bus.sync_read.assert_called_once_with("Present_Position", num_retry=7)
def test_send_action_uses_read_retries(follower):
follower.config.max_relative_target = 10.0
follower.config.num_read_retries = 7
follower.connect()
action = {f"{motor}.pos": value * 10 for value, motor in enumerate(follower.bus.motors, 1)}
follower.send_action(action)
follower.bus.sync_read.assert_called_once_with("Present_Position", num_retry=7)
def test_send_action(follower):
follower.connect()
-36
View File
@@ -1,36 +0,0 @@
#!/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 unittest.mock import patch
import pytest
import torch
from lerobot.utils.device_utils import get_safe_torch_device, is_torch_device_available
def test_cpu_always_available():
assert get_safe_torch_device("cpu") == torch.device("cpu")
assert is_torch_device_available("cpu")
def test_missing_cuda_raises_valueerror():
with patch("torch.cuda.is_available", return_value=False), pytest.raises(ValueError, match="CUDA"):
get_safe_torch_device("cuda")
def test_missing_mps_raises_valueerror():
with patch("torch.backends.mps.is_available", return_value=False), pytest.raises(ValueError, match="MPS"):
get_safe_torch_device("mps")
-46
View File
@@ -1,46 +0,0 @@
#!/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.
import numpy as np
import pytest
from lerobot.utils.rotation import Rotation
def test_zero_quaternion_rejected():
with pytest.raises(ValueError, match="non-zero"):
Rotation(np.zeros(4))
def test_non_finite_quaternion_rejected():
with pytest.raises(ValueError, match="non-zero|finite"):
Rotation(np.array([np.nan, 0.0, 0.0, 1.0]))
def test_wrong_shape_rejected():
with pytest.raises(ValueError, match="shape"):
Rotation(np.array([1.0, 0.0, 0.0]))
def test_identity_roundtrip():
r = Rotation.from_rotvec(np.zeros(3))
assert np.allclose(r.as_rotvec(), 0.0)
assert np.allclose(r.as_matrix(), np.eye(3))
def test_rotvec_roundtrip():
rotvec = np.array([0.1, -0.2, 0.3])
r = Rotation.from_rotvec(rotvec)
assert np.allclose(r.as_rotvec(), rotvec, atol=1e-6)
-230
View File
@@ -1,230 +0,0 @@
#!/usr/bin/env python3
"""Provision the SONIC decoder checkpoint at ``lerobot/sonic_decoder``.
Takes NVIDIA's ``nvidia/GEAR-SONIC/model_decoder.onnx``, embeds the SONIC deploy constants
(``kp``/``kd`` PD gains, ``default_angles`` standing pose, the residual ``action_scale``, and
the ``neutral_token`` idle latent) into the ONNX ``metadata_props`` (the convention Holosoma
uses for its gains), and pushes the result to ``lerobot/sonic_decoder``. After this runs, the
runtime loads the decoder *and* every one of these constants straight from the checkpoint --
no motor-physics math at deploy time, so ``sonic_whole_body.py`` carries none of the
armature/bandwidth machinery nor any hardcoded deploy constants.
The constants here are derived once from Unitree motor physics (armature + target bandwidth).
That derivation is intentionally kept in this one-off provisioning script (not the runtime);
the shared/harmonic helper is a separate PR.
Build only (no network/auth needed if the source ONNX is already cached):
python upload_sonic_decoder.py --out ./sonic_decoder
Build + upload:
huggingface-cli login # or export HF_TOKEN=...
python upload_sonic_decoder.py --upload
"""
from __future__ import annotations
import argparse
import json
import pathlib
import numpy as np
import onnx
from huggingface_hub import hf_hub_download
SRC_REPO_ID = "nvidia/GEAR-SONIC"
SRC_FILENAME = "model_decoder.onnx"
DST_REPO_ID = "lerobot/sonic_decoder"
# ── SONIC deploy-constant derivation (provisioning-time only) ─────────────────
# All constants are (29,) in IsaacLab joint order: legs, waist, arms.
# kp = armature * w**2, kd = 4 * armature * w, with a x2 factor on the stiff joints
# (ankles + waist). action_scale = 0.25 * effort / (armature * w**2) is the residual
# scaling that maps decoder output to a joint-angle delta on top of default_angles.
NATURAL_FREQ = 10.0 * 2.0 * np.pi
MOTOR_ARMATURE = {"5020": 0.003609725, "7520_14": 0.010177520, "7520_22": 0.025101925, "4010": 0.00425}
EFFORT = {"5020": 25.0, "7520_14": 88.0, "7520_22": 139.0, "4010": 5.0}
MOTOR_MODELS = (
["7520_22", "7520_22", "7520_14", "7520_22", "5020", "5020"] * 2
+ ["7520_14", "5020", "5020"]
+ ["5020", "5020", "5020", "5020", "5020", "4010", "4010"] * 2
)
DOUBLE_INDICES = {4, 5, 10, 11, 13, 14} # ankles + waist
# Nominal standing pose (rad), 29 joints in IsaacLab order. Decoder actions are residuals
# added on top of this.
DEFAULT_ANGLES = [
-0.312,
0.0,
0.0,
0.669,
-0.363,
0.0, # left leg
-0.312,
0.0,
0.0,
0.669,
-0.363,
0.0, # right leg
0.0,
0.0,
0.0, # waist
0.2,
0.2,
0.0,
0.6,
0.0,
0.0,
0.0, # left arm
0.2,
-0.2,
0.0,
0.6,
0.0,
0.0,
0.0, # right arm
]
# Neutral idle token (64-D), held until the first real token arrives. Captured from the
# encoder while the robot stood idle in sim: the encoder is an FSQ bottleneck (~5 bit/dim,
# Div(16)), so tokens live on the 1/16 grid. We store the integer FSQ codes and rescale by
# 1/16 -> an exact on-grid token that decodes to a stable, natural standing pose (unlike the
# literal all-zero token, which is off-manifold and decodes to a slightly goofy stance).
NEUTRAL_TOKEN_CODES = [
-1,
3,
1,
-1,
1,
-3,
6,
1,
1,
1,
-2,
-4,
-2,
0,
-3,
-1,
2,
-1,
-3,
-5,
3,
1,
1,
-4,
-1,
-1,
1,
-7,
0,
1,
2,
-2,
5,
-2,
-2,
-4,
0,
-1,
3,
-1,
0,
-5,
-1,
0,
-4,
0,
0,
-1,
-1,
2,
-2,
1,
3,
3,
1,
0,
0,
6,
0,
-7,
3,
0,
2,
-2,
]
def compute_kp_kd() -> tuple[list[float], list[float]]:
"""Return (kp, kd) as plain float lists, (29,) in IsaacLab joint order."""
def stiffness(k):
return MOTOR_ARMATURE[k] * NATURAL_FREQ**2
def damping(k):
return 4.0 * MOTOR_ARMATURE[k] * NATURAL_FREQ
kp = [(2 if i in DOUBLE_INDICES else 1) * stiffness(k) for i, k in enumerate(MOTOR_MODELS)]
kd = [(2 if i in DOUBLE_INDICES else 1) * damping(k) for i, k in enumerate(MOTOR_MODELS)]
return kp, kd
def compute_action_scale() -> list[float]:
"""Return the per-joint residual action scale, (29,) in IsaacLab joint order."""
return [0.25 * EFFORT[k] / (MOTOR_ARMATURE[k] * NATURAL_FREQ**2) for k in MOTOR_MODELS]
def build(out_dir: pathlib.Path) -> pathlib.Path:
"""Download the source decoder, embed the deploy-constant metadata, save to ``out_dir``."""
src = hf_hub_download(repo_id=SRC_REPO_ID, filename=SRC_FILENAME)
model = onnx.load(src)
kp, kd = compute_kp_kd()
neutral_token = [c / 16.0 for c in NEUTRAL_TOKEN_CODES] # FSQ Div(16): codes -> on-grid token
meta = {prop.key: prop.value for prop in model.metadata_props}
meta["kp"] = json.dumps(kp)
meta["kd"] = json.dumps(kd)
meta["action_scale"] = json.dumps(compute_action_scale())
meta["default_angles"] = json.dumps(DEFAULT_ANGLES)
meta["neutral_token"] = json.dumps(neutral_token)
# Rewrite metadata_props with the merged dict.
del model.metadata_props[:]
for key, value in meta.items():
model.metadata_props.add(key=key, value=value)
out_dir.mkdir(parents=True, exist_ok=True)
out_path = out_dir / SRC_FILENAME
onnx.save(model, out_path)
print(f"Wrote {out_path} with kp/kd/action_scale/default_angles/neutral_token metadata.")
return out_path
def upload(out_path: pathlib.Path) -> None:
from huggingface_hub import HfApi
api = HfApi()
api.create_repo(repo_id=DST_REPO_ID, repo_type="model", exist_ok=True)
api.upload_file(
path_or_fileobj=str(out_path),
path_in_repo=SRC_FILENAME,
repo_id=DST_REPO_ID,
repo_type="model",
)
print(f"Uploaded {out_path.name} -> {DST_REPO_ID}")
def main() -> None:
p = argparse.ArgumentParser()
p.add_argument("--out", type=pathlib.Path, default=pathlib.Path("./sonic_decoder"))
p.add_argument("--upload", action="store_true", help="Push the built ONNX to the hub")
args = p.parse_args()
out_path = build(args.out)
if args.upload:
upload(out_path)
if __name__ == "__main__":
main()