mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-29 20:49:42 +00:00
Compare commits
16 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 0d788abd85 | |||
| 7b78e751a6 | |||
| 6e5f6df6e7 | |||
| 265abe6c79 | |||
| b4e2d0b610 | |||
| 5594eba06a | |||
| 207183c2f8 | |||
| 7d615acf9a | |||
| 09572babee | |||
| 35339d31e5 | |||
| f37be3edbe | |||
| 4d076845ac | |||
| 413972c812 | |||
| 0449aa02f6 | |||
| a05c0833e1 | |||
| 7b76d94c5b |
@@ -68,17 +68,16 @@ ENV HOME=/home/user_lerobot \
|
||||
# issues with MuJoCo and OpenGL drivers.
|
||||
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 src/ src/
|
||||
|
||||
RUN uv sync --locked --extra all --no-cache
|
||||
RUN uv sync --locked --extra all --no-install-project --no-cache
|
||||
|
||||
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
|
||||
COPY --chown=user_lerobot:user_lerobot . .
|
||||
RUN uv sync --locked --extra all --no-cache
|
||||
|
||||
# Set the default command
|
||||
CMD ["/bin/bash"]
|
||||
|
||||
@@ -60,15 +60,14 @@ ENV HOME=/home/user_lerobot \
|
||||
# run other Python projects in the same container without dependency conflicts.
|
||||
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 src/ src/
|
||||
RUN uv sync --locked --extra all --no-install-project --no-cache
|
||||
|
||||
RUN uv sync --locked --extra all --no-cache
|
||||
|
||||
# Copy the rest of the application code
|
||||
# Copy the application code and install the local project
|
||||
# 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"]
|
||||
|
||||
@@ -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 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?
|
||||
|
||||
@@ -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 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):
|
||||
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):
|
||||
|
||||
```diff
|
||||
action_values = predict_action(
|
||||
|
||||
@@ -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 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**:
|
||||
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**:
|
||||
- The update progress will be displayed
|
||||
|
||||
## Step 6: Verify Update
|
||||
|
||||
+9
-163
@@ -1,177 +1,23 @@
|
||||
# LeRobot
|
||||
|
||||
<div class="flex justify-center">
|
||||
<a target="_blank" href="https://huggingface.co/lerobot">
|
||||
<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"
|
||||
style="width: 100%"
|
||||
></img>
|
||||
</a>
|
||||
</div>
|
||||
|
||||
# LeRobot
|
||||
|
||||
**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;">
|
||||
<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>
|
||||
🤗 LeRobot hosts pretrained models and datasets on the LeRobot HuggingFace page.
|
||||
|
||||
<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.
|
||||
Join the LeRobot community on [Discord](https://discord.gg/s3KuuzsPFb)
|
||||
|
||||
@@ -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 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)
|
||||
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.
|
||||
- `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 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 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
|
||||
def transform_features(
|
||||
|
||||
+2
-2
@@ -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 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
|
||||
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 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.
|
||||
|
||||
|
||||
@@ -50,11 +50,11 @@ lerobot-edit-dataset \
|
||||
Divide a dataset into multiple subsets.
|
||||
|
||||
```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 \
|
||||
--repo_id lerobot/pusht \
|
||||
--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
|
||||
lerobot-edit-dataset \
|
||||
|
||||
@@ -494,6 +494,19 @@ 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.*"
|
||||
|
||||
@@ -19,6 +19,7 @@ import copy
|
||||
import logging
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from typing import Any, NotRequired, TypedDict
|
||||
|
||||
import datasets
|
||||
import pandas as pd
|
||||
@@ -49,8 +50,32 @@ from .utils import (
|
||||
)
|
||||
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.
|
||||
|
||||
Args:
|
||||
@@ -59,14 +84,14 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
|
||||
Returns:
|
||||
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"]
|
||||
|
||||
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 = {}
|
||||
merged_encoder_info: dict[str, Any] = {}
|
||||
fallback_keys: list[str] = []
|
||||
for info_key in VIDEO_ENCODER_INFO_KEYS:
|
||||
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
|
||||
|
||||
if fallback_keys:
|
||||
logging.warning(
|
||||
logger.warning(
|
||||
f"Merging heterogeneous or incomplete video encoder metadata for feature {vk}. "
|
||||
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
|
||||
|
||||
|
||||
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.
|
||||
|
||||
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
|
||||
|
||||
|
||||
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.
|
||||
|
||||
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(
|
||||
df,
|
||||
dst_meta,
|
||||
meta_idx,
|
||||
data_idx,
|
||||
videos_idx,
|
||||
):
|
||||
df: pd.DataFrame,
|
||||
dst_meta: LeRobotDatasetMetadata,
|
||||
meta_idx: IndexState,
|
||||
data_idx: IndexState,
|
||||
videos_idx: VideoIndexState,
|
||||
) -> pd.DataFrame:
|
||||
"""Updates metadata DataFrame with new chunk, file, and timestamp indices.
|
||||
|
||||
Adjusts all indices and timestamps to account for previously aggregated
|
||||
@@ -289,7 +316,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:
|
||||
@@ -309,7 +336,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.
|
||||
"""
|
||||
logging.info("Start aggregate_datasets")
|
||||
logger.info("Start aggregate_datasets")
|
||||
|
||||
if data_files_size_in_mb is None:
|
||||
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,
|
||||
)
|
||||
|
||||
logging.info("Find all tasks")
|
||||
logger.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 = {"chunk": 0, "file": 0}
|
||||
data_idx = {"chunk": 0, "file": 0}
|
||||
videos_idx = {
|
||||
meta_idx: IndexState = {"chunk": 0, "file": 0}
|
||||
data_idx: IndexState = {"chunk": 0, "file": 0}
|
||||
videos_idx: VideoIndexState = {
|
||||
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
|
||||
|
||||
finalize_aggregation(dst_meta, all_metadata)
|
||||
logging.info("Aggregation complete.")
|
||||
logger.info("Aggregation complete.")
|
||||
|
||||
|
||||
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.
|
||||
|
||||
Handles video file concatenation and rotation based on file size limits.
|
||||
@@ -406,15 +438,16 @@ def aggregate_videos(
|
||||
videos_idx[key]["dst_file_durations"] = {}
|
||||
|
||||
for key, video_idx in videos_idx.items():
|
||||
unique_chunk_file_pairs = {
|
||||
(chunk, file)
|
||||
for chunk, file in zip(
|
||||
src_meta.episodes[f"videos/{key}/chunk_index"],
|
||||
src_meta.episodes[f"videos/{key}/file_index"],
|
||||
strict=False,
|
||||
)
|
||||
}
|
||||
unique_chunk_file_pairs = sorted(unique_chunk_file_pairs)
|
||||
unique_chunk_file_pairs: list[ChunkFile] = sorted(
|
||||
{
|
||||
(chunk, file)
|
||||
for chunk, file in zip(
|
||||
src_meta.episodes[f"videos/{key}/chunk_index"],
|
||||
src_meta.episodes[f"videos/{key}/file_index"],
|
||||
strict=False,
|
||||
)
|
||||
}
|
||||
)
|
||||
|
||||
chunk_idx = video_idx["chunk"]
|
||||
file_idx = video_idx["file"]
|
||||
@@ -489,7 +522,14 @@ def aggregate_videos(
|
||||
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.
|
||||
|
||||
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:
|
||||
dict: Updated data_idx with current chunk and file indices.
|
||||
"""
|
||||
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
|
||||
)
|
||||
}
|
||||
|
||||
unique_chunk_file_ids = sorted(unique_chunk_file_ids)
|
||||
unique_chunk_file_ids: list[ChunkFile] = sorted(
|
||||
{
|
||||
(c, f)
|
||||
for c, f in zip(
|
||||
src_meta.episodes["data/chunk_index"],
|
||||
src_meta.episodes["data/file_index"],
|
||||
strict=False,
|
||||
)
|
||||
}
|
||||
)
|
||||
contains_images = len(dst_meta.image_keys) > 0
|
||||
|
||||
# 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
|
||||
# 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:
|
||||
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
|
||||
|
||||
|
||||
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.
|
||||
|
||||
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:
|
||||
dict: Updated meta_idx with current chunk and file indices.
|
||||
"""
|
||||
chunk_file_ids = {
|
||||
(c, f)
|
||||
for c, f in zip(
|
||||
src_meta.episodes["meta/episodes/chunk_index"],
|
||||
src_meta.episodes["meta/episodes/file_index"],
|
||||
strict=False,
|
||||
)
|
||||
}
|
||||
|
||||
chunk_file_ids = sorted(chunk_file_ids)
|
||||
chunk_file_ids: list[ChunkFile] = sorted(
|
||||
{
|
||||
(c, f)
|
||||
for c, f in zip(
|
||||
src_meta.episodes["meta/episodes/chunk_index"],
|
||||
src_meta.episodes["meta/episodes/file_index"],
|
||||
strict=False,
|
||||
)
|
||||
}
|
||||
)
|
||||
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)
|
||||
@@ -622,16 +670,16 @@ def aggregate_metadata(src_meta, dst_meta, meta_idx, data_idx, videos_idx):
|
||||
def append_or_create_parquet_file(
|
||||
df: pd.DataFrame,
|
||||
src_path: Path,
|
||||
idx: dict[str, int],
|
||||
idx: IndexState,
|
||||
max_mb: float,
|
||||
chunk_size: int,
|
||||
default_path: str,
|
||||
contains_images: bool = False,
|
||||
aggr_root: Path = None,
|
||||
aggr_root: Path | None = None,
|
||||
hf_features: datasets.Features | None = None,
|
||||
concatenate: bool = True,
|
||||
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.
|
||||
|
||||
Manages file rotation when size limits are exceeded to prevent individual files
|
||||
@@ -654,7 +702,13 @@ 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)
|
||||
|
||||
@@ -698,7 +752,9 @@ def append_or_create_parquet_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.
|
||||
|
||||
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.
|
||||
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)
|
||||
|
||||
logging.info("write info")
|
||||
logger.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)
|
||||
|
||||
logging.info("write stats")
|
||||
logger.info("write stats")
|
||||
aggr_meta.stats = aggregate_stats([m.stats for m in all_metadata])
|
||||
write_stats(aggr_meta.stats, aggr_meta.root)
|
||||
|
||||
@@ -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)
|
||||
self.episodes = load_episodes(self.root)
|
||||
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.stats = load_stats(self.root)
|
||||
|
||||
def ensure_readable(self) -> None:
|
||||
|
||||
@@ -384,7 +384,12 @@ class LiberoEnv(gym.Env):
|
||||
|
||||
def close(self):
|
||||
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(
|
||||
|
||||
@@ -384,7 +384,9 @@ 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.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 = {
|
||||
cam: spaces.Box(
|
||||
|
||||
@@ -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.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]
|
||||
return padded
|
||||
|
||||
|
||||
@@ -302,6 +302,33 @@ 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,
|
||||
@@ -309,16 +336,19 @@ def reconcile_evo1_processors(
|
||||
) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]:
|
||||
"""Reconcile checkpoint-loaded pipelines with the current EVO1 config.
|
||||
|
||||
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.
|
||||
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.
|
||||
"""
|
||||
# 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,
|
||||
|
||||
@@ -46,6 +46,12 @@ 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")
|
||||
|
||||
@@ -142,11 +142,25 @@ class SOFollower(Robot):
|
||||
range_mins[full_turn_motor] = 0
|
||||
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 = {}
|
||||
for motor, m in self.bus.motors.items():
|
||||
self.calibration[motor] = MotorCalibration(
|
||||
id=m.id,
|
||||
drive_mode=0,
|
||||
drive_mode=drive_modes[motor],
|
||||
homing_offset=homing_offsets[motor],
|
||||
range_min=range_mins[motor],
|
||||
range_max=range_maxes[motor],
|
||||
@@ -180,7 +194,7 @@ class SOFollower(Robot):
|
||||
def get_observation(self) -> RobotObservation:
|
||||
# Read arm position
|
||||
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()}
|
||||
dt_ms = (time.perf_counter() - start) * 1e3
|
||||
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.
|
||||
# /!\ 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")
|
||||
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_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 concurrent.futures
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
@@ -52,6 +53,7 @@ from lerobot.datasets import (
|
||||
get_feature_stats,
|
||||
write_stats,
|
||||
)
|
||||
from lerobot.datasets.compute_stats import sample_indices
|
||||
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
|
||||
|
||||
|
||||
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.
|
||||
|
||||
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
|
||||
@@ -92,16 +96,31 @@ def process_single_episode(dataset: LeRobotDataset, episode_idx: int) -> dict:
|
||||
start_idx = dataset.meta.episodes[episode_idx]["dataset_from_index"]
|
||||
end_idx = dataset.meta.episodes[episode_idx]["dataset_to_index"]
|
||||
|
||||
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
|
||||
episode_len = end_idx - start_idx
|
||||
|
||||
if key not in collected_data:
|
||||
collected_data[key] = []
|
||||
collected_data[key].append(value)
|
||||
# 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] = {}
|
||||
|
||||
# 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 = {}
|
||||
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
|
||||
|
||||
|
||||
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.
|
||||
|
||||
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
|
||||
@@ -153,15 +174,15 @@ def compute_quantile_stats_for_dataset(dataset: LeRobotDataset) -> dict[str, dic
|
||||
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)
|
||||
ep_stats = process_single_episode(dataset, episode_idx, use_sampling)
|
||||
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, 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:
|
||||
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)
|
||||
}
|
||||
|
||||
@@ -188,6 +209,7 @@ 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.
|
||||
|
||||
@@ -195,6 +217,8 @@ 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(
|
||||
@@ -208,7 +232,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)
|
||||
new_stats = compute_quantile_stats_for_dataset(dataset, use_sampling=use_sampling)
|
||||
|
||||
logging.info("Updating dataset metadata with new quantile statistics")
|
||||
dataset.meta.stats = new_stats
|
||||
@@ -248,6 +272,14 @@ 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
|
||||
@@ -258,6 +290,7 @@ def main():
|
||||
repo_id=args.repo_id,
|
||||
root=root,
|
||||
overwrite=args.overwrite,
|
||||
use_sampling=not args.no_sampling,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -564,7 +564,7 @@ def eval_policy(
|
||||
if seeds:
|
||||
all_seeds.extend(seeds)
|
||||
else:
|
||||
all_seeds.append(None)
|
||||
all_seeds.extend([None] * env.num_envs)
|
||||
|
||||
# FIXME: episode_data is either None or it doesn't exist
|
||||
if return_episode_data:
|
||||
|
||||
@@ -22,7 +22,8 @@ import dataclasses
|
||||
import logging
|
||||
import sys
|
||||
import time
|
||||
from contextlib import nullcontext
|
||||
from collections.abc import Iterator
|
||||
from contextlib import contextmanager, nullcontext
|
||||
from pprint import pformat
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
@@ -76,6 +77,20 @@ 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
|
||||
@@ -280,14 +295,6 @@ 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")
|
||||
@@ -695,7 +702,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 torch.no_grad(), accelerator.autocast():
|
||||
with _make_eval_envs(cfg) as eval_env, torch.no_grad(), accelerator.autocast():
|
||||
eval_info = eval_policy_all(
|
||||
envs=eval_env, # dict[suite][task_id] -> vec_env
|
||||
policy=accelerator.unwrap_model(policy),
|
||||
@@ -743,9 +750,6 @@ 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:
|
||||
|
||||
@@ -29,6 +29,12 @@ 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")
|
||||
|
||||
@@ -110,11 +110,25 @@ class SOLeader(Teleoperator):
|
||||
range_mins[full_turn_motor] = 0
|
||||
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 = {}
|
||||
for motor, m in self.bus.motors.items():
|
||||
self.calibration[motor] = MotorCalibration(
|
||||
id=m.id,
|
||||
drive_mode=0,
|
||||
drive_mode=drive_modes[motor],
|
||||
homing_offset=homing_offsets[motor],
|
||||
range_min=range_mins[motor],
|
||||
range_max=range_maxes[motor],
|
||||
@@ -145,7 +159,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")
|
||||
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()}
|
||||
dt_ms = (time.perf_counter() - start) * 1e3
|
||||
logger.debug(f"{self} read action: {dt_ms:.1f}ms")
|
||||
|
||||
@@ -41,7 +41,7 @@ class RandomSubsetApply(Transform):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
transforms: Sequence[Callable],
|
||||
transforms: Sequence[Callable[..., Any]],
|
||||
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] * len(transforms)
|
||||
p = [1.0] * 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 = None
|
||||
self.selected_transforms: list[Callable[..., Any]] = []
|
||||
|
||||
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):
|
||||
def _check_input(self, sharpness: float | Sequence[float]) -> tuple[float, float]:
|
||||
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):
|
||||
def make_transform_from_config(cfg: ImageTransformConfig) -> Transform:
|
||||
if cfg.type == "SharpnessJitter":
|
||||
return SharpnessJitter(**cfg.kwargs)
|
||||
|
||||
@@ -236,8 +236,8 @@ class ImageTransforms(Transform):
|
||||
super().__init__()
|
||||
self._cfg = cfg
|
||||
|
||||
self.weights = []
|
||||
self.transforms = {}
|
||||
self.weights: list[float] = []
|
||||
self.transforms: dict[str, Transform] = {}
|
||||
for tf_name, tf_cfg in cfg.tfs.items():
|
||||
if tf_cfg.weight <= 0.0:
|
||||
continue
|
||||
|
||||
@@ -133,10 +133,13 @@ def say(text: str, blocking: bool = False):
|
||||
else:
|
||||
raise RuntimeError("Unsupported operating system for text-to-speech.")
|
||||
|
||||
if blocking:
|
||||
subprocess.run(cmd, check=True)
|
||||
else:
|
||||
subprocess.Popen(cmd, creationflags=subprocess.CREATE_NO_WINDOW if system == "Windows" else 0)
|
||||
try:
|
||||
if blocking:
|
||||
subprocess.run(cmd, check=True, timeout=5)
|
||||
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):
|
||||
|
||||
@@ -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)
|
||||
@@ -482,6 +482,20 @@ 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"
|
||||
|
||||
@@ -294,6 +294,19 @@ 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})
|
||||
|
||||
@@ -496,6 +496,60 @@ 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())
|
||||
|
||||
@@ -12,7 +12,6 @@
|
||||
# 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
|
||||
|
||||
@@ -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)
|
||||
adapted_policy = torch.nn.Linear(1, 1)
|
||||
peft_model_from_pretrained = MagicMock(return_value=adapted_policy)
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"peft",
|
||||
SimpleNamespace(
|
||||
PeftConfig=SimpleNamespace(from_pretrained=peft_config_from_pretrained),
|
||||
PeftModel=SimpleNamespace(from_pretrained=peft_model_from_pretrained),
|
||||
),
|
||||
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),
|
||||
)
|
||||
|
||||
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",
|
||||
|
||||
@@ -49,7 +49,7 @@ def _make_bus_mock() -> MagicMock:
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def follower():
|
||||
def follower(tmp_path):
|
||||
bus_mock = _make_bus_mock()
|
||||
|
||||
def _bus_side_effect(*_args, **kwargs):
|
||||
@@ -71,7 +71,7 @@ def follower():
|
||||
),
|
||||
patch.object(SO100Follower, "configure", lambda self: None),
|
||||
):
|
||||
cfg = SO100FollowerConfig(port="/dev/null")
|
||||
cfg = SO100FollowerConfig(port="/dev/null", calibration_dir=tmp_path)
|
||||
robot = SO100Follower(cfg)
|
||||
yield robot
|
||||
if robot.is_connected:
|
||||
@@ -99,6 +99,27 @@ 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()
|
||||
|
||||
@@ -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("I_Coefficient", "shoulder_pan", 1)
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user