Compare commits

..

16 Commits

Author SHA1 Message Date
Steven Palma 0d788abd85 chore(robots): drive mode calibration gripper 2026-07-29 17:49:41 +02:00
johnnynunez 7b78e751a6 fix(calibration): detect inverted gripper drive_mode on SO follower/leader calibration
The SO follower/leader calibration hardcoded drive_mode=0 for all motors.
When the gripper motor is mounted mirrored relative to the other arm's
(raw position increases when closing instead of decreasing), the
0=closed/100=open normalization convention flips: a 'closed' (0%)
command from the leader unnormalizes to the follower's fully-open hard
stop, slamming the gripper open and triggering the Feetech overload
protection (torque off until power cycle).

FeetechMotorsBus already supports per-motor inversion via drive_mode
(apply_drive_mode=True); calibration just never set it. This adds a
step after range recording that asks the user to fully close the
gripper, reads the raw position, and sets drive_mode=1 when the closed
position sits at range_max. No calibration file format change.

Verified against the exact calibration values reported in #3942 using
the real FeetechMotorsBus normalization: with drive_mode=0 a 0% goal
maps to raw 2035 (open, bug reproduced); with drive_mode=1 it maps to
raw 3528 (closed, expected).

Fixes #3942
2026-07-29 17:34:54 +02:00
Martino Russi 6e5f6df6e7 fix(evo1): re-pad normalizer stats when loading from checkpoint (#3945)
* fix(evo1): re-pad normalizer stats when loading from checkpoint

reconcile_evo1_processors did not re-pad the (un)normalizer stats to
max_state_dim/max_action_dim on the checkpoint-load path. When
lerobot-train loads a checkpoint (e.g. stage2 from a stage1 checkpoint)
it injects the raw dataset stats via processor overrides, so LIBERO's
8-dim state stats normalized a 24-dim padded state and crashed with
"size of tensor a (24) must match tensor b (8)".

Restore _refresh_evo1_normalization_steps (removed in the "remove legacy
codepaths" refactor) and call it from reconcile_evo1_processors so the
loaded stats/features are re-padded to EVO1's fixed widths. Padding is a
no-op when stats are already at the target width.

Co-authored-by: Cursor <cursoragent@cursor.com>

* test(evo1): cover reconcile re-padding of overridden normalizer stats

Regression test for the stage2-from-checkpoint crash: reloading a
checkpoint with raw (unpadded) dataset stats injected via processor
overrides must be re-padded to max_state_dim/max_action_dim by
reconcile_evo1_processors, otherwise normalizing the padded state
raises a shape mismatch.

Co-authored-by: Cursor <cursoragent@cursor.com>

---------

Co-authored-by: Martino Russi <martino@huggingface.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-29 17:26:39 +02:00
Steven Palma 265abe6c79 chore(datasets): add typing to aggregate helpers (#4211)
* chore(datasets): add typing to aggregate helpers

Signed-off-by: nathon-lee <leejianwoo@gmail.com>

* chore(dataset): add more typing aggregate

* chore(test): remove panda test

---------

Signed-off-by: nathon-lee <leejianwoo@gmail.com>
Co-authored-by: nathon-lee <leejianwoo@gmail.com>
2026-07-29 17:07:34 +02:00
Old-Ding b4e2d0b610 docs: fix wording in guides (#3939)
Generated-by: OpenAI Codex

Signed-off-by: aineoae86-sys <ai.neo.ae86@gmail.com>
Co-authored-by: aineoae86-sys <ai.neo.ae86@gmail.com>
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-29 16:24:03 +02:00
Old-Ding 5594eba06a docs: fix repeated word in backward compatibility guide (#3938)
Generated-by: OpenAI Codex

Signed-off-by: aineoae86-sys <ai.neo.ae86@gmail.com>
Co-authored-by: aineoae86-sys <ai.neo.ae86@gmail.com>
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-29 16:23:01 +02:00
saime428 207183c2f8 docs: fix dataset split fraction example (#3936)
* docs: fix dataset split fraction example

* docs: preserve three-way dataset split example

---------

Co-authored-by: saime <2286263079@qq.com>
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-29 16:21:12 +02:00
Steven Palma 7d615acf9a fix(robots): retry SO follower/leader bus reads on transient Feetech errors (#4207)
* fix(robots): retry SO follower/leader bus reads on transient Feetech errors

SO-100/SO-101 teleoperation aborts when a sync_read of Present_Position
returns a corrupted status packet ("Incorrect status packet!"), which the
Feetech bus emits intermittently under load. The read path already supports
a num_retry argument but the SO follower and leader never used it, so a single
transient failure crashed the control loop.

Add a max_read_retry config option (default 3) to SOFollowerConfig and
SOLeaderConfig and forward it to every Present_Position sync_read. Retries are
immediate and only happen on failure, so the steady-state read cost is
unchanged; set max_read_retry=0 to restore the previous behavior.

Fixes #3131

* chore(robots): change defaults

---------

Co-authored-by: isaka1022 <isaka1022@gmail.com>
2026-07-29 15:20:14 +02:00
Steven Palma 09572babee perf(docker): split dependency install from source copy for CI layer caching (#4208)
* perf(docker): split dep install from src copy for CI layer caching

Install third-party deps (torch + all extras) in a layer keyed only on
pyproject.toml + uv.lock via --no-install-project, then copy src and install
the local package. Editing src/ no longer busts the heavy dependency layer,
so BuildKit layer cache hits across CI builds.

Applied to both Dockerfile.user and Dockerfile.internal.

* chore(ci): less verbose comments + copy all files

---------

Co-authored-by: dongmao.zhang <dongmao.zhang@bytedance.com>
2026-07-29 15:04:46 +02:00
Predrag Cvetkovic 35339d31e5 fix(datasets): bound memory of augment_dataset_quantile_stats by sampling frames (#3749)
* fix(datasets): bound memory of augment_dataset_quantile_stats by sampling frames

Per-episode stats previously materialized every frame (and decoded up to 16
episodes in parallel), so peak memory scaled with episode length and OOM'd on
large datasets (#2889). Numeric features are now read in full from the table
(exact), while only image/video frames are sub-sampled per episode using the
existing sample_indices heuristic. Worker count is configurable via
LEROBOT_STATS_MAX_WORKERS; --no-sampling restores exact behavior.

* Update tests/datasets/test_augment_quantile_stats.py

Co-authored-by: Haoming Song <1847575517@qq.com>
Signed-off-by: Pepijn <138571049+pkooij@users.noreply.github.com>

---------

Signed-off-by: Pepijn <138571049+pkooij@users.noreply.github.com>
Co-authored-by: Pepijn <138571049+pkooij@users.noreply.github.com>
Co-authored-by: Haoming Song <1847575517@qq.com>
2026-07-29 12:32:03 +02:00
Steven Palma f37be3edbe fix(eval): prevent eval_policy crash when start_seed is None and num_envs>1 (#4203)
* fix(eval): align seed list length with num_envs when unseeded

eval_policy appended a single None per batch to all_seeds on the unseeded path while the reward and success lists grew by num_envs. The per-episode zip(..., strict=True) then raised ValueError for num_envs > 1. Extend all_seeds by num_envs so the lists stay aligned.

* chore(tests): delete lerobot_eval test

---------

Co-authored-by: Devin Lai <markauto75@gmail.com>
2026-07-28 18:41:28 +02:00
Khalil Meftah 4d076845ac fix peft factory test mocking (#4201) 2026-07-28 17:54:58 +02:00
Steven Palma 413972c812 fix(env): eval env lifecycle (#4194)
Co-authored-by: itxaiohanglover <1531137510@qq.com>
Co-authored-by: nickndeng <nickndeng@gmail.com>
Co-authored-by: nickndeng <107904079+nickndeng@users.noreply.github.com>
Co-authored-by: Pepijn <138571049+pkooij@users.noreply.github.com>
2026-07-28 16:45:48 +02:00
Steven Palma 0449aa02f6 fix(utils): handle missing/unresponsive TTS on Linux (#4199)
* fix: handle missing/unresponsive TTS on Linux

spd-say may be installed but hang indefinitely when speech-dispatcher
is not running. Add a 5s timeout and catch TimeoutExpired alongside
FileNotFoundError so recording continues without audio.

* chore(utils): add log warning for say

---------

Co-authored-by: Jiwen Cai <jiwenc@nvidia.com>
2026-07-28 16:45:32 +02:00
Alexandre Edmond a05c0833e1 chore(mypy): cover annotations and transforms (#3860)
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-28 16:25:45 +02:00
Alexandre Edmond 7b76d94c5b Handle resuming empty local datasets (#3859)
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
2026-07-28 16:25:42 +02:00
30 changed files with 599 additions and 311 deletions
+4 -5
View File
@@ -68,17 +68,16 @@ ENV HOME=/home/user_lerobot \
# issues with MuJoCo and OpenGL drivers. # issues with MuJoCo and OpenGL drivers.
RUN uv venv --python python${PYTHON_VERSION} RUN uv venv --python python${PYTHON_VERSION}
# Install Python dependencies for caching # Install third-party dependencies separately for layer caching
COPY --chown=user_lerobot:user_lerobot setup.py pyproject.toml uv.lock README.md MANIFEST.in ./ COPY --chown=user_lerobot:user_lerobot setup.py pyproject.toml uv.lock README.md MANIFEST.in ./
COPY --chown=user_lerobot:user_lerobot src/ src/ RUN uv sync --locked --extra all --no-install-project --no-cache
RUN uv sync --locked --extra all --no-cache
RUN chmod +x /lerobot/.venv/lib/python${PYTHON_VERSION}/site-packages/triton/backends/nvidia/bin/ptxas RUN chmod +x /lerobot/.venv/lib/python${PYTHON_VERSION}/site-packages/triton/backends/nvidia/bin/ptxas
# Copy the rest of the application source code # Copy the application source code and install the local project
# Make sure to have the git-LFS files for testing # Make sure to have the git-LFS files for testing
COPY --chown=user_lerobot:user_lerobot . . COPY --chown=user_lerobot:user_lerobot . .
RUN uv sync --locked --extra all --no-cache
# Set the default command # Set the default command
CMD ["/bin/bash"] CMD ["/bin/bash"]
+4 -5
View File
@@ -60,15 +60,14 @@ ENV HOME=/home/user_lerobot \
# run other Python projects in the same container without dependency conflicts. # run other Python projects in the same container without dependency conflicts.
RUN uv venv RUN uv venv
# Install Python dependencies for caching # Install third-party dependencies separately for layer caching
COPY --chown=user_lerobot:user_lerobot setup.py pyproject.toml uv.lock README.md MANIFEST.in ./ COPY --chown=user_lerobot:user_lerobot setup.py pyproject.toml uv.lock README.md MANIFEST.in ./
COPY --chown=user_lerobot:user_lerobot src/ src/ RUN uv sync --locked --extra all --no-install-project --no-cache
RUN uv sync --locked --extra all --no-cache # Copy the application code and install the local project
# Copy the rest of the application code
# Make sure to have the git-LFS files for testing # Make sure to have the git-LFS files for testing
COPY --chown=user_lerobot:user_lerobot . . COPY --chown=user_lerobot:user_lerobot . .
RUN uv sync --locked --extra all --no-cache
# Set the default command # Set the default command
CMD ["/bin/bash"] CMD ["/bin/bash"]
+3 -3
View File
@@ -58,7 +58,7 @@ final_action = postprocessor(action)
## Hardware API redesign ## Hardware API redesign
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. 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.
### What changed? ### 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. 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 to first try and and replay an episode of the dataset your policy was trained on using the section above. 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 on your inference script (shown here in the `record.py` script): Then, add these same transformations to your inference script (shown here in the `record.py` script):
```diff ```diff
action_values = predict_action( action_values = predict_action(
+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: For each motor you want to update:
1. **Select the motor** from the list by clicking on it 1. **Select the motor** from the list by clicking on it
2. **Click on Upgrade tab**: 2. **Click the Upgrade tab**:
3. **Click on Online button**: 3. **Click the Online button**:
- If an potential firmware update is found, it will be displayed in the box - If a potential firmware update is found, it will be displayed in the box
4. **Click on Upgrade button**: 4. **Click the Upgrade button**:
- The update progress will be displayed - The update progress will be displayed
## Step 6: Verify Update ## Step 6: Verify Update
+9 -163
View File
@@ -1,177 +1,23 @@
# LeRobot
<div class="flex justify-center"> <div class="flex justify-center">
<a target="_blank" href="https://huggingface.co/lerobot"> <a target="_blank" href="https://huggingface.co/lerobot">
<img <img
alt="LeRobot, Hugging Face Robotics Library" alt="HuggingFace Expert Acceleration Program"
src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/lerobot/lerobot-logo-thumbnail.png" src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/lerobot/lerobot-logo-thumbnail.png"
style="width: 100%"
></img> ></img>
</a> </a>
</div> </div>
# LeRobot
**State-of-the-art machine learning for real-world robotics** **State-of-the-art machine learning for real-world robotics**
🤗 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 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.
🤗 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 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.
The goal: lower the barrier to entry for robotics, so that everyone can contribute to, and benefit from, shared datasets and pretrained models. 🤗 LeRobot already provides a set of pretrained models, datasets with human collected demonstrations, and simulated environments so that everyone can get started.
<div align="center" style="display: flex; justify-content: center; gap: 8px; flex-wrap: wrap; margin: 20px 0;"> 🤗 LeRobot hosts pretrained models and datasets on the LeRobot HuggingFace page.
<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>
<div align="center"> Join the LeRobot community on [Discord](https://discord.gg/s3KuuzsPFb)
<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.
+4 -4
View File
@@ -22,7 +22,7 @@ With processors, you choose the learning features you want to use for your polic
## Three pipelines ## Three pipelines
We often compose three pipelines. Depending on your setup, some can be empty if action and observation spaces already match. We often compose three pipelines. Depending on your setup, some can be empty if action and observation spaces already match.
Each of these pipelines handle different conversions between different action and observation spaces. Below is a quick explanation of each pipeline. Each of these pipelines handles 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) 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) 2. Pipeline 2: Dataset action space → robot command space (EE targets → joints)
@@ -74,15 +74,15 @@ 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. - `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. - `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. - `observation_to_transition`: transforms the robot observation dict to a pipeline transition.
- `transition_to_observation`: transforms the pipeline transition to a observation dict. - `transition_to_observation`: transforms the pipeline transition to an observation dict.
Checkout [src/lerobot/processor/converters.py](https://github.com/huggingface/lerobot/blob/main/src/lerobot/processor/converters.py) for more details. Check out [src/lerobot/processor/converters.py](https://github.com/huggingface/lerobot/blob/main/src/lerobot/processor/converters.py) for more details.
## Dataset feature contracts ## Dataset feature contracts
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(...)`. 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 and example of how we declare features with the `transform_features` method in the phone to SO-100 follower examples: Below is an example of how we declare features with the `transform_features` method in the phone to SO-100 follower examples:
```python ```python
def transform_features( 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") policy = PI0Policy.from_pretrained("lerobot/pi0_base", policy_cfg=policy_cfg, device="cuda")
# Now use predict_action_chunk with RTC parameters # Now use predict_action_chunk with RTC parameters
inference_delay = 4 # How many steps of inference latency, this values should be calculated based on the inference latency of the policy inference_delay = 4 # How many steps of inference latency, this value should be calculated based on the inference latency of the policy
# Initialize the action queue # Initialize the action queue
action_queue = ActionQueue(policy_cfg.rtc_config) action_queue = ActionQueue(policy_cfg.rtc_config)
@@ -100,7 +100,7 @@ Typical values: 8-12 steps
RTCConfig(execution_horizon=10) 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 a 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 an optimal value.
**`prefix_attention_schedule`**: How to weight consistency across the overlap region. **`prefix_attention_schedule`**: How to weight consistency across the overlap region.
+2 -2
View File
@@ -50,11 +50,11 @@ lerobot-edit-dataset \
Divide a dataset into multiple subsets. Divide a dataset into multiple subsets.
```bash ```bash
# Split by fractions (e.g. 80% train, 20% test, 20% val) # Split by fractions (e.g. 60% train, 20% val, 20% test)
lerobot-edit-dataset \ lerobot-edit-dataset \
--repo_id lerobot/pusht \ --repo_id lerobot/pusht \
--operation.type split \ --operation.type split \
--operation.splits '{"train": 0.8, "test": 0.2, "val": 0.2}' --operation.splits '{"train": 0.6, "val": 0.2, "test": 0.2}'
# Split by specific episode indices # Split by specific episode indices
lerobot-edit-dataset \ lerobot-edit-dataset \
+13
View File
@@ -494,6 +494,19 @@ ignore_errors = true
module = "lerobot.envs.*" module = "lerobot.envs.*"
ignore_errors = false 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]] # [[tool.mypy.overrides]]
# module = "lerobot.utils.*" # module = "lerobot.utils.*"
+114 -58
View File
@@ -19,6 +19,7 @@ import copy
import logging import logging
import shutil import shutil
from pathlib import Path from pathlib import Path
from typing import Any, NotRequired, TypedDict
import datasets import datasets
import pandas as pd import pandas as pd
@@ -49,8 +50,32 @@ from .utils import (
) )
from .video_utils import concatenate_video_files, get_video_duration_in_s from .video_utils import concatenate_video_files, get_video_duration_in_s
logger = logging.getLogger(__name__)
def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMetadata]) -> dict[str, dict]: 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:
"""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. """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: Args:
@@ -59,14 +84,14 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
Returns: Returns:
dict: A dictionary of merged video feature info. dict: A dictionary of merged video feature info.
""" """
merged_info = copy.deepcopy(all_metadata[0].features) merged_info: FeatureDict = copy.deepcopy(all_metadata[0].features)
video_keys = [k for k in merged_info if merged_info[k].get("dtype") == "video"] video_keys = [k for k in merged_info if merged_info[k].get("dtype") == "video"]
for vk in video_keys: for vk in video_keys:
video_infos = [m.features.get(vk, {}).get("info") or {} for m in all_metadata] video_infos = [m.features.get(vk, {}).get("info") or {} for m in all_metadata]
base_video_info = video_infos[0] base_video_info = video_infos[0]
merged_encoder_info: dict = {} merged_encoder_info: dict[str, Any] = {}
fallback_keys: list[str] = [] fallback_keys: list[str] = []
for info_key in VIDEO_ENCODER_INFO_KEYS: for info_key in VIDEO_ENCODER_INFO_KEYS:
values = [info.get(info_key, None) for info in video_infos] values = [info.get(info_key, None) for info in video_infos]
@@ -80,7 +105,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 merged_encoder_info[info_key] = {} if info_key == "video.extra_options" else None
if fallback_keys: if fallback_keys:
logging.warning( logger.warning(
f"Merging heterogeneous or incomplete video encoder metadata for feature {vk}. " f"Merging heterogeneous or incomplete video encoder metadata for feature {vk}. "
f"Setting these keys to null: {fallback_keys}.", f"Setting these keys to null: {fallback_keys}.",
) )
@@ -92,7 +117,7 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
return merged_info return merged_info
def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]): def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]) -> tuple[int, str | None, FeatureDict]:
"""Validates that all dataset metadata have consistent properties. """Validates that all dataset metadata have consistent properties.
Ensures all datasets have the same fps, robot_type, and features to guarantee Ensures all datasets have the same fps, robot_type, and features to guarantee
@@ -129,7 +154,9 @@ def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]):
return fps, robot_type, features return fps, robot_type, features
def update_data_df(df, src_meta, dst_meta): def update_data_df(
df: pd.DataFrame, src_meta: LeRobotDatasetMetadata, dst_meta: LeRobotDatasetMetadata
) -> pd.DataFrame:
"""Updates a data DataFrame with new indices and task mappings for aggregation. """Updates a data DataFrame with new indices and task mappings for aggregation.
Adjusts episode indices, frame indices, and task indices to account for Adjusts episode indices, frame indices, and task indices to account for
@@ -154,12 +181,12 @@ def update_data_df(df, src_meta, dst_meta):
def update_meta_data( def update_meta_data(
df, df: pd.DataFrame,
dst_meta, dst_meta: LeRobotDatasetMetadata,
meta_idx, meta_idx: IndexState,
data_idx, data_idx: IndexState,
videos_idx, videos_idx: VideoIndexState,
): ) -> pd.DataFrame:
"""Updates metadata DataFrame with new chunk, file, and timestamp indices. """Updates metadata DataFrame with new chunk, file, and timestamp indices.
Adjusts all indices and timestamps to account for previously aggregated Adjusts all indices and timestamps to account for previously aggregated
@@ -289,7 +316,7 @@ def aggregate_datasets(
chunk_size: int | None = None, chunk_size: int | None = None,
concatenate_videos: bool = True, concatenate_videos: bool = True,
concatenate_data: bool = True, concatenate_data: bool = True,
): ) -> None:
"""Aggregates multiple LeRobot datasets into a single unified dataset. """Aggregates multiple LeRobot datasets into a single unified dataset.
This is the main function that orchestrates the aggregation process by: This is the main function that orchestrates the aggregation process by:
@@ -309,7 +336,7 @@ def aggregate_datasets(
concatenate_videos: When False, keep one mp4 per source file instead of packing into shards. 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. concatenate_data: When False, keep one parquet per source file instead of packing into shards.
""" """
logging.info("Start aggregate_datasets") logger.info("Start aggregate_datasets")
if data_files_size_in_mb is None: if data_files_size_in_mb is None:
data_files_size_in_mb = DEFAULT_DATA_FILE_SIZE_IN_MB data_files_size_in_mb = DEFAULT_DATA_FILE_SIZE_IN_MB
@@ -341,15 +368,15 @@ def aggregate_datasets(
video_files_size_in_mb=video_files_size_in_mb, video_files_size_in_mb=video_files_size_in_mb,
) )
logging.info("Find all tasks") logger.info("Find all tasks")
unique_tasks = pd.concat([m.tasks for m in all_metadata]).index.unique() unique_tasks = pd.concat([m.tasks for m in all_metadata]).index.unique()
dst_meta.tasks = pd.DataFrame( dst_meta.tasks = pd.DataFrame(
{"task_index": range(len(unique_tasks))}, index=pd.Index(unique_tasks, name="task") {"task_index": range(len(unique_tasks))}, index=pd.Index(unique_tasks, name="task")
) )
meta_idx = {"chunk": 0, "file": 0} meta_idx: IndexState = {"chunk": 0, "file": 0}
data_idx = {"chunk": 0, "file": 0} data_idx: IndexState = {"chunk": 0, "file": 0}
videos_idx = { videos_idx: VideoIndexState = {
key: {"chunk": 0, "file": 0, "latest_duration": 0, "episode_duration": 0} for key in video_keys key: {"chunk": 0, "file": 0, "latest_duration": 0, "episode_duration": 0} for key in video_keys
} }
@@ -373,12 +400,17 @@ def aggregate_datasets(
dst_meta.info.total_frames += src_meta.total_frames dst_meta.info.total_frames += src_meta.total_frames
finalize_aggregation(dst_meta, all_metadata) finalize_aggregation(dst_meta, all_metadata)
logging.info("Aggregation complete.") logger.info("Aggregation complete.")
def aggregate_videos( def aggregate_videos(
src_meta, dst_meta, videos_idx, video_files_size_in_mb, chunk_size, concatenate_videos=True src_meta: LeRobotDatasetMetadata,
): dst_meta: LeRobotDatasetMetadata,
videos_idx: VideoIndexState,
video_files_size_in_mb: float,
chunk_size: int,
concatenate_videos: bool = True,
) -> VideoIndexState:
"""Aggregates video chunks from a source dataset into the destination dataset. """Aggregates video chunks from a source dataset into the destination dataset.
Handles video file concatenation and rotation based on file size limits. Handles video file concatenation and rotation based on file size limits.
@@ -406,15 +438,16 @@ def aggregate_videos(
videos_idx[key]["dst_file_durations"] = {} videos_idx[key]["dst_file_durations"] = {}
for key, video_idx in videos_idx.items(): for key, video_idx in videos_idx.items():
unique_chunk_file_pairs = { unique_chunk_file_pairs: list[ChunkFile] = sorted(
(chunk, file) {
for chunk, file in zip( (chunk, file)
src_meta.episodes[f"videos/{key}/chunk_index"], for chunk, file in zip(
src_meta.episodes[f"videos/{key}/file_index"], src_meta.episodes[f"videos/{key}/chunk_index"],
strict=False, src_meta.episodes[f"videos/{key}/file_index"],
) strict=False,
} )
unique_chunk_file_pairs = sorted(unique_chunk_file_pairs) }
)
chunk_idx = video_idx["chunk"] chunk_idx = video_idx["chunk"]
file_idx = video_idx["file"] file_idx = video_idx["file"]
@@ -489,7 +522,14 @@ def aggregate_videos(
return videos_idx return videos_idx
def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_size, concatenate_data=True): 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:
"""Aggregates data chunks from a source dataset into the destination dataset. """Aggregates data chunks from a source dataset into the destination dataset.
Reads source data files, updates indices to match the aggregated dataset, Reads source data files, updates indices to match the aggregated dataset,
@@ -510,14 +550,16 @@ def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_si
Returns: Returns:
dict: Updated data_idx with current chunk and file indices. dict: Updated data_idx with current chunk and file indices.
""" """
unique_chunk_file_ids = { unique_chunk_file_ids: list[ChunkFile] = sorted(
(c, f) {
for c, f in zip( (c, f)
src_meta.episodes["data/chunk_index"], src_meta.episodes["data/file_index"], strict=False for c, f in zip(
) 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 contains_images = len(dst_meta.image_keys) > 0
# retrieve features schema for proper image typing in parquet # retrieve features schema for proper image typing in parquet
@@ -525,7 +567,7 @@ def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_si
# Track source to destination file mapping for metadata update # Track source to destination file mapping for metadata update
# This is critical for handling datasets that are already results of a merge # This is critical for handling datasets that are already results of a merge
src_to_dst: dict[tuple[int, int], tuple[int, int]] = {} src_to_dst: dict[ChunkFile, ChunkFile] = {}
for src_chunk_idx, src_file_idx in unique_chunk_file_ids: for src_chunk_idx, src_file_idx in unique_chunk_file_ids:
src_path = src_meta.root / DEFAULT_DATA_PATH.format( src_path = src_meta.root / DEFAULT_DATA_PATH.format(
@@ -564,7 +606,13 @@ def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_si
return data_idx return data_idx
def aggregate_metadata(src_meta, dst_meta, meta_idx, data_idx, videos_idx): def aggregate_metadata(
src_meta: LeRobotDatasetMetadata,
dst_meta: LeRobotDatasetMetadata,
meta_idx: IndexState,
data_idx: IndexState,
videos_idx: VideoIndexState,
) -> IndexState:
"""Aggregates metadata from a source dataset into the destination dataset. """Aggregates metadata from a source dataset into the destination dataset.
Reads source metadata files, updates all indices and timestamps, Reads source metadata files, updates all indices and timestamps,
@@ -580,16 +628,16 @@ def aggregate_metadata(src_meta, dst_meta, meta_idx, data_idx, videos_idx):
Returns: Returns:
dict: Updated meta_idx with current chunk and file indices. dict: Updated meta_idx with current chunk and file indices.
""" """
chunk_file_ids = { chunk_file_ids: list[ChunkFile] = sorted(
(c, f) {
for c, f in zip( (c, f)
src_meta.episodes["meta/episodes/chunk_index"], for c, f in zip(
src_meta.episodes["meta/episodes/file_index"], src_meta.episodes["meta/episodes/chunk_index"],
strict=False, src_meta.episodes["meta/episodes/file_index"],
) strict=False,
} )
}
chunk_file_ids = sorted(chunk_file_ids) )
for chunk_idx, file_idx in 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) src_path = src_meta.root / DEFAULT_EPISODES_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
df = pd.read_parquet(src_path) df = pd.read_parquet(src_path)
@@ -622,16 +670,16 @@ def aggregate_metadata(src_meta, dst_meta, meta_idx, data_idx, videos_idx):
def append_or_create_parquet_file( def append_or_create_parquet_file(
df: pd.DataFrame, df: pd.DataFrame,
src_path: Path, src_path: Path,
idx: dict[str, int], idx: IndexState,
max_mb: float, max_mb: float,
chunk_size: int, chunk_size: int,
default_path: str, default_path: str,
contains_images: bool = False, contains_images: bool = False,
aggr_root: Path = None, aggr_root: Path | None = None,
hf_features: datasets.Features | None = None, hf_features: datasets.Features | None = None,
concatenate: bool = True, concatenate: bool = True,
one_row_group_per_episode: bool = False, one_row_group_per_episode: bool = False,
) -> tuple[dict[str, int], tuple[int, int]]: ) -> tuple[IndexState, ChunkFile]:
"""Appends data to an existing parquet file or creates a new one based on size constraints. """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 Manages file rotation when size limits are exceeded to prevent individual files
@@ -654,7 +702,13 @@ def append_or_create_parquet_file(
Returns: Returns:
tuple: (updated_idx, (dst_chunk, dst_file)) where updated_idx is the index dict 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. 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_chunk, dst_file = idx["chunk"], idx["file"]
dst_path = aggr_root / default_path.format(chunk_index=dst_chunk, file_index=dst_file) dst_path = aggr_root / default_path.format(chunk_index=dst_chunk, file_index=dst_file)
@@ -698,7 +752,9 @@ def append_or_create_parquet_file(
return idx, (dst_chunk, dst_file) return idx, (dst_chunk, dst_file)
def finalize_aggregation(aggr_meta, all_metadata): def finalize_aggregation(
aggr_meta: LeRobotDatasetMetadata, all_metadata: list[LeRobotDatasetMetadata]
) -> None:
"""Finalizes the dataset aggregation by writing summary files and statistics. """Finalizes the dataset aggregation by writing summary files and statistics.
Writes the tasks file, info file with total counts and splits, and Writes the tasks file, info file with total counts and splits, and
@@ -708,16 +764,16 @@ def finalize_aggregation(aggr_meta, all_metadata):
aggr_meta: Aggregated dataset metadata. aggr_meta: Aggregated dataset metadata.
all_metadata: List of all source dataset metadata objects. all_metadata: List of all source dataset metadata objects.
""" """
logging.info("write tasks") logger.info("write tasks")
write_tasks(aggr_meta.tasks, aggr_meta.root) write_tasks(aggr_meta.tasks, aggr_meta.root)
logging.info("write info") logger.info("write info")
aggr_meta.info.total_tasks = len(aggr_meta.tasks) 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_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.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)}"} aggr_meta.info.splits = {"train": f"0:{sum(m.total_episodes for m in all_metadata)}"}
write_info(aggr_meta.info, aggr_meta.root) write_info(aggr_meta.info, aggr_meta.root)
logging.info("write stats") logger.info("write stats")
aggr_meta.stats = aggregate_stats([m.stats for m in all_metadata]) aggr_meta.stats = aggregate_stats([m.stats for m in all_metadata])
write_stats(aggr_meta.stats, aggr_meta.root) write_stats(aggr_meta.stats, aggr_meta.root)
+2 -2
View File
@@ -188,8 +188,8 @@ class LeRobotDatasetMetadata:
def _load_metadata(self): def _load_metadata(self):
self.info = load_info(self.root) self.info = load_info(self.root)
check_version_compatibility(self.repo_id, self._version, CODEBASE_VERSION) check_version_compatibility(self.repo_id, self._version, CODEBASE_VERSION)
self.tasks = load_tasks(self.root) self.tasks = load_tasks(self.root) if self.total_tasks > 0 else None
self.episodes = load_episodes(self.root) self.episodes = load_episodes(self.root) if self.total_episodes > 0 else None
self.stats = load_stats(self.root) self.stats = load_stats(self.root)
def ensure_readable(self) -> None: def ensure_readable(self) -> None:
+6 -1
View File
@@ -384,7 +384,12 @@ class LiberoEnv(gym.Env):
def close(self): def close(self):
if self._env is not None: if self._env is not None:
self._env.close() 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( def _make_env_fns(
+3 -1
View File
@@ -384,7 +384,9 @@ class RoboTwinEnv(gym.Env):
self._env: Any | None = None # deferred — created on first reset() inside worker self._env: Any | None = None # deferred — created on first reset() inside worker
self._step_count: int = 0 self._step_count: int = 0
self._black_frame = np.zeros((self.observation_height, self.observation_width, 3), dtype=np.uint8) self._black_frame: np.ndarray = np.zeros(
(self.observation_height, self.observation_width, 3), dtype=np.uint8
)
image_spaces = { image_spaces = {
cam: spaces.Box( cam: spaces.Box(
+1 -1
View File
@@ -373,7 +373,7 @@ class VLABenchEnv(gym.Env):
if action.shape[0] != 7: if action.shape[0] != 7:
# Unknown layout — fall back to zero-pad so the sim doesn't crash. # Unknown layout — fall back to zero-pad so the sim doesn't crash.
padded = np.zeros(ctrl_dim, dtype=np.float64) padded: np.ndarray = np.zeros(ctrl_dim, dtype=np.float64)
padded[: min(action.shape[0], ctrl_dim)] = action[:ctrl_dim] padded[: min(action.shape[0], ctrl_dim)] = action[:ctrl_dim]
return padded return padded
+35 -5
View File
@@ -302,6 +302,33 @@ def _pad_evo1_stats(
return padded_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( def reconcile_evo1_processors(
config: Evo1Config, config: Evo1Config,
preprocessor: PolicyProcessorPipeline, preprocessor: PolicyProcessorPipeline,
@@ -309,16 +336,19 @@ def reconcile_evo1_processors(
) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]: ) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]:
"""Reconcile checkpoint-loaded pipelines with the current EVO1 config. """Reconcile checkpoint-loaded pipelines with the current EVO1 config.
Two things cannot be restored from a serialized pipeline alone: the EVO1 batch converter Three 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 (converters are plain functions and are never serialized), eval-time CLI overrides of the
action postprocessing flags (`postprocess_action_dim`, `binarize_gripper`, `gripper_*`). This action postprocessing flags (`postprocess_action_dim`, `binarize_gripper`, `gripper_*`), and the
restores the converter and rebuilds the action step from the current config so those overrides (un)normalizer stats/features when the generic override path injects raw, unpadded dataset
take effect. 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.
""" """
# Pipelines reloaded from a checkpoint come back with the default batch converter, which drops # 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. # non-observation extras (embodiment_id, state_mask, custom task fields) needed by EVO1.
preprocessor.to_transition = evo1_batch_to_transition preprocessor.to_transition = evo1_batch_to_transition
_refresh_evo1_normalization_steps(config, preprocessor, postprocessor)
action_step = Evo1ActionProcessorStep( action_step = Evo1ActionProcessorStep(
action_dim=_evo1_action_dim(config), action_dim=_evo1_action_dim(config),
binarize_gripper=config.binarize_gripper, binarize_gripper=config.binarize_gripper,
@@ -46,6 +46,12 @@ class SOFollowerConfig:
position_i_coefficient: int = 0 position_i_coefficient: int = 0
position_d_coefficient: int = 32 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("so101_follower")
@RobotConfig.register_subclass("so100_follower") @RobotConfig.register_subclass("so100_follower")
+17 -3
View File
@@ -142,11 +142,25 @@ class SOFollower(Robot):
range_mins[full_turn_motor] = 0 range_mins[full_turn_motor] = 0
range_maxes[full_turn_motor] = 4095 range_maxes[full_turn_motor] = 4095
drive_modes = dict.fromkeys(self.bus.motors, 0)
input(f"Fully close the gripper of {self} and press ENTER....")
gripper_closed_pos = self.bus.read(
"Present_Position", "gripper", normalize=False, num_retry=self.config.num_read_retries
)
distance_to_min = abs(gripper_closed_pos - range_mins["gripper"])
distance_to_max = abs(gripper_closed_pos - range_maxes["gripper"])
if min(distance_to_min, distance_to_max) > (range_maxes["gripper"] - range_mins["gripper"]) * 0.2:
raise ValueError("Gripper is not fully closed. Run calibration again.")
drive_modes["gripper"] = int(distance_to_max < distance_to_min)
if drive_modes["gripper"]:
logger.info("Gripper motor is inverted, setting drive_mode=1 to compensate.")
self.calibration = {} self.calibration = {}
for motor, m in self.bus.motors.items(): for motor, m in self.bus.motors.items():
self.calibration[motor] = MotorCalibration( self.calibration[motor] = MotorCalibration(
id=m.id, id=m.id,
drive_mode=0, drive_mode=drive_modes[motor],
homing_offset=homing_offsets[motor], homing_offset=homing_offsets[motor],
range_min=range_mins[motor], range_min=range_mins[motor],
range_max=range_maxes[motor], range_max=range_maxes[motor],
@@ -180,7 +194,7 @@ class SOFollower(Robot):
def get_observation(self) -> RobotObservation: def get_observation(self) -> RobotObservation:
# Read arm position # Read arm position
start = time.perf_counter() start = time.perf_counter()
obs_dict = self.bus.sync_read("Present_Position") obs_dict = self.bus.sync_read("Present_Position", num_retry=self.config.num_read_retries)
obs_dict = {f"{motor}.pos": val for motor, val in obs_dict.items()} obs_dict = {f"{motor}.pos": val for motor, val in obs_dict.items()}
dt_ms = (time.perf_counter() - start) * 1e3 dt_ms = (time.perf_counter() - start) * 1e3
logger.debug(f"{self} read state: {dt_ms:.1f}ms") logger.debug(f"{self} read state: {dt_ms:.1f}ms")
@@ -221,7 +235,7 @@ class SOFollower(Robot):
# Cap goal position when too far away from present position. # Cap goal position when too far away from present position.
# /!\ Slower fps expected due to reading from the follower. # /!\ Slower fps expected due to reading from the follower.
if self.config.max_relative_target is not None: if self.config.max_relative_target is not None:
present_pos = self.bus.sync_read("Present_Position") present_pos = self.bus.sync_read("Present_Position", num_retry=self.config.num_read_retries)
goal_present_pos = {key: (g_pos, present_pos[key]) for key, g_pos in goal_pos.items()} 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) goal_pos = ensure_safe_goal_position(goal_present_pos, self.config.max_relative_target)
@@ -36,6 +36,7 @@ python src/lerobot/scripts/augment_dataset_quantile_stats.py \
import argparse import argparse
import concurrent.futures import concurrent.futures
import logging import logging
import os
from pathlib import Path from pathlib import Path
import numpy as np import numpy as np
@@ -52,6 +53,7 @@ from lerobot.datasets import (
get_feature_stats, get_feature_stats,
write_stats, write_stats,
) )
from lerobot.datasets.compute_stats import sample_indices
from lerobot.utils.utils import init_logging from lerobot.utils.utils import init_logging
@@ -77,12 +79,14 @@ def has_quantile_stats(stats: dict[str, dict] | None, quantile_list_keys: list[s
return False return False
def process_single_episode(dataset: LeRobotDataset, episode_idx: int) -> dict: def process_single_episode(dataset: LeRobotDataset, episode_idx: int, use_sampling: bool = True) -> dict:
"""Process a single episode and return its statistics. """Process a single episode and return its statistics.
Args: Args:
dataset: The LeRobot dataset dataset: The LeRobot dataset
episode_idx: Index of the episode to process 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: Returns:
Dictionary containing episode statistics Dictionary containing episode statistics
@@ -92,16 +96,31 @@ def process_single_episode(dataset: LeRobotDataset, episode_idx: int) -> dict:
start_idx = dataset.meta.episodes[episode_idx]["dataset_from_index"] start_idx = dataset.meta.episodes[episode_idx]["dataset_from_index"]
end_idx = dataset.meta.episodes[episode_idx]["dataset_to_index"] end_idx = dataset.meta.episodes[episode_idx]["dataset_to_index"]
collected_data: dict[str, list] = {} episode_len = end_idx - start_idx
for idx in range(start_idx, end_idx):
item = dataset[idx]
for key, value in item.items():
if key not in dataset.features:
continue
if key not in collected_data: # Images/video are the memory hog, so sub-sample those frames per episode;
collected_data[key] = [] # numeric columns are cheap, so read them in full (exact).
collected_data[key].append(value) 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] = {}
# 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])
ep_stats = {} ep_stats = {}
for key, data_list in collected_data.items(): for key, data_list in collected_data.items():
@@ -131,11 +150,13 @@ def process_single_episode(dataset: LeRobotDataset, episode_idx: int) -> dict:
return ep_stats return ep_stats
def compute_quantile_stats_for_dataset(dataset: LeRobotDataset) -> dict[str, dict]: def compute_quantile_stats_for_dataset(dataset: LeRobotDataset, use_sampling: bool = True) -> dict[str, dict]:
"""Compute quantile statistics for all episodes in the dataset. """Compute quantile statistics for all episodes in the dataset.
Args: Args:
dataset: The LeRobot dataset to compute statistics for 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: Returns:
Dictionary containing aggregated statistics with quantiles Dictionary containing aggregated statistics with quantiles
@@ -153,15 +174,15 @@ def compute_quantile_stats_for_dataset(dataset: LeRobotDataset) -> dict[str, dic
if has_videos: if has_videos:
logging.info("Dataset contains video keys - using sequential processing for thread safety") logging.info("Dataset contains video keys - using sequential processing for thread safety")
for episode_idx in tqdm(range(dataset.num_episodes), desc="Processing episodes"): for episode_idx in tqdm(range(dataset.num_episodes), desc="Processing episodes"):
ep_stats = process_single_episode(dataset, episode_idx) ep_stats = process_single_episode(dataset, episode_idx, use_sampling)
episode_stats_list.append(ep_stats) episode_stats_list.append(ep_stats)
else: else:
logging.info("Dataset has no video keys - using parallel processing for better performance") logging.info("Dataset has no video keys - using parallel processing for better performance")
max_workers = min(dataset.num_episodes, 16) max_workers = min(dataset.num_episodes, int(os.environ.get("LEROBOT_STATS_MAX_WORKERS", 16)))
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
future_to_episode = { future_to_episode = {
executor.submit(process_single_episode, dataset, episode_idx): episode_idx executor.submit(process_single_episode, dataset, episode_idx, use_sampling): episode_idx
for episode_idx in range(dataset.num_episodes) for episode_idx in range(dataset.num_episodes)
} }
@@ -188,6 +209,7 @@ def augment_dataset_with_quantile_stats(
repo_id: str, repo_id: str,
root: str | Path | None = None, root: str | Path | None = None,
overwrite: bool = False, overwrite: bool = False,
use_sampling: bool = True,
) -> None: ) -> None:
"""Augment a dataset with quantile statistics if they are missing. """Augment a dataset with quantile statistics if they are missing.
@@ -195,6 +217,8 @@ def augment_dataset_with_quantile_stats(
repo_id: Repository ID of the dataset repo_id: Repository ID of the dataset
root: Local root directory for the dataset root: Local root directory for the dataset
overwrite: Overwrite existing quantile statistics if they already exist 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}") logging.info(f"Loading dataset: {repo_id}")
dataset = LeRobotDataset( dataset = LeRobotDataset(
@@ -208,7 +232,7 @@ def augment_dataset_with_quantile_stats(
logging.info("Dataset does not contain quantile statistics. Computing them now...") logging.info("Dataset does not contain quantile statistics. Computing them now...")
new_stats = compute_quantile_stats_for_dataset(dataset) new_stats = compute_quantile_stats_for_dataset(dataset, use_sampling=use_sampling)
logging.info("Updating dataset metadata with new quantile statistics") logging.info("Updating dataset metadata with new quantile statistics")
dataset.meta.stats = new_stats dataset.meta.stats = new_stats
@@ -248,6 +272,14 @@ def main():
action="store_true", action="store_true",
help="Overwrite existing quantile statistics if they already exist", 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() args = parser.parse_args()
root = Path(args.root) if args.root else None root = Path(args.root) if args.root else None
@@ -258,6 +290,7 @@ def main():
repo_id=args.repo_id, repo_id=args.repo_id,
root=root, root=root,
overwrite=args.overwrite, overwrite=args.overwrite,
use_sampling=not args.no_sampling,
) )
+1 -1
View File
@@ -564,7 +564,7 @@ def eval_policy(
if seeds: if seeds:
all_seeds.extend(seeds) all_seeds.extend(seeds)
else: else:
all_seeds.append(None) all_seeds.extend([None] * env.num_envs)
# FIXME: episode_data is either None or it doesn't exist # FIXME: episode_data is either None or it doesn't exist
if return_episode_data: if return_episode_data:
+17 -13
View File
@@ -22,7 +22,8 @@ import dataclasses
import logging import logging
import sys import sys
import time import time
from contextlib import nullcontext from collections.abc import Iterator
from contextlib import contextmanager, nullcontext
from pprint import pformat from pprint import pformat
from typing import TYPE_CHECKING, Any from typing import TYPE_CHECKING, Any
@@ -76,6 +77,20 @@ else:
from .lerobot_eval import eval_policy_all 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]: def _dataloader_worker_kwargs(cfg: TrainPipelineConfig) -> dict[str, Any]:
"""Return worker-only DataLoader options, disabling them for single-process loading.""" """Return worker-only DataLoader options, disabling them for single-process loading."""
workers_enabled = cfg.num_workers > 0 workers_enabled = cfg.num_workers > 0
@@ -280,14 +295,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if not is_main_process: if not is_main_process:
dataset, eval_dataset = make_train_eval_datasets(cfg) 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 cfg.is_reward_model_training:
if is_main_process: if is_main_process:
logging.info("Creating reward model") logging.info("Creating reward model")
@@ -695,7 +702,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if is_main_process: if is_main_process:
step_id = get_step_identifier(step, cfg.steps) step_id = get_step_identifier(step, cfg.steps)
logging.info(f"Eval policy at step {step}") logging.info(f"Eval policy at step {step}")
with torch.no_grad(), accelerator.autocast(): with _make_eval_envs(cfg) as eval_env, torch.no_grad(), accelerator.autocast():
eval_info = eval_policy_all( eval_info = eval_policy_all(
envs=eval_env, # dict[suite][task_id] -> vec_env envs=eval_env, # dict[suite][task_id] -> vec_env
policy=accelerator.unwrap_model(policy), policy=accelerator.unwrap_model(policy),
@@ -743,9 +750,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if is_main_process: if is_main_process:
progbar.close() progbar.close()
if eval_env:
close_envs(eval_env)
is_fsdp = accelerator.distributed_type == DistributedType.FSDP is_fsdp = accelerator.distributed_type == DistributedType.FSDP
model_state_dict = accelerator.get_state_dict(policy) if is_fsdp else None model_state_dict = accelerator.get_state_dict(policy) if is_fsdp else None
if is_main_process: if is_main_process:
@@ -29,6 +29,12 @@ class SOLeaderConfig:
# Whether to use degrees for angles # Whether to use degrees for angles
use_degrees: bool = True 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("so101_leader")
@TeleoperatorConfig.register_subclass("so100_leader") @TeleoperatorConfig.register_subclass("so100_leader")
@@ -110,11 +110,25 @@ class SOLeader(Teleoperator):
range_mins[full_turn_motor] = 0 range_mins[full_turn_motor] = 0
range_maxes[full_turn_motor] = 4095 range_maxes[full_turn_motor] = 4095
drive_modes = dict.fromkeys(self.bus.motors, 0)
input(f"Fully close the gripper of {self} and press ENTER....")
gripper_closed_pos = self.bus.read(
"Present_Position", "gripper", normalize=False, num_retry=self.config.num_read_retries
)
distance_to_min = abs(gripper_closed_pos - range_mins["gripper"])
distance_to_max = abs(gripper_closed_pos - range_maxes["gripper"])
if min(distance_to_min, distance_to_max) > (range_maxes["gripper"] - range_mins["gripper"]) * 0.2:
raise ValueError("Gripper is not fully closed. Run calibration again.")
drive_modes["gripper"] = int(distance_to_max < distance_to_min)
if drive_modes["gripper"]:
logger.info("Gripper motor is inverted, setting drive_mode=1 to compensate.")
self.calibration = {} self.calibration = {}
for motor, m in self.bus.motors.items(): for motor, m in self.bus.motors.items():
self.calibration[motor] = MotorCalibration( self.calibration[motor] = MotorCalibration(
id=m.id, id=m.id,
drive_mode=0, drive_mode=drive_modes[motor],
homing_offset=homing_offsets[motor], homing_offset=homing_offsets[motor],
range_min=range_mins[motor], range_min=range_mins[motor],
range_max=range_maxes[motor], range_max=range_maxes[motor],
@@ -145,7 +159,7 @@ class SOLeader(Teleoperator):
@check_if_not_connected @check_if_not_connected
def get_action(self) -> dict[str, float]: def get_action(self) -> dict[str, float]:
start = time.perf_counter() start = time.perf_counter()
action = self.bus.sync_read("Present_Position") action = self.bus.sync_read("Present_Position", num_retry=self.config.num_read_retries)
action = {f"{motor}.pos": val for motor, val in action.items()} action = {f"{motor}.pos": val for motor, val in action.items()}
dt_ms = (time.perf_counter() - start) * 1e3 dt_ms = (time.perf_counter() - start) * 1e3
logger.debug(f"{self} read action: {dt_ms:.1f}ms") logger.debug(f"{self} read action: {dt_ms:.1f}ms")
+7 -7
View File
@@ -41,7 +41,7 @@ class RandomSubsetApply(Transform):
def __init__( def __init__(
self, self,
transforms: Sequence[Callable], transforms: Sequence[Callable[..., Any]],
p: list[float] | None = None, p: list[float] | None = None,
n_subset: int | None = None, n_subset: int | None = None,
random_order: bool = False, random_order: bool = False,
@@ -50,7 +50,7 @@ class RandomSubsetApply(Transform):
if not isinstance(transforms, Sequence): if not isinstance(transforms, Sequence):
raise TypeError("Argument transforms should be a sequence of callables") raise TypeError("Argument transforms should be a sequence of callables")
if p is None: if p is None:
p = [1] * len(transforms) p = [1.0] * len(transforms)
elif len(p) != len(transforms): elif len(p) != len(transforms):
raise ValueError( raise ValueError(
f"Length of p doesn't match the number of transforms: {len(p)} != {len(transforms)}" 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.n_subset = n_subset
self.random_order = random_order self.random_order = random_order
self.selected_transforms = None self.selected_transforms: list[Callable[..., Any]] = []
def forward(self, *inputs: Any) -> Any: def forward(self, *inputs: Any) -> Any:
needs_unpacking = len(inputs) > 1 needs_unpacking = len(inputs) > 1
@@ -119,7 +119,7 @@ class SharpnessJitter(Transform):
super().__init__() super().__init__()
self.sharpness = self._check_input(sharpness) self.sharpness = self._check_input(sharpness)
def _check_input(self, sharpness): def _check_input(self, sharpness: float | Sequence[float]) -> tuple[float, float]:
if isinstance(sharpness, (int | float)): if isinstance(sharpness, (int | float)):
if sharpness < 0: if sharpness < 0:
raise ValueError("If sharpness is a single number, it must be non negative.") 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): def make_transform_from_config(cfg: ImageTransformConfig) -> Transform:
if cfg.type == "SharpnessJitter": if cfg.type == "SharpnessJitter":
return SharpnessJitter(**cfg.kwargs) return SharpnessJitter(**cfg.kwargs)
@@ -236,8 +236,8 @@ class ImageTransforms(Transform):
super().__init__() super().__init__()
self._cfg = cfg self._cfg = cfg
self.weights = [] self.weights: list[float] = []
self.transforms = {} self.transforms: dict[str, Transform] = {}
for tf_name, tf_cfg in cfg.tfs.items(): for tf_name, tf_cfg in cfg.tfs.items():
if tf_cfg.weight <= 0.0: if tf_cfg.weight <= 0.0:
continue continue
+7 -4
View File
@@ -133,10 +133,13 @@ def say(text: str, blocking: bool = False):
else: else:
raise RuntimeError("Unsupported operating system for text-to-speech.") raise RuntimeError("Unsupported operating system for text-to-speech.")
if blocking: try:
subprocess.run(cmd, check=True) if blocking:
else: subprocess.run(cmd, check=True, timeout=5)
subprocess.Popen(cmd, creationflags=subprocess.CREATE_NO_WINDOW if system == "Windows" else 0) 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): def log_say(text: str, play_sounds: bool = True, blocking: bool = False):
@@ -0,0 +1,104 @@
# 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,6 +482,20 @@ def test_add_frame_works_in_write_mode(tmp_path):
# ── Resume mode ────────────────────────────────────────────────────── # ── 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): def test_resume_creates_writer(tmp_path):
"""After resume(), writer is a DatasetWriter.""" """After resume(), writer is a DatasetWriter."""
root = tmp_path / "resume_ds" root = tmp_path / "resume_ds"
+13
View File
@@ -294,6 +294,19 @@ def test__sync_read(addr, length, ids_values, mock_motors, dummy_motors):
assert read_values == ids_values 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)) @pytest.mark.parametrize("raise_on_error", (True, False))
def test__sync_read_comm(raise_on_error, mock_motors, dummy_motors): def test__sync_read_comm(raise_on_error, mock_motors, dummy_motors):
addr, length, ids_values = (10, 4, {1: 1337}) addr, length, ids_values = (10, 4, {1: 1337})
+54
View File
@@ -496,6 +496,60 @@ def test_evo1_processor_save_load_round_trip_applies_config_overrides(tmp_path):
assert "embodiment_id" in processed 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): def test_evo1_policy_forward_and_inference_use_batched_embedding(monkeypatch):
monkeypatch.setattr(modeling_evo1, "Evo1Model", DummyEvo1Model) monkeypatch.setattr(modeling_evo1, "Evo1Model", DummyEvo1Model)
policy = modeling_evo1.Evo1Policy(make_config()) policy = modeling_evo1.Evo1Policy(make_config())
+12 -8
View File
@@ -12,7 +12,6 @@
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
import sys
from types import SimpleNamespace from types import SimpleNamespace
from unittest.mock import MagicMock from unittest.mock import MagicMock
@@ -47,18 +46,23 @@ def test_make_policy_keeps_peft_adapter_and_base_revisions_separate(monkeypatch)
peft_config_from_pretrained = MagicMock(return_value=peft_config) peft_config_from_pretrained = MagicMock(return_value=peft_config)
adapted_policy = torch.nn.Linear(1, 1) adapted_policy = torch.nn.Linear(1, 1)
peft_model_from_pretrained = MagicMock(return_value=adapted_policy) peft_model_from_pretrained = MagicMock(return_value=adapted_policy)
monkeypatch.setitem( require_package = MagicMock()
sys.modules, monkeypatch.setattr(policy_factory, "require_package", require_package)
"peft", monkeypatch.setattr(
SimpleNamespace( policy_factory,
PeftConfig=SimpleNamespace(from_pretrained=peft_config_from_pretrained), "PeftConfig",
PeftModel=SimpleNamespace(from_pretrained=peft_model_from_pretrained), SimpleNamespace(from_pretrained=peft_config_from_pretrained),
), )
monkeypatch.setattr(
policy_factory,
"PeftModel",
SimpleNamespace(from_pretrained=peft_model_from_pretrained),
) )
policy = policy_factory.make_policy(cfg, ds_meta=dataset_meta) policy = policy_factory.make_policy(cfg, ds_meta=dataset_meta)
assert policy is adapted_policy assert policy is adapted_policy
require_package.assert_called_once_with("peft", extra="peft")
peft_config_from_pretrained.assert_called_once_with( peft_config_from_pretrained.assert_called_once_with(
"user/adapter", "user/adapter",
revision="adapter-sha", revision="adapter-sha",
+71 -2
View File
@@ -49,7 +49,7 @@ def _make_bus_mock() -> MagicMock:
@pytest.fixture @pytest.fixture
def follower(): def follower(tmp_path):
bus_mock = _make_bus_mock() bus_mock = _make_bus_mock()
def _bus_side_effect(*_args, **kwargs): def _bus_side_effect(*_args, **kwargs):
@@ -71,7 +71,7 @@ def follower():
), ),
patch.object(SO100Follower, "configure", lambda self: None), patch.object(SO100Follower, "configure", lambda self: None),
): ):
cfg = SO100FollowerConfig(port="/dev/null") cfg = SO100FollowerConfig(port="/dev/null", calibration_dir=tmp_path)
robot = SO100Follower(cfg) robot = SO100Follower(cfg)
yield robot yield robot
if robot.is_connected: if robot.is_connected:
@@ -99,6 +99,27 @@ def test_get_observation(follower):
assert obs[f"{motor}.pos"] == idx 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): def test_send_action(follower):
follower.connect() follower.connect()
@@ -128,3 +149,51 @@ def test_configure_writes_position_pid_coefficients():
bus_mock.write.assert_any_call("P_Coefficient", "shoulder_pan", 32) bus_mock.write.assert_any_call("P_Coefficient", "shoulder_pan", 32)
bus_mock.write.assert_any_call("I_Coefficient", "shoulder_pan", 1) bus_mock.write.assert_any_call("I_Coefficient", "shoulder_pan", 1)
bus_mock.write.assert_any_call("D_Coefficient", "shoulder_pan", 16) bus_mock.write.assert_any_call("D_Coefficient", "shoulder_pan", 16)
@pytest.mark.parametrize(
"gripper_closed_pos, expected_drive_mode",
[
(2035, 0), # closed position at range_min -> raw increases when opening -> not inverted
(3528, 1), # closed position at range_max -> raw increases when closing -> inverted
(2781, None), # not near either end stop -> unsafe to infer
],
)
def test_calibrate_detects_gripper_drive_mode(follower, gripper_closed_pos, expected_drive_mode):
"""Regression test for #3942: the follower gripper can be mounted mirrored with respect to the
leader's, in which case its raw position increases when closing. Calibration must detect this
and set drive_mode=1 so that normalized values follow the 0=closed/100=open convention."""
follower.connect()
motors = list(follower.bus.motors)
follower.bus.set_half_turn_homings.return_value = dict.fromkeys(motors, 0)
follower.bus.record_ranges_of_motion.return_value = (
dict.fromkeys(motors, 2035),
dict.fromkeys(motors, 3528),
)
follower.bus.read.return_value = gripper_closed_pos
with (
patch("builtins.input", return_value=""),
patch.object(type(follower), "_save_calibration", lambda self: None),
):
follower.calibration = {}
if expected_drive_mode is None:
with pytest.raises(ValueError, match="Gripper is not fully closed"):
follower.calibrate()
else:
follower.calibrate()
follower.bus.read.assert_called_with(
"Present_Position",
"gripper",
normalize=False,
num_retry=follower.config.num_read_retries,
)
if expected_drive_mode is None:
return
assert follower.calibration["gripper"].drive_mode == expected_drive_mode
for motor in motors:
if motor != "gripper":
assert follower.calibration[motor].drive_mode == 0