mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-07 18:11:50 +00:00
Compare commits
201 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| aa517b5780 | |||
| 2c17433f4d | |||
| d31283cc5d | |||
| d8feb22f93 | |||
| 9073d64050 | |||
| d6ce7bd330 | |||
| 43eb0e375f | |||
| fa8be1c4fe | |||
| a7d1179aab | |||
| 772da63a8e | |||
| 27a229ea64 | |||
| 0d1e57f032 | |||
| e2eff72ec0 | |||
| d0d9036304 | |||
| 5d79869934 | |||
| 7403060ad6 | |||
| 89aaaf1556 | |||
| c7e976d812 | |||
| 1a4a47e804 | |||
| 4fe5c3ab70 | |||
| cf7946e602 | |||
| b12a386334 | |||
| cee5a3fec5 | |||
| 8fb18109ef | |||
| 170d8be7c2 | |||
| 8063cd5ed3 | |||
| 03891f66da | |||
| 99213daa3e | |||
| 42cffd6f2e | |||
| 7b1b37b696 | |||
| 4382742681 | |||
| e8d79b5191 | |||
| 066308ceb8 | |||
| 40e9ddd1ed | |||
| c69f23723e | |||
| 50293bb17b | |||
| 839ac5f2aa | |||
| 0479eb8f69 | |||
| a877c596ba | |||
| 1ccdf365d2 | |||
| 6bdcd460e0 | |||
| 2005a28a00 | |||
| 58b91dc886 | |||
| 8e0f5cd052 | |||
| f51272362c | |||
| cd0098a5f7 | |||
| efde42d4a9 | |||
| aeb70812c1 | |||
| 376a6457cf | |||
| a2489ab0da | |||
| 014486999e | |||
| cda44e5a52 | |||
| 9183083e75 | |||
| 6745958362 | |||
| 51588f741b | |||
| df4292f6ed | |||
| 7e30090e97 | |||
| e881fb6678 | |||
| acf0ba7fb3 | |||
| a74b90edd1 | |||
| 846677f9cc | |||
| af9ddcf9a2 | |||
| d32006440c | |||
| f1cfdfced9 | |||
| 888a5b6249 | |||
| f247aa0701 | |||
| 1ac6a6d3fe | |||
| e698c709d8 | |||
| a988da4789 | |||
| 99963b6968 | |||
| 332ca4ccc5 | |||
| fc43246942 | |||
| 793ad86fc9 | |||
| a6dbb65917 | |||
| 6c7169c4af | |||
| f125d5e3bf | |||
| 75dcfd4886 | |||
| ff3cbaa872 | |||
| ce793cde64 | |||
| 029c4a9a76 | |||
| d893bf1e30 | |||
| 8c796b39f5 | |||
| 4ebe482a7e | |||
| 2fcc358e98 | |||
| b052843f08 | |||
| ebb464c255 | |||
| 2914ae2a96 | |||
| 645c87e3a9 | |||
| 2c802ac134 | |||
| 15ffc01fb3 | |||
| a837685bf8 | |||
| d32b76cc66 | |||
| 08fb310eaa | |||
| 574a708950 | |||
| ce665160ae | |||
| 35c5d43255 | |||
| 95c1e32aa5 | |||
| e4db65a127 | |||
| 0053defa2e | |||
| fd5d8b3d5f | |||
| 5bf82f8229 | |||
| 5ca3920611 | |||
| 8bde9d0ab7 | |||
| abcbc16126 | |||
| e4fd30a8d4 | |||
| 5f759b1637 | |||
| 6a75b4761a | |||
| e5ade5565d | |||
| 0524551f52 | |||
| 862bc7ef85 | |||
| d38792d6e5 | |||
| db3cf0158c | |||
| 0535f2a59a | |||
| 2805ae347c | |||
| 28ef6fcd14 | |||
| 7fc7ec75bb | |||
| 87890cbf38 | |||
| 5326ffe77e | |||
| a1734cf575 | |||
| 82f300e880 | |||
| 3e7c9d7afc | |||
| e9cb779eab | |||
| 8ff95be04c | |||
| f02ce69df0 | |||
| 1feb7b5d88 | |||
| fbe9009db2 | |||
| c0013b130b | |||
| c4763f61a1 | |||
| b95c219d96 | |||
| 9b1138171e | |||
| 023b8f3466 | |||
| 1cad87ebd2 | |||
| 99de7567e6 | |||
| 21baa8fa02 | |||
| 8b4a5368b3 | |||
| f5c6b03b61 | |||
| e7be2fd113 | |||
| b632490b4b | |||
| 9a9c7208d2 | |||
| 427b97d198 | |||
| 2c2bb1e8bf | |||
| 4b24f94225 | |||
| 670a278cbc | |||
| fc74001202 | |||
| f14ac5d486 | |||
| 7bd0d62ce5 | |||
| 7eccefe235 | |||
| b72274066e | |||
| 20f2910b63 | |||
| fd4ae3466b | |||
| 7beb040e8e | |||
| 05bd18f453 | |||
| 8077456c00 | |||
| 5595887fd0 | |||
| 41959389b6 | |||
| 2c4e888c7f | |||
| 5ced72e6b8 | |||
| 907023f9f7 | |||
| 4ba23ea029 | |||
| 409ac0baca | |||
| 699363f9fc | |||
| ae7a54de57 | |||
| fb9139b882 | |||
| 9fe3a3fb17 | |||
| 26cb9a24c3 | |||
| 77106697c3 | |||
| 75bc44c166 | |||
| f2b79656eb | |||
| 14c2ece004 | |||
| 35612c61e1 | |||
| f7bb3e2d90 | |||
| 1e0d667a22 | |||
| 33969a0337 | |||
| fa26290e8c | |||
| e9f7f5127b | |||
| 097842c70f | |||
| 3b8a3a32a0 | |||
| 1c56779dd9 | |||
| 83a4338f8b | |||
| 730c7b2f35 | |||
| 116059a43e | |||
| b08149a113 | |||
| c227107f60 | |||
| 01dc289f3d | |||
| 6830ca7645 | |||
| ed42c71fc3 | |||
| e0139065bd | |||
| e509f255af | |||
| e2fcd140b0 | |||
| 2a7a0e6129 | |||
| 9f33791b19 | |||
| 453e0a995f | |||
| 8ebf79c494 | |||
| 8774aec304 | |||
| ac742c9f0d | |||
| cd13f1ecfd | |||
| 9aa632968f | |||
| 62caaf07b0 | |||
| 3355f04ca6 | |||
| 769f531603 | |||
| f6c7287ae7 |
@@ -25,7 +25,7 @@ body:
|
||||
id: system-info
|
||||
attributes:
|
||||
label: System Info
|
||||
description: Please share your LeRobot configuration by running `lerobot-info` (if installed) or `python -m lerobot.scripts.display_sys_info` (if not installed) and pasting the output below.
|
||||
description: If needed, you can share your lerobot configuration with us by running `python -m lerobot.scripts.display_sys_info` and copy-pasting its outputs below
|
||||
render: Shell
|
||||
placeholder: lerobot version, OS, python version, numpy version, torch version, and lerobot's configuration
|
||||
validations:
|
||||
|
||||
@@ -1,68 +0,0 @@
|
||||
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
# This workflow handles closing stale issues and PRs.
|
||||
name: Stale
|
||||
on:
|
||||
# Allows running this workflow manually from the Actions tab
|
||||
workflow_dispatch:
|
||||
|
||||
# Runs at 02:00
|
||||
schedule:
|
||||
- cron: "0 2 * * *"
|
||||
|
||||
env:
|
||||
CLOSE_ISSUE_MESSAGE: >
|
||||
This issue was closed because it has been stalled for 14 days with no activity.
|
||||
Feel free to reopen if is still relevant, or to ping a collaborator if you have any questions.
|
||||
CLOSE_PR_MESSAGE: >
|
||||
This PR was closed because it has been stalled for 14 days with no activity.
|
||||
Feel free to reopen if is still relevant, or to ping a collaborator if you have any questions.
|
||||
WARN_ISSUE_MESSAGE: >
|
||||
This issue has been automatically marked as stale because it has not had
|
||||
recent activity (1 year). It will be closed if no further activity occurs.
|
||||
Thank you for your contributions.
|
||||
WARN_PR_MESSAGE: >
|
||||
This PR has been automatically marked as stale because it has not had
|
||||
recent activity (1 year). It will be closed if no further activity occurs.
|
||||
Thank you for your contributions.
|
||||
|
||||
jobs:
|
||||
# This job runs the actions/stale action to close stale issues and PRs.
|
||||
stale:
|
||||
name: Close Stale Issues and PRs
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
actions: write
|
||||
contents: write # only for delete-branch option
|
||||
issues: write
|
||||
pull-requests: write
|
||||
steps:
|
||||
- uses: actions/stale@v10
|
||||
with:
|
||||
repo-token: ${{ secrets.GITHUB_TOKEN }}
|
||||
stale-issue-label: stale
|
||||
stale-pr-label: stale
|
||||
exempt-issue-labels: never-stale
|
||||
exempt-pr-labels: never-stale
|
||||
days-before-issue-stale: 180 # TODO(Steven): Will modify this to 90 after initial cleanup
|
||||
days-before-issue-close: 14
|
||||
days-before-pr-stale: 180
|
||||
days-before-pr-close: 14
|
||||
delete-branch: true
|
||||
close-issue-message: ${{ env.CLOSE_ISSUE_MESSAGE }}
|
||||
close-pr-message: ${{ env.CLOSE_PR_MESSAGE }}
|
||||
stale-issue-message: ${{ env.WARN_ISSUE_MESSAGE }}
|
||||
stale-pr-message: ${{ env.WARN_PR_MESSAGE }}
|
||||
operations-per-run: 500
|
||||
@@ -202,7 +202,7 @@ Check out [example 1](https://github.com/huggingface/lerobot/blob/main/examples/
|
||||
You can also locally visualize episodes from a dataset on the hub by executing our script from the command line:
|
||||
|
||||
```bash
|
||||
lerobot-dataset-viz \
|
||||
python -m lerobot.scripts.visualize_dataset \
|
||||
--repo-id lerobot/pusht \
|
||||
--episode-index 0
|
||||
```
|
||||
@@ -210,7 +210,7 @@ lerobot-dataset-viz \
|
||||
or from a dataset in a local folder with the `root` option and the `--local-files-only` (in the following case the dataset will be searched for in `./my_local_data_dir/lerobot/pusht`)
|
||||
|
||||
```bash
|
||||
lerobot-dataset-viz \
|
||||
python -m lerobot.scripts.visualize_dataset \
|
||||
--repo-id lerobot/pusht \
|
||||
--root ./my_local_data_dir \
|
||||
--local-files-only 1 \
|
||||
@@ -221,19 +221,19 @@ It will open `rerun.io` and display the camera streams, robot states and actions
|
||||
|
||||
https://github-production-user-asset-6210df.s3.amazonaws.com/4681518/328035972-fd46b787-b532-47e2-bb6f-fd536a55a7ed.mov?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAVCODYLSA53PQK4ZA%2F20240505%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20240505T172924Z&X-Amz-Expires=300&X-Amz-Signature=d680b26c532eeaf80740f08af3320d22ad0b8a4e4da1bcc4f33142c15b509eda&X-Amz-SignedHeaders=host&actor_id=24889239&key_id=0&repo_id=748713144
|
||||
|
||||
Our script can also visualize datasets stored on a distant server. See `lerobot-dataset-viz --help` for more instructions.
|
||||
Our script can also visualize datasets stored on a distant server. See `python -m lerobot.scripts.visualize_dataset --help` for more instructions.
|
||||
|
||||
### The `LeRobotDataset` format
|
||||
|
||||
A dataset in `LeRobotDataset` format is very simple to use. It can be loaded from a repository on the Hugging Face hub or a local folder simply with e.g. `dataset = LeRobotDataset("lerobot/aloha_static_coffee")` and can be indexed into like any Hugging Face and PyTorch dataset. For instance `dataset[0]` will retrieve a single temporal frame from the dataset containing observation(s) and an action as PyTorch tensors ready to be fed to a model.
|
||||
|
||||
A specificity of `LeRobotDataset` is that, rather than retrieving a single frame by its index, we can retrieve several frames based on their temporal relationship with the indexed frame, by setting `delta_timestamps` to a list of relative times with respect to the indexed frame. For example, with `delta_timestamps = {"observation.image": [-1, -0.5, -0.2, 0]}` one can retrieve, for a given index, 4 frames: 3 "previous" frames 1 second, 0.5 seconds, and 0.2 seconds before the indexed frame, and the indexed frame itself (corresponding to the 0 entry). See example [1_load_lerobot_dataset.py](https://github.com/huggingface/lerobot/blob/main/examples/dataset/load_lerobot_dataset.py) for more details on `delta_timestamps`.
|
||||
A specificity of `LeRobotDataset` is that, rather than retrieving a single frame by its index, we can retrieve several frames based on their temporal relationship with the indexed frame, by setting `delta_timestamps` to a list of relative times with respect to the indexed frame. For example, with `delta_timestamps = {"observation.image": [-1, -0.5, -0.2, 0]}` one can retrieve, for a given index, 4 frames: 3 "previous" frames 1 second, 0.5 seconds, and 0.2 seconds before the indexed frame, and the indexed frame itself (corresponding to the 0 entry). See example [1_load_lerobot_dataset.py](https://github.com/huggingface/lerobot/blob/main/examples/1_load_lerobot_dataset.py) for more details on `delta_timestamps`.
|
||||
|
||||
Under the hood, the `LeRobotDataset` format makes use of several ways to serialize data which can be useful to understand if you plan to work more closely with this format. We tried to make a flexible yet simple dataset format that would cover most type of features and specificities present in reinforcement learning and robotics, in simulation and in real-world, with a focus on cameras and robot states but easily extended to other types of sensory inputs as long as they can be represented by a tensor.
|
||||
|
||||
Here are the important details and internal structure organization of a typical `LeRobotDataset` instantiated with `dataset = LeRobotDataset("lerobot/aloha_static_coffee")`. The exact features will change from dataset to dataset but not the main aspects:
|
||||
|
||||
```
|
||||
````
|
||||
dataset attributes:
|
||||
├ hf_dataset: a Hugging Face dataset (backed by Arrow/parquet). Typical features example:
|
||||
│ ├ observation.images.cam_high (VideoFrame):
|
||||
@@ -269,7 +269,7 @@ dataset attributes:
|
||||
├ root (Path): local directory where the dataset is stored
|
||||
├ image_transforms (Callable): optional image transformations to apply to visual modalities
|
||||
└ delta_timestamps (dict): optional delta timestamps for temporal queries
|
||||
```
|
||||
decoding videos (e.g., 'pyav', 'torchcodec')
|
||||
|
||||
A `LeRobotDataset` is serialised using several widespread file formats for each of its parts, namely:
|
||||
|
||||
@@ -279,6 +279,42 @@ A `LeRobotDataset` is serialised using several widespread file formats for each
|
||||
|
||||
Dataset can be uploaded/downloaded from the HuggingFace hub seamlessly. To work on a local dataset, you can specify its location with the `root` argument if it's not in the default `~/.cache/huggingface/lerobot` location.
|
||||
|
||||
### Evaluate a pretrained policy
|
||||
|
||||
Check out [example 2](https://github.com/huggingface/lerobot/blob/main/examples/2_evaluate_pretrained_policy.py) that illustrates how to download a pretrained policy from Hugging Face hub, and run an evaluation on its corresponding environment.
|
||||
|
||||
We also provide a more capable script to parallelize the evaluation over multiple environments during the same rollout. Here is an example with a pretrained model hosted on [lerobot/diffusion_pusht](https://huggingface.co/lerobot/diffusion_pusht):
|
||||
|
||||
```bash
|
||||
lerobot-eval \
|
||||
--policy.path=lerobot/diffusion_pusht \
|
||||
--env.type=pusht \
|
||||
--eval.batch_size=10 \
|
||||
--eval.n_episodes=10 \
|
||||
--policy.use_amp=false \
|
||||
--policy.device=cuda
|
||||
````
|
||||
|
||||
Note: After training your own policy, you can re-evaluate the checkpoints with:
|
||||
|
||||
```bash
|
||||
lerobot-eval --policy.path={OUTPUT_DIR}/checkpoints/last/pretrained_model
|
||||
```
|
||||
|
||||
See `lerobot-eval --help` for more instructions.
|
||||
|
||||
### Train your own policy
|
||||
|
||||
Check out [example 3](https://github.com/huggingface/lerobot/blob/main/examples/3_train_policy.py) that illustrates how to train a model using our core library in python, and [example 4](https://github.com/huggingface/lerobot/blob/main/examples/4_train_policy_with_script.md) that shows how to use our training script from command line.
|
||||
|
||||
To use wandb for logging training and evaluation curves, make sure you've run `wandb login` as a one-time setup step. Then, when running the training command above, enable WandB in the configuration by adding `--wandb.enable=true`.
|
||||
|
||||
A link to the wandb logs for the run will also show up in yellow in your terminal. Here is an example of what they look like in your browser. Please also check [here](https://github.com/huggingface/lerobot/blob/main/examples/4_train_policy_with_script.md#typical-logs-and-metrics) for the explanation of some commonly used metrics in logs.
|
||||
|
||||
\<img src="https://raw.githubusercontent.com/huggingface/lerobot/main/media/wandb.png" alt="WandB logs example"\>
|
||||
|
||||
Note: For efficiency, during training every checkpoint is evaluated on a low number of episodes. You may use `--eval.n_episodes=500` to evaluate on more episodes than the default. Or, after training, you may want to re-evaluate your best checkpoints on more episodes or change the evaluation settings. See `lerobot-eval --help` for more instructions.
|
||||
|
||||
#### Reproduce state-of-the-art (SOTA)
|
||||
|
||||
We provide some pretrained policies on our [hub page](https://huggingface.co/lerobot) that can achieve state-of-the-art performances.
|
||||
@@ -337,7 +373,3 @@ If you want, you can cite this work with:
|
||||
## Star History
|
||||
|
||||
[](https://star-history.com/#huggingface/lerobot&Timeline)
|
||||
|
||||
```
|
||||
|
||||
```
|
||||
|
||||
@@ -1,378 +0,0 @@
|
||||
"""
|
||||
Benchmark memory footprint and inference latency of a policy on arbitrary devices.
|
||||
|
||||
This script loads a pretrained policy directly (similar to the async inference server)
|
||||
and generates dummy input data based on the policy's input_features to perform
|
||||
accurate benchmarking without requiring datasets.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import signal
|
||||
import statistics
|
||||
from contextlib import contextmanager
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
import psutil
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
|
||||
from lerobot.configs.types import FeatureType
|
||||
from lerobot.policies.factory import get_policy_class
|
||||
from lerobot.policies.pretrained import PreTrainedPolicy
|
||||
|
||||
|
||||
class TimeoutException:
|
||||
pass
|
||||
|
||||
|
||||
@contextmanager
|
||||
def timeout(seconds):
|
||||
def signal_handler(signum, frame):
|
||||
raise TimeoutException(f"Timed out after {seconds} seconds")
|
||||
|
||||
# On Windows, signal is not available, so we can't use this timeout mechanism
|
||||
if not hasattr(signal, "SIGALRM"):
|
||||
yield
|
||||
return
|
||||
|
||||
old_handler = signal.signal(signal.SIGALRM, signal_handler)
|
||||
try:
|
||||
# signal.alarm expects integer seconds
|
||||
# for float seconds, we can use setitimer
|
||||
signal.setitimer(signal.ITIMER_REAL, seconds)
|
||||
yield
|
||||
finally:
|
||||
signal.setitimer(signal.ITIMER_REAL, 0)
|
||||
signal.signal(signal.SIGALRM, old_handler)
|
||||
|
||||
|
||||
def bytes_to_human(n: int) -> str:
|
||||
for unit in ["B", "KB", "MB", "GB", "TB"]:
|
||||
if n < 1024:
|
||||
return f"{n:.2f} {unit}"
|
||||
n /= 1024
|
||||
return f"{n:.2f} PB"
|
||||
|
||||
|
||||
def percentile(values: list[float], p: float) -> float:
|
||||
if not values:
|
||||
return float("nan")
|
||||
k = (len(values) - 1) * (p / 100.0)
|
||||
f = int(k)
|
||||
c = min(f + 1, len(values) - 1)
|
||||
if f == c:
|
||||
return values[f]
|
||||
return values[f] + (values[c] - values[f]) * (k - f)
|
||||
|
||||
|
||||
def generate_dummy_observation(input_features: dict, device: str = "cpu") -> dict:
|
||||
"""Generate dummy observation data based on policy input features."""
|
||||
dummy_obs = {}
|
||||
|
||||
for key, feature in input_features.items():
|
||||
shape = feature.shape
|
||||
|
||||
if feature.type == FeatureType.VISUAL:
|
||||
# Images: random values in [0, 1] range (already normalized)
|
||||
dummy_obs[key] = torch.rand(shape, dtype=torch.float32, device=device)
|
||||
elif feature.type in [FeatureType.STATE, FeatureType.ACTION, FeatureType.ENV]:
|
||||
# State/action/env: random normal distribution
|
||||
dummy_obs[key] = torch.randn(shape, dtype=torch.float32, device=device)
|
||||
else:
|
||||
# Default: random normal for unknown types
|
||||
dummy_obs[key] = torch.randn(shape, dtype=torch.float32, device=device)
|
||||
|
||||
# Add batch dimension
|
||||
for key in dummy_obs:
|
||||
dummy_obs[key] = dummy_obs[key].unsqueeze(0)
|
||||
|
||||
# Add task string for language-conditioned policies
|
||||
dummy_obs["task"] = ""
|
||||
|
||||
return dummy_obs
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Policy inference benchmark")
|
||||
parser.add_argument(
|
||||
"--policy-id", type=str, required=True, help="Model ID or local path to pretrained policy"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--policy-type", type=str, required=True, help="Type of policy (smolvla, act, diffusion, etc.)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--device", type=str, default="mps", choices=["cuda", "cpu", "mps"], help="Device to run on"
|
||||
)
|
||||
parser.add_argument("--seed", type=int, default=42, help="Random seed")
|
||||
parser.add_argument(
|
||||
"--num-samples", type=int, default=100, help="Number of inference samples to benchmark"
|
||||
)
|
||||
parser.add_argument("--warmup", type=int, default=10, help="Number of warmup samples (not timed)")
|
||||
parser.add_argument(
|
||||
"--output-dir", type=str, default="outputs/benchmarks", help="Directory to save benchmark results"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--timeout",
|
||||
type=float,
|
||||
default=0.3,
|
||||
help="Timeout for each inference pass in seconds (default: 0.3s = 300ms)",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
# Seed & deterministic-ish setup
|
||||
torch.manual_seed(args.seed)
|
||||
if args.device == "cuda":
|
||||
torch.cuda.manual_seed_all(args.seed)
|
||||
torch.backends.cudnn.benchmark = False
|
||||
torch.backends.cudnn.deterministic = False # leave False to avoid perf cliffs
|
||||
|
||||
# Resolve device availability
|
||||
device = args.device.lower()
|
||||
if device == "cuda" and not torch.cuda.is_available():
|
||||
print("[!] CUDA requested but unavailable. Falling back to CPU.")
|
||||
device = "cpu"
|
||||
elif device == "mps" and not (hasattr(torch.backends, "mps") and torch.backends.mps.is_available()):
|
||||
print("[!] MPS requested but unavailable. Falling back to CPU.")
|
||||
device = "cpu"
|
||||
|
||||
use_cuda = device == "cuda"
|
||||
|
||||
# Create output directory and log file
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
policy_name = args.policy_id.replace("/", "_").replace("\\", "_")
|
||||
log_file = output_dir / f"benchmark_{args.policy_type}_{policy_name}_{device}_{timestamp}.txt"
|
||||
|
||||
# Load policy directly from pretrained (similar to async inference server)
|
||||
print(f"Loading policy {args.policy_type} from {args.policy_id}...")
|
||||
policy_class = get_policy_class(args.policy_type)
|
||||
policy: PreTrainedPolicy = policy_class.from_pretrained(args.policy_id)
|
||||
policy.eval()
|
||||
policy.to(device)
|
||||
|
||||
print(f"Policy loaded on {device}")
|
||||
print(f"Input features: {list(policy.config.input_features.keys())}")
|
||||
print(f"Output features: {list(policy.config.output_features.keys())}")
|
||||
|
||||
# Generate dummy observation based on policy input features
|
||||
dummy_observation = generate_dummy_observation(policy.config.input_features, device)
|
||||
dummy_observation["task"] = ""
|
||||
|
||||
# Helper to sync for fair timings
|
||||
def _sync(dev_=device):
|
||||
if dev_ == "cuda" and torch.cuda.is_available():
|
||||
torch.cuda.synchronize()
|
||||
elif dev_ == "mps" and hasattr(torch, "mps"):
|
||||
try:
|
||||
torch.mps.synchronize()
|
||||
except AttributeError:
|
||||
pass # MPS sync not available in this PyTorch version
|
||||
|
||||
# Warmup (to stabilize kernels/caches)
|
||||
print("Warming up...")
|
||||
with torch.no_grad():
|
||||
policy.reset()
|
||||
for _ in range(args.warmup):
|
||||
_ = policy.select_action(dummy_observation)
|
||||
_sync()
|
||||
|
||||
# Memory footprint before timing
|
||||
process = psutil.Process(os.getpid())
|
||||
rss_before = process.memory_info().rss
|
||||
if use_cuda:
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
|
||||
# PyTorch timing with Event objects for more accurate GPU timing
|
||||
print(f"Running benchmark: {args.num_samples} samples...")
|
||||
|
||||
if use_cuda:
|
||||
# Use CUDA Events for precise GPU timing
|
||||
start_events = []
|
||||
end_events = []
|
||||
timeout_count = 0
|
||||
|
||||
with torch.no_grad():
|
||||
for forward in tqdm(range(args.num_samples), desc="Trials"):
|
||||
start_event = torch.cuda.Event(enable_timing=True)
|
||||
end_event = torch.cuda.Event(enable_timing=True)
|
||||
try:
|
||||
with timeout(args.timeout):
|
||||
start_event.record()
|
||||
_ = policy.select_action(dummy_observation)
|
||||
end_event.record()
|
||||
|
||||
start_events.append(start_event)
|
||||
end_events.append(end_event)
|
||||
except TimeoutException:
|
||||
timeout_count += 1
|
||||
# Add placeholder for timeout
|
||||
start_events.append(None)
|
||||
end_events.append(None)
|
||||
print(f"\n[!] Timeout on forward {forward + 1}")
|
||||
continue
|
||||
|
||||
# Synchronize and collect timing results
|
||||
torch.cuda.synchronize()
|
||||
per_forward_ms = []
|
||||
for start_event, end_event in zip(start_events, end_events, strict=True):
|
||||
if start_event is None:
|
||||
per_forward_ms.append(args.timeout * 1000)
|
||||
else:
|
||||
per_forward_ms.append(start_event.elapsed_time(end_event))
|
||||
|
||||
if timeout_count > 0:
|
||||
print(f"[!] {timeout_count} inference passes timed out (>{args.timeout * 1000:.1f}ms)")
|
||||
|
||||
else:
|
||||
# Use simple time.perf_counter for CPU/MPS timing with timeout
|
||||
import time
|
||||
|
||||
per_forward_ms = []
|
||||
timeout_count = 0
|
||||
|
||||
with torch.no_grad():
|
||||
for sample in tqdm(range(args.num_samples), desc="Samples"):
|
||||
try:
|
||||
with timeout(args.timeout):
|
||||
start_time = time.perf_counter()
|
||||
_ = policy.select_action(dummy_observation)
|
||||
end_time = time.perf_counter()
|
||||
|
||||
per_forward_ms.append((end_time - start_time) * 1000) # Convert to ms
|
||||
except TimeoutException:
|
||||
timeout_count += 1
|
||||
per_forward_ms.append(args.timeout * 1000)
|
||||
print(f"\n[!] Timeout on sample {sample + 1}")
|
||||
continue
|
||||
|
||||
if timeout_count > 0:
|
||||
print(f"[!] {timeout_count} inference passes timed out (>{args.timeout * 1000:.1f}ms)")
|
||||
|
||||
# Memory footprint after timing
|
||||
rss_after = process.memory_info().rss
|
||||
rss_delta = rss_after - rss_before
|
||||
cuda_peak = torch.cuda.max_memory_allocated() if use_cuda else 0
|
||||
|
||||
# Sort timing results for percentile calculations
|
||||
per_forward_ms_sorted = sorted(per_forward_ms)
|
||||
|
||||
mean_ms = statistics.fmean(per_forward_ms) if per_forward_ms else float("nan")
|
||||
std_ms = statistics.pstdev(per_forward_ms) if len(per_forward_ms) > 1 else 0.0
|
||||
min_ms = per_forward_ms_sorted[0] if per_forward_ms_sorted else float("nan")
|
||||
max_ms = per_forward_ms_sorted[-1] if per_forward_ms_sorted else float("nan")
|
||||
p50_ms = percentile(per_forward_ms_sorted, 50)
|
||||
p95_ms = percentile(per_forward_ms_sorted, 95)
|
||||
|
||||
# Model size
|
||||
num_params = sum(p.numel() for p in policy.parameters())
|
||||
|
||||
# Prepare results for logging
|
||||
results = {
|
||||
"timestamp": datetime.now().isoformat(),
|
||||
"policy_type": args.policy_type,
|
||||
"policy_id": args.policy_id,
|
||||
"device": device,
|
||||
"num_trials": args.num_samples,
|
||||
"forwards_per_trial": 1,
|
||||
"warmup": args.warmup,
|
||||
"timeout_ms": args.timeout * 1000,
|
||||
"seed": args.seed,
|
||||
"num_params": num_params,
|
||||
"timeout_count": timeout_count,
|
||||
"latency_mean_ms": mean_ms,
|
||||
"latency_std_ms": std_ms,
|
||||
"latency_min_ms": min_ms,
|
||||
"latency_max_ms": max_ms,
|
||||
"latency_p50_ms": p50_ms,
|
||||
"latency_p95_ms": p95_ms,
|
||||
"cpu_rss_before": rss_before,
|
||||
"cpu_rss_after": rss_after,
|
||||
"cpu_rss_delta": rss_delta,
|
||||
"cuda_peak_alloc": cuda_peak,
|
||||
"input_features": list(policy.config.input_features.keys()),
|
||||
"output_features": list(policy.config.output_features.keys()),
|
||||
}
|
||||
|
||||
# Format and write results to log file
|
||||
log_content = f"""
|
||||
=== LeRobot Policy Inference Benchmark ===
|
||||
Timestamp: {results["timestamp"]}
|
||||
Policy: {results["policy_type"]} ({results["policy_id"]})
|
||||
Device: {results["device"]}
|
||||
Seed: {results["seed"]}
|
||||
|
||||
=== Model Information ===
|
||||
Parameters: {results["num_params"]:,}
|
||||
Input Features: {", ".join(results["input_features"])}
|
||||
Output Features: {", ".join(results["output_features"])}
|
||||
|
||||
=== Benchmark Configuration ===
|
||||
Samples: {results["num_trials"]}
|
||||
Warmup: {results["warmup"]}
|
||||
Total Measurements: {len(per_forward_ms)}
|
||||
Timeout: {results["timeout_ms"]:.1f}ms
|
||||
Timeouts: {results["timeout_count"]} / {results["num_trials"]}
|
||||
|
||||
=== Latency Results (ms) ===
|
||||
Mean: {results["latency_mean_ms"]:.3f}
|
||||
Std Dev: {results["latency_std_ms"]:.3f}
|
||||
Min: {results["latency_min_ms"]:.3f}
|
||||
Max: {results["latency_max_ms"]:.3f}
|
||||
P50: {results["latency_p50_ms"]:.3f}
|
||||
P95: {results["latency_p95_ms"]:.3f}
|
||||
|
||||
=== Memory Footprint ===
|
||||
CPU RSS Before: {bytes_to_human(results["cpu_rss_before"])}
|
||||
CPU RSS After: {bytes_to_human(results["cpu_rss_after"])} (Δ {bytes_to_human(results["cpu_rss_delta"])})
|
||||
"""
|
||||
|
||||
if use_cuda:
|
||||
log_content += f"CUDA Peak: {bytes_to_human(results['cuda_peak_alloc'])} (reset before timing)\n"
|
||||
|
||||
log_content += f"""
|
||||
=== Raw Timing Data (first 20 measurements, ms) ===
|
||||
{", ".join(f"{t:.3f}" for t in per_forward_ms[:20])}
|
||||
{"..." if len(per_forward_ms) > 20 else ""}
|
||||
|
||||
=== Summary Statistics ===
|
||||
Timing Method: {"CUDA Events" if use_cuda else "torch.utils.benchmark.Timer"}
|
||||
Device Available: {torch.cuda.is_available() if device == "cuda" else torch.backends.mps.is_available() if device == "mps" else True}
|
||||
PyTorch Version: {torch.__version__}
|
||||
|
||||
Benchmark completed successfully at {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
|
||||
"""
|
||||
|
||||
# Write to log file
|
||||
with open(log_file, "w") as f:
|
||||
f.write(log_content)
|
||||
|
||||
# Print to console (shorter version)
|
||||
print("\n=== Inference Benchmark Results ===")
|
||||
print(f"Policy: {args.policy_type} ({args.policy_id})")
|
||||
print(f"Device: {device}")
|
||||
print(f"Samples: {args.num_samples} | Warmup: {args.warmup}")
|
||||
print(f"Model params: {num_params:,}")
|
||||
|
||||
print("\nLatency per forward (ms):")
|
||||
print(f" mean: {mean_ms:.3f} std: {std_ms:.3f}")
|
||||
print(f" min: {min_ms:.3f} max: {max_ms:.3f}")
|
||||
print(f" p50: {p50_ms:.3f} p95: {p95_ms:.3f}")
|
||||
|
||||
print("\nMemory footprint:")
|
||||
print(f" CPU RSS before: {bytes_to_human(rss_before)}")
|
||||
print(f" CPU RSS after : {bytes_to_human(rss_after)} (Δ {bytes_to_human(rss_delta)})")
|
||||
if use_cuda:
|
||||
print(
|
||||
f" CUDA peak allocated: {bytes_to_human(cuda_peak)} "
|
||||
f"(reset by reset_peak_memory_stats before timing)"
|
||||
)
|
||||
|
||||
print(f"\nResults saved to: {log_file}")
|
||||
print("Benchmark completed successfully!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -23,14 +23,14 @@
|
||||
- sections:
|
||||
- local: lerobot-dataset-v3
|
||||
title: Using LeRobotDataset
|
||||
- local: libero
|
||||
title: Using Libero
|
||||
- local: porting_datasets_v3
|
||||
title: Porting Large Datasets
|
||||
title: "Datasets"
|
||||
- sections:
|
||||
- local: smolvla
|
||||
title: Finetune SmolVLA
|
||||
- local: libero
|
||||
title: Using Libero
|
||||
title: "Policies"
|
||||
|
||||
- sections:
|
||||
|
||||
+15
-14
@@ -62,7 +62,7 @@ pip install -e ".[hilserl]"
|
||||
|
||||
### Understanding Configuration
|
||||
|
||||
The training process begins with proper configuration for the HILSerl environment. The main configuration class is `GymManipulatorConfig` in `lerobot/rl/gym_manipulator.py`, which contains nested `HILSerlRobotEnvConfig` and `DatasetConfig`. The configuration is organized into focused, nested sub-configs:
|
||||
The training process begins with proper configuration for the HILSerl environment. The main configuration class is `GymManipulatorConfig` in `lerobot/scripts/rl/gym_manipulator.py`, which contains nested `HILSerlRobotEnvConfig` and `DatasetConfig`. The configuration is organized into focused, nested sub-configs:
|
||||
|
||||
<!-- prettier-ignore-start -->
|
||||
```python
|
||||
@@ -160,8 +160,9 @@ The action processor (`action_processor`) handles outgoing actions and human int
|
||||
|
||||
1. **AddTeleopActionAsComplimentaryDataStep**: Captures teleoperator actions for logging
|
||||
2. **AddTeleopEventsAsInfoStep**: Records intervention events and episode control signals
|
||||
3. **InterventionActionProcessorStep**: Handles human interventions and episode termination
|
||||
4. **Inverse Kinematics Pipeline** (when enabled):
|
||||
3. **AddRobotObservationAsComplimentaryData**: Stores raw robot state for processing
|
||||
4. **InterventionActionProcessorStep**: Handles human interventions and episode termination
|
||||
5. **Inverse Kinematics Pipeline** (when enabled):
|
||||
- **MapDeltaActionToRobotActionStep**: Converts delta actions to robot action format
|
||||
- **EEReferenceAndDelta**: Computes end-effector reference and delta movements
|
||||
- **EEBoundsAndSafety**: Enforces workspace safety bounds
|
||||
@@ -346,7 +347,7 @@ With the bounds defined, you can safely collect demonstrations for training. Tra
|
||||
|
||||
**Setting Up Record Mode**
|
||||
|
||||
Create a configuration file for recording demonstrations (or edit an existing one like [env_config.json](https://huggingface.co/datasets/lerobot/config_examples/resolve/main/rl/env_config.json)):
|
||||
Create a configuration file for recording demonstrations (or edit an existing one like [env_config_so100.json](https://huggingface.co/datasets/aractingi/lerobot-example-config-files/blob/main/env_config_so100.json)):
|
||||
|
||||
1. Set `mode` to `"record"` at the root level
|
||||
2. Specify a unique `repo_id` for your dataset in the `dataset` section (e.g., "username/task_name")
|
||||
@@ -518,7 +519,7 @@ During the online training, press `space` to take over the policy and `space` ag
|
||||
Start the recording process, an example of the config file can be found [here](https://huggingface.co/datasets/aractingi/lerobot-example-config-files/blob/main/env_config_so100.json):
|
||||
|
||||
```bash
|
||||
python -m lerobot.rl.gym_manipulator --config_path src/lerobot/configs/env_config_so100.json
|
||||
python -m lerobot.scripts.rl.gym_manipulator --config_path src/lerobot/configs/env_config_so100.json
|
||||
```
|
||||
|
||||
During recording:
|
||||
@@ -549,7 +550,7 @@ Note: If you already know the crop parameters, you can skip this step and just s
|
||||
Use the `crop_dataset_roi.py` script to interactively select regions of interest in your camera images:
|
||||
|
||||
```bash
|
||||
python -m lerobot.rl.crop_dataset_roi --repo-id username/pick_lift_cube
|
||||
python -m lerobot.scripts.rl.crop_dataset_roi --repo-id username/pick_lift_cube
|
||||
```
|
||||
|
||||
1. For each camera view, the script will display the first frame
|
||||
@@ -618,7 +619,7 @@ Before training, you need to collect a dataset with labeled examples. The `recor
|
||||
To collect a dataset, you need to modify some parameters in the environment configuration based on HILSerlRobotEnvConfig.
|
||||
|
||||
```bash
|
||||
python -m lerobot.rl.gym_manipulator --config_path src/lerobot/configs/reward_classifier_train_config.json
|
||||
python -m lerobot.scripts.rl.gym_manipulator --config_path src/lerobot/configs/reward_classifier_train_config.json
|
||||
```
|
||||
|
||||
**Key Parameters for Data Collection**
|
||||
@@ -764,7 +765,7 @@ or set the argument in the json config file.
|
||||
Run `gym_manipulator.py` to test the model.
|
||||
|
||||
```bash
|
||||
python -m lerobot.rl.gym_manipulator --config_path path/to/env_config.json
|
||||
python -m lerobot.scripts.rl.gym_manipulator --config_path path/to/env_config.json
|
||||
```
|
||||
|
||||
The reward classifier will automatically provide rewards based on the visual input from the robot's cameras.
|
||||
@@ -772,12 +773,12 @@ The reward classifier will automatically provide rewards based on the visual inp
|
||||
**Example Workflow for training the reward classifier**
|
||||
|
||||
1. **Create the configuration files**:
|
||||
Create the necessary json configuration files for the reward classifier and the environment. Check the examples [here](https://huggingface.co/datasets/lerobot/config_examples/resolve/main/reward_classifier/config.json).
|
||||
Create the necessary json configuration files for the reward classifier and the environment. Check the examples [here](https://huggingface.co/datasets/aractingi/lerobot-example-config-files/tree/main).
|
||||
|
||||
2. **Collect a dataset**:
|
||||
|
||||
```bash
|
||||
python -m lerobot.rl.gym_manipulator --config_path src/lerobot/configs/env_config.json
|
||||
python -m lerobot.scripts.rl.gym_manipulator --config_path src/lerobot/configs/env_config.json
|
||||
```
|
||||
|
||||
3. **Train the classifier**:
|
||||
@@ -788,7 +789,7 @@ The reward classifier will automatically provide rewards based on the visual inp
|
||||
|
||||
4. **Test the classifier**:
|
||||
```bash
|
||||
python -m lerobot.rl.gym_manipulator --config_path src/lerobot/configs/env_config.json
|
||||
python -m lerobot.scripts.rl.gym_manipulator --config_path src/lerobot/configs/env_config.json
|
||||
```
|
||||
|
||||
### Training with Actor-Learner
|
||||
@@ -797,7 +798,7 @@ The LeRobot system uses a distributed actor-learner architecture for training. T
|
||||
|
||||
**Configuration Setup**
|
||||
|
||||
Create a training configuration file (example available [here](https://huggingface.co/datasets/lerobot/config_examples/resolve/main/rl/train_config.json)). The training config is based on the main `TrainRLServerPipelineConfig` class in `lerobot/configs/train.py`.
|
||||
Create a training configuration file (example available [here](https://huggingface.co/datasets/aractingi/lerobot-example-config-files/blob/main/train_config_hilserl_so100.json)). The training config is based on the main `TrainRLServerPipelineConfig` class in `lerobot/configs/train.py`.
|
||||
|
||||
1. Configure the policy settings (`type="sac"`, `device`, etc.)
|
||||
2. Set `dataset` to your cropped dataset
|
||||
@@ -810,7 +811,7 @@ Create a training configuration file (example available [here](https://huggingfa
|
||||
First, start the learner server process:
|
||||
|
||||
```bash
|
||||
python -m lerobot.rl.learner --config_path src/lerobot/configs/train_config_hilserl_so100.json
|
||||
python -m lerobot.scripts.rl.learner --config_path src/lerobot/configs/train_config_hilserl_so100.json
|
||||
```
|
||||
|
||||
The learner:
|
||||
@@ -825,7 +826,7 @@ The learner:
|
||||
In a separate terminal, start the actor process with the same configuration:
|
||||
|
||||
```bash
|
||||
python -m lerobot.rl.actor --config_path src/lerobot/configs/train_config_hilserl_so100.json
|
||||
python -m lerobot.scripts.rl.actor --config_path src/lerobot/configs/train_config_hilserl_so100.json
|
||||
```
|
||||
|
||||
The actor:
|
||||
|
||||
@@ -26,7 +26,7 @@ pip install -e ".[hilserl]"
|
||||
|
||||
## Configuration
|
||||
|
||||
To use `gym_hil` with LeRobot, you need to create a configuration file. An example is provided [here](https://huggingface.co/datasets/lerobot/config_examples/resolve/main/rl/gym_hil/env_config.json). Key configuration sections include:
|
||||
To use `gym_hil` with LeRobot, you need to create a configuration file. An example is provided [here](https://huggingface.co/datasets/aractingi/lerobot-example-config-files/blob/main/gym_hil_env.json). Key configuration sections include:
|
||||
|
||||
### Environment Type and Task
|
||||
|
||||
@@ -91,7 +91,7 @@ Important parameters:
|
||||
To run the environment, set mode to null:
|
||||
|
||||
```bash
|
||||
python -m lerobot.rl.gym_manipulator --config_path path/to/gym_hil_env.json
|
||||
python -m lerobot.scripts.rl.gym_manipulator --config_path path/to/gym_hil_env.json
|
||||
```
|
||||
|
||||
### Recording a Dataset
|
||||
@@ -118,21 +118,21 @@ To collect a dataset, set the mode to `record` whilst defining the repo_id and n
|
||||
```
|
||||
|
||||
```bash
|
||||
python -m lerobot.rl.gym_manipulator --config_path path/to/gym_hil_env.json
|
||||
python -m lerobot.scripts.rl.gym_manipulator --config_path path/to/gym_hil_env.json
|
||||
```
|
||||
|
||||
### Training a Policy
|
||||
|
||||
To train a policy, checkout the configuration example available [here](https://huggingface.co/datasets/lerobot/config_examples/resolve/main/rl/gym_hil/train_config.json) and run the actor and learner servers:
|
||||
To train a policy, checkout the configuration example available [here](https://huggingface.co/datasets/aractingi/lerobot-example-config-files/blob/main/train_gym_hil_env.json) and run the actor and learner servers:
|
||||
|
||||
```bash
|
||||
python -m lerobot.rl.actor --config_path path/to/train_gym_hil_env.json
|
||||
python -m lerobot.scripts.rl.actor --config_path path/to/train_gym_hil_env.json
|
||||
```
|
||||
|
||||
In a different terminal, run the learner server:
|
||||
|
||||
```bash
|
||||
python -m lerobot.rl.learner --config_path path/to/train_gym_hil_env.json
|
||||
python -m lerobot.scripts.rl.learner --config_path path/to/train_gym_hil_env.json
|
||||
```
|
||||
|
||||
The simulation environment provides a safe and repeatable way to develop and test your Human-In-the-Loop reinforcement learning components before deploying to real robots.
|
||||
|
||||
@@ -22,7 +22,7 @@ pip install -e ".[hilserl]"
|
||||
|
||||
## Teleoperate and Record a Dataset
|
||||
|
||||
To use `gym_hil` with LeRobot, you need to use a configuration file. An example config file can be found [here](https://huggingface.co/datasets/lerobot/config_examples/resolve/main/sim_il/env_config.json).
|
||||
To use `gym_hil` with LeRobot, you need to use a configuration file. An example config file can be found [here](https://huggingface.co/datasets/aractingi/lerobot-example-config-files/blob/main/env_config_gym_hil_il.json).
|
||||
|
||||
To teleoperate and collect a dataset, we need to modify this config file. Here's an example configuration for imitation learning data collection:
|
||||
|
||||
@@ -61,14 +61,14 @@ Then we can run this command to start:
|
||||
<hfoption id="Linux">
|
||||
|
||||
```bash
|
||||
python -m lerobot.rl.gym_manipulator --config_path path/to/env_config_gym_hil_il.json
|
||||
python -m lerobot.scripts.rl.gym_manipulator --config_path path/to/env_config_gym_hil_il.json
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
<hfoption id="MacOS">
|
||||
|
||||
```bash
|
||||
mjpython -m lerobot.rl.gym_manipulator --config_path path/to/env_config_gym_hil_il.json
|
||||
mjpython -m lerobot.scripts.rl.gym_manipulator --config_path path/to/env_config_gym_hil_il.json
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
@@ -165,7 +165,7 @@ huggingface-cli upload ${HF_USER}/il_sim_test${CKPT} \
|
||||
|
||||
## Evaluate your policy in Sim
|
||||
|
||||
To evaluate your policy we have to use a configuration file. An example can be found [here](https://huggingface.co/datasets/lerobot/config_examples/resolve/main/sim_il/eval_config.json).
|
||||
To evaluate your policy we have to use a configuration file. An example can be found [here](https://huggingface.co/datasets/aractingi/lerobot-example-config-files/blob/main/eval_config_gym_hil.json).
|
||||
|
||||
Here's an example evaluation configuration:
|
||||
|
||||
@@ -198,14 +198,14 @@ Then you can run this command to visualize your trained policy
|
||||
<hfoption id="Linux">
|
||||
|
||||
```bash
|
||||
python -m lerobot.rl.eval_policy --config_path=path/to/eval_config_gym_hil.json
|
||||
python -m lerobot.scripts.rl.eval_policy --config_path=path/to/eval_config_gym_hil.json
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
<hfoption id="MacOS">
|
||||
|
||||
```bash
|
||||
mjpython -m lerobot.rl.eval_policy --config_path=path/to/eval_config_gym_hil.json
|
||||
mjpython -m lerobot.scripts.rl.eval_policy --config_path=path/to/eval_config_gym_hil.json
|
||||
```
|
||||
|
||||
</hfoption>
|
||||
|
||||
@@ -14,7 +14,7 @@ Training data from one robot setup needs adaptation for deployment on different
|
||||
**That's where processors come in.** They serve as universal translators that bridge these gaps, ensuring seamless data flow from sensors to models to actuators.
|
||||
Processors handle all the preprocessing and postprocessing steps needed to convert raw environment data into model-ready inputs and vice versa.
|
||||
|
||||
This means that your favorite policy can be used like this:
|
||||
Now your favorite policy can be used like this:
|
||||
|
||||
```python
|
||||
import torch
|
||||
@@ -40,7 +40,7 @@ postprocessed_action = postprocessor(action)
|
||||
|
||||
## What are Processors?
|
||||
|
||||
In robotics, data comes in many forms: images from cameras, joint positions from sensors, text instructions from users, and more. Each type of data requires specific transformations before a model can use it effectively. Models need this data to be:
|
||||
In robotics, data comes in many forms - images from cameras, joint positions from sensors, text instructions from users, and more. Each type of data requires specific transformations before a model can use it effectively. Models need this data to be:
|
||||
|
||||
- **Normalized**: Scaled to appropriate ranges for neural network processing
|
||||
- **Batched**: Organized with proper dimensions for batch processing
|
||||
@@ -181,8 +181,8 @@ training_preprocessor = PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
RenameObservationsProcessorStep(rename_map={}), # Standardize keys
|
||||
AddBatchDimensionProcessorStep(), # Add batch dims
|
||||
TokenizerProcessorStep(tokenizer_name="...", ...), # Tokenize language
|
||||
DeviceProcessorStep(device="cuda"), # Move to GPU first
|
||||
NormalizerProcessorStep(features=..., stats=...), # Normalize on GPU
|
||||
DeviceProcessorStep(device="cuda"), # Move to GPU first ⚡
|
||||
NormalizerProcessorStep(features=..., stats=...), # Normalize on GPU ⚡
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@@ -8,7 +8,6 @@ This docs will guide you to:
|
||||
- Record a dataset and push it to the Hub
|
||||
- Load datasets for training with `LeRobotDataset`
|
||||
- Stream datasets without downloading using `StreamingLeRobotDataset`
|
||||
- Apply image transforms for data augmentation during training
|
||||
- Migrate existing `v2.1` datasets to `v3.0`
|
||||
|
||||
## What’s new in `v3`
|
||||
@@ -151,117 +150,6 @@ dataset = StreamingLeRobotDataset(repo_id) # streams directly from the Hub
|
||||
</figure>
|
||||
</div>
|
||||
|
||||
## Image transforms
|
||||
|
||||
Image transforms are data augmentations applied to camera frames during training to improve model robustness and generalization. LeRobot supports various transforms including brightness, contrast, saturation, hue, and sharpness adjustments.
|
||||
|
||||
### Using transforms during dataset creation/recording
|
||||
|
||||
Currently, transforms are applied during **training time only**, not during recording. When you create or record a dataset, the raw images are stored without transforms. This allows you to experiment with different augmentations later without re-recording data.
|
||||
|
||||
### Adding transforms to existing datasets (API)
|
||||
|
||||
Use the `image_transforms` parameter when loading a dataset for training:
|
||||
|
||||
```python
|
||||
from lerobot.datasets.lerobot_dataset import LeRobotDataset
|
||||
from lerobot.datasets.transforms import ImageTransforms, ImageTransformsConfig, ImageTransformConfig
|
||||
|
||||
# Option 1: Use default transform configuration (disabled by default)
|
||||
transforms_config = ImageTransformsConfig(
|
||||
enable=True, # Enable transforms
|
||||
max_num_transforms=3, # Apply up to 3 transforms per frame
|
||||
random_order=False, # Apply in standard order
|
||||
)
|
||||
transforms = ImageTransforms(transforms_config)
|
||||
|
||||
dataset = LeRobotDataset(
|
||||
repo_id="your-username/your-dataset",
|
||||
image_transforms=transforms
|
||||
)
|
||||
|
||||
# Option 2: Create custom transform configuration
|
||||
custom_transforms_config = ImageTransformsConfig(
|
||||
enable=True,
|
||||
max_num_transforms=2,
|
||||
random_order=True,
|
||||
tfs={
|
||||
"brightness": ImageTransformConfig(
|
||||
weight=1.0,
|
||||
type="ColorJitter",
|
||||
kwargs={"brightness": (0.7, 1.3)} # Adjust brightness range
|
||||
),
|
||||
"contrast": ImageTransformConfig(
|
||||
weight=2.0, # Higher weight = more likely to be selected
|
||||
type="ColorJitter",
|
||||
kwargs={"contrast": (0.8, 1.2)}
|
||||
),
|
||||
"sharpness": ImageTransformConfig(
|
||||
weight=0.5, # Lower weight = less likely to be selected
|
||||
type="SharpnessJitter",
|
||||
kwargs={"sharpness": (0.3, 2.0)}
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
dataset = LeRobotDataset(
|
||||
repo_id="your-username/your-dataset",
|
||||
image_transforms=ImageTransforms(custom_transforms_config)
|
||||
)
|
||||
|
||||
# Option 3: Use pure torchvision transforms
|
||||
from torchvision.transforms import v2
|
||||
|
||||
torchvision_transforms = v2.Compose([
|
||||
v2.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
|
||||
v2.GaussianBlur(kernel_size=3, sigma=(0.1, 2.0)),
|
||||
])
|
||||
|
||||
dataset = LeRobotDataset(
|
||||
repo_id="your-username/your-dataset",
|
||||
image_transforms=torchvision_transforms
|
||||
)
|
||||
```
|
||||
|
||||
### Available transform types
|
||||
|
||||
LeRobot provides several transform types:
|
||||
|
||||
- **`ColorJitter`**: Adjusts brightness, contrast, saturation, and hue
|
||||
- **`SharpnessJitter`**: Randomly adjusts image sharpness
|
||||
- **`Identity`**: No transformation (useful for testing)
|
||||
|
||||
You can also use any `torchvision.transforms.v2` transform by passing it directly to the `image_transforms` parameter.
|
||||
|
||||
### Configuration options
|
||||
|
||||
- **`enable`**: Enable/disable transforms (default: `False`)
|
||||
- **`max_num_transforms`**: Maximum number of transforms applied per frame (default: `3`)
|
||||
- **`random_order`**: Apply transforms in random order vs. standard order (default: `False`)
|
||||
- **`weight`**: Sampling probability for each transform (higher = more likely, if sum of weights is not 1, they will be normalized)
|
||||
- **`kwargs`**: Transform-specific parameters (e.g., brightness range)
|
||||
|
||||
### Visualizing transforms
|
||||
|
||||
Use the visualization script to preview how transforms affect your data:
|
||||
|
||||
```bash
|
||||
lerobot-imgtransform-viz \
|
||||
--repo-id=your-username/your-dataset \
|
||||
--output-dir=./transform_examples \
|
||||
--n-examples=5
|
||||
```
|
||||
|
||||
This saves example images showing the effect of each transform, helping you tune parameters.
|
||||
|
||||
### Best practices
|
||||
|
||||
- **Start conservative**: Begin with small ranges (e.g., brightness 0.9-1.1) and increase gradually
|
||||
- **Test first**: Use the visualization script to ensure transforms look reasonable
|
||||
- **Monitor training**: Strong augmentations can hurt performance if too aggressive
|
||||
- **Match your domain**: If your robot operates in varying lighting, use brightness/contrast transforms
|
||||
- **Combine wisely**: Using too many transforms simultaneously can make training unstable
|
||||
|
||||
## Migrate `v2.1` → `v3.0`
|
||||
|
||||
A converter aggregates per‑episode files into larger shards and writes episode offsets/metadata. Convert your dataset using the instructions below.
|
||||
|
||||
@@ -28,7 +28,7 @@ Links:
|
||||
|
||||
### Phone orientation and controls
|
||||
|
||||
- Orientation: hold the phone with the screen facing up and the top edge pointing in the same direction as the robot gripper. This ensures calibration aligns the phone’s frame with the robot frame so motion feels natural, see the image below for reference.
|
||||
- Orientation: hold the phone with the screen facing up and the top edge pointing in the same direction as the robot gripper. This ensures calibration aligns the phone’s frame with the robot frame so motion feels natural.
|
||||
- Enable/disable:
|
||||
- iOS: Hold `B1` to enable teleoperation, release to stop. The first press captures a reference pose.
|
||||
- Android: Press and hold the `Move` button, release to stop. The first press captures a reference pose.
|
||||
@@ -36,8 +36,6 @@ Links:
|
||||
- iOS: Analog input `A3` controls the gripper as velocity input.
|
||||
- Android: Buttons `A` and `B` act like increment/decrement (A opens, B closes). You can tune velocity in the `GripperVelocityToJoint` step.
|
||||
|
||||
<img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/lerobot/phone_teleop.webp" alt="Phone teleop orientation" title="Phone teleop orientation" width="40%">
|
||||
|
||||
### Step 1: Choose the platform
|
||||
|
||||
Modify the examples to use `PhoneOS.IOS` or `PhoneOS.ANDROID` in `PhoneConfig`. The API is identical across platforms, only the input source differs. All examples are under `examples/` and have `phone_so100_*.py` variants.
|
||||
@@ -66,79 +64,80 @@ Run on of the examples scripts to teleoperate, record a dataset, replay a datase
|
||||
|
||||
All scripts assume you configured your robot (e.g., SO-100 follower) and set the correct serial port.
|
||||
|
||||
Additionally you need to **copy the urdf of the robot to the examples folder**. For the examples in this tutorial (Using SO100/SO101) it is highly recommended to use the urdf in the [SO-ARM100 repo](https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf)
|
||||
- Android: after starting the script, open the printed local URL on your phone, tap Start, then press and hold Move.
|
||||
- iOS: open HEBI Mobile I/O first; B1 enables motion. A3 controls the gripper.
|
||||
|
||||
You can customize mapping or safety limits by editing the processor steps shown in the examples.
|
||||
|
||||
You can also remap inputs (e.g., use a different analog input) or adapt the pipeline to other robots (e.g., LeKiwi) by modifying the input and kinematics steps. More about this in the [Processors for Robots and Teleoperators](./processors_robots_teleop.mdx) guide.
|
||||
|
||||
- Run this example to teleoperate:
|
||||
|
||||
```bash
|
||||
python examples/phone_to_so100/teleoperate.py
|
||||
python examples/phone_so100_teleop.py
|
||||
```
|
||||
|
||||
After running the example:
|
||||
|
||||
- Android: after starting the script, open the printed local URL on your phone, tap Start, then press and hold Move.
|
||||
- iOS: open HEBI Mobile I/O first; B1 enables motion. A3 controls the gripper.
|
||||
|
||||
Additionally you can customize mapping or safety limits by editing the processor steps shown in the examples. You can also remap inputs (e.g., use a different analog input) or adapt the pipeline to other robots (e.g., LeKiwi) by modifying the input and kinematics steps. More about this in the [Processors for Robots and Teleoperators](./processors_robots_teleop.mdx) guide.
|
||||
|
||||
- Run this example to record a dataset, which saves absolute end effector observations and actions:
|
||||
|
||||
```bash
|
||||
python examples/phone_to_so100/record.py
|
||||
python examples/phone_so100_record.py
|
||||
```
|
||||
|
||||
- Run this example to replay recorded episodes:
|
||||
|
||||
```bash
|
||||
python examples/phone_to_so100/replay.py
|
||||
python examples/phone_so100_replay.py
|
||||
```
|
||||
|
||||
- Run this example to evaluate a pretrained policy:
|
||||
|
||||
```bash
|
||||
python examples/phone_to_so100/evaluate.py
|
||||
python examples/phone_so100_eval.py
|
||||
```
|
||||
|
||||
### Important pipeline steps and options
|
||||
|
||||
- Kinematics are used in multiple steps. We use [Placo](https://github.com/Rhoban/placo) which is a wrapper around Pinocchio for handling our kinematics. We construct the kinematics object by passing the robot's URDF and target frame. We set `target_frame_name` to the gripper frame.
|
||||
|
||||
```examples/phone_to_so100/teleoperate.py
|
||||
kinematics_solver = RobotKinematics(
|
||||
urdf_path="./SO101/so101_new_calib.urdf",
|
||||
target_frame_name="gripper_frame_link",
|
||||
joint_names=list(robot.bus.motors.keys()),
|
||||
```44:49:examples/phone_so100_teleop.py
|
||||
RobotKinematics(
|
||||
urdf_path="./src/lerobot/teleoperators/sim/so101_new_calib.urdf",
|
||||
target_frame_name="gripper_frame_link",
|
||||
joint_names=list(robot.bus.motors.keys()),
|
||||
)
|
||||
|
||||
```
|
||||
|
||||
- The `MapPhoneActionToRobotAction` step converts the calibrated phone pose and inputs into target deltas and gripper commands, below is shown what the step outputs.
|
||||
|
||||
```src/lerobot/teleoperators/phone/phone_processor.py
|
||||
action["enabled"] = enabled
|
||||
action["target_x"] = -pos[1] if enabled else 0.0
|
||||
action["target_y"] = pos[0] if enabled else 0.0
|
||||
action["target_z"] = pos[2] if enabled else 0.0
|
||||
action["target_wx"] = rotvec[1] if enabled else 0.0
|
||||
action["target_wy"] = rotvec[0] if enabled else 0.0
|
||||
action["target_wz"] = -rotvec[2] if enabled else 0.0
|
||||
action["gripper_vel"] = gripper_vel # Still send gripper action when disabled
|
||||
```72:83:src/lerobot/teleoperators/phone/phone_processor.py
|
||||
# Map calibrated phone pose to robot targets (enabled gates the motion)
|
||||
act.update(
|
||||
{
|
||||
"action.enabled": enabled,
|
||||
"action.target_x": -pos[1] if enabled else 0.0,
|
||||
"action.target_y": pos[0] if enabled else 0.0,
|
||||
"action.target_z": pos[2] if enabled else 0.0,
|
||||
"action.target_wx": rotvec[1] if enabled else 0.0,
|
||||
"action.target_wy": rotvec[0] if enabled else 0.0,
|
||||
"action.target_wz": -rotvec[2] if enabled else 0.0,
|
||||
"action.gripper": gripper,
|
||||
}
|
||||
)
|
||||
```
|
||||
|
||||
- The `EEReferenceAndDelta` step converts target deltas to an absolute desired EE pose, storing a reference on enable, the `end_effector_step_sizes` are the step sizes for the EE pose and can be modified to change the motion speed.
|
||||
|
||||
```examples/phone_to_so100/teleoperate.py
|
||||
```56:65:examples/phone_so100_teleop.py
|
||||
EEReferenceAndDelta(
|
||||
kinematics=kinematics_solver,
|
||||
end_effector_step_sizes={"x": 0.5, "y": 0.5, "z": 0.5},
|
||||
motor_names=list(robot.bus.motors.keys()),
|
||||
use_latched_reference=True,
|
||||
),
|
||||
)
|
||||
```
|
||||
|
||||
- The `EEBoundsAndSafety` step clamps EE motion to a workspace and checks for large ee step jumps to ensure safety. The `end_effector_bounds` are the bounds for the EE pose and can be modified to change the workspace. The `max_ee_step_m` and `max_ee_twist_step_rad` are the step limits for the EE pose and can be modified to change the safety limits.
|
||||
|
||||
```examples/phone_to_so100/teleoperate.py
|
||||
```61:66:examples/phone_so100_teleop.py
|
||||
EEBoundsAndSafety(
|
||||
end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
|
||||
max_ee_step_m=0.10,
|
||||
@@ -148,8 +147,11 @@ Additionally you can customize mapping or safety limits by editing the processor
|
||||
|
||||
- The `GripperVelocityToJoint` step turns a velocity‑like gripper input into absolute gripper position using the current measured state. The `speed_factor` is the factor by which the velocity is multiplied.
|
||||
|
||||
```examples/phone_to_so100/teleoperate.py
|
||||
GripperVelocityToJoint(speed_factor=20.0)
|
||||
```78:81:examples/phone_so100_teleop.py
|
||||
GripperVelocityToJoint(
|
||||
motor_names=list(robot.bus.motors.keys()),
|
||||
speed_factor=20.0,
|
||||
)
|
||||
```
|
||||
|
||||
#### Different IK initial guesses
|
||||
@@ -158,7 +160,7 @@ We use different IK initial guesses in the kinematic steps. As initial guess eit
|
||||
|
||||
- Closed loop (used in record/eval): sets `initial_guess_current_joints=True` so IK starts from the measured joints each frame.
|
||||
|
||||
```examples/phone_to_so100/record.py
|
||||
```71:76:examples/phone_so100_eval.py
|
||||
InverseKinematicsEEToJoints(
|
||||
kinematics=kinematics_solver,
|
||||
motor_names=list(robot.bus.motors.keys()),
|
||||
@@ -168,7 +170,7 @@ We use different IK initial guesses in the kinematic steps. As initial guess eit
|
||||
|
||||
- Open loop (used in replay): sets `initial_guess_current_joints=False` so IK continues from the previous IK solution rather than the measured state. This preserves action stability when we replay without feedback.
|
||||
|
||||
```examples/phone_to_so100/replay.py
|
||||
```80:86:examples/phone_so100_replay.py
|
||||
InverseKinematicsEEToJoints(
|
||||
kinematics=kinematics_solver,
|
||||
motor_names=list(robot.bus.motors.keys()),
|
||||
@@ -179,6 +181,7 @@ We use different IK initial guesses in the kinematic steps. As initial guess eit
|
||||
### Pipeline steps explained
|
||||
|
||||
- MapPhoneActionToRobotAction: converts calibrated phone pose and inputs into target deltas and a gripper command. Motion is gated by an enable signal (B1 on iOS, Move on Android).
|
||||
- AddRobotObservationAsComplimentaryData: reads current robot joints and inserts them under `complementary_data.raw_joint_positions` for FK/IK steps to use.
|
||||
- EEReferenceAndDelta: latches a reference EE pose on enable and combines it with target deltas to produce an absolute desired EE pose each frame. When disabled, it keeps sending the last commanded pose.
|
||||
- EEBoundsAndSafety: clamps the EE pose to a workspace and rate‑limits jumps for safety. Also declares `action.ee.*` features.
|
||||
- InverseKinematicsEEToJoints: turns an EE pose into joint positions with IK. `initial_guess_current_joints=True` is recommended for closed‑loop control; set `False` for open‑loop replay for stability.
|
||||
|
||||
@@ -17,7 +17,7 @@ We use the Phone to SO‑100 follower examples for concreteness, but the same pa
|
||||
|
||||
The examples in this guide use absolute end effector (EE) poses because they are easy to reason about. In practice, relative EE deltas or joint position are often preferred as learning features.
|
||||
|
||||
With processors, you choose the learning features you want to use for your policy. This could be joints positions/velocities, absolute EE, or relative EE positions. You can also choose to store other features, such as joint torques, motor currents, etc.
|
||||
You can choose what you save and learn from the teleop and robot action spaces, joints, absolute EE, or relative EE by using/implementing the right steps (and `transform_features()`) in your pipelines.
|
||||
|
||||
## Three pipelines
|
||||
|
||||
@@ -31,102 +31,99 @@ Each of these pipelines handle different conversions between different action an
|
||||
Below is an example of the three pipelines that we use in the phone to SO-100 follower examples:
|
||||
|
||||
```69:90:examples/phone_so100_record.py
|
||||
phone_to_robot_ee_pose_processor = RobotProcessorPipeline[RobotAction, RobotAction]( # teleop -> dataset action
|
||||
steps=[
|
||||
MapPhoneActionToRobotAction(platform=teleop_config.phone_os),
|
||||
EEReferenceAndDelta(
|
||||
kinematics=kinematics_solver, end_effector_step_sizes={"x": 0.5, "y": 0.5, "z": 0.5}, motor_names=list(robot.bus.motors.keys()),
|
||||
),
|
||||
EEBoundsAndSafety(
|
||||
end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]}, max_ee_step_m=0.20, max_ee_twist_step_rad=0.50,
|
||||
),
|
||||
GripperVelocityToJoint(),
|
||||
],
|
||||
to_transition=robot_action_to_transition,
|
||||
to_output=transition_to_robot_action,
|
||||
phone_to_robot_ee_pose = RobotProcessor( # teleop -> dataset action
|
||||
steps=[MapPhoneActionToRobotAction(platform=teleop_config.phone_os),
|
||||
AddRobotObservationAsComplimentaryData(robot=robot),
|
||||
EEReferenceAndDelta(kinematics=kinematics_solver,
|
||||
end_effector_step_sizes={"x": 0.5, "y": 0.5, "z": 0.5},
|
||||
motor_names=list(robot.bus.motors.keys())),
|
||||
EEBoundsAndSafety(end_effector_bounds={"min": [-1, -1, -1], "max": [1, 1, 1]},
|
||||
max_ee_step_m=0.20, max_ee_twist_step_rad=0.50)],
|
||||
to_transition=to_transition_teleop_action,
|
||||
to_output=lambda tr: tr,
|
||||
)
|
||||
|
||||
robot_ee_to_joints_processor = RobotProcessorPipeline[RobotAction, RobotAction]( # dataset action -> robot
|
||||
steps=[
|
||||
InverseKinematicsEEToJoints(
|
||||
kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys()), initial_guess_current_joints=True,
|
||||
),
|
||||
],
|
||||
to_transition=robot_action_to_transition,
|
||||
to_output=transition_to_robot_action,
|
||||
robot_ee_to_joints = RobotProcessor( # dataset action -> robot
|
||||
steps=[InverseKinematicsEEToJoints(kinematics=kinematics_solver,
|
||||
motor_names=list(robot.bus.motors.keys()),
|
||||
initial_guess_current_joints=True),
|
||||
GripperVelocityToJoint(motor_names=list(robot.bus.motors.keys()), speed_factor=20.0)],
|
||||
to_transition=lambda tr: tr,
|
||||
to_output=to_output_robot_action,
|
||||
)
|
||||
|
||||
robot_joints_to_ee_pose = RobotProcessorPipeline[RobotObservation, RobotObservation]( # robot obs -> dataset obs
|
||||
steps=[
|
||||
ForwardKinematicsJointsToEE(kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys()))
|
||||
],
|
||||
to_transition=observation_to_transition,
|
||||
to_output=transition_to_observation,
|
||||
robot_joints_to_ee_pose = RobotProcessor( # robot obs -> dataset obs
|
||||
steps=[ForwardKinematicsJointsToEE(kinematics=kinematics_solver,
|
||||
motor_names=list(robot.bus.motors.keys()))],
|
||||
to_transition=to_transition_robot_observation,
|
||||
to_output=lambda tr: tr,
|
||||
)
|
||||
```
|
||||
|
||||
## Why to_transition / to_output
|
||||
|
||||
To convert from robot/teleoperator to pipeline and back, we use the `to_transition` and `to_output` pipeline adapters.
|
||||
They standardize conversions to reduce boilerplate code, and form the bridge between the robot and teleoperators raw dictionaries and the pipeline’s `EnvTransition` format.
|
||||
They standardize conversions to reduce boilerplate code, and form the bridge between the robot and teleoperators raw dicts and the pipeline’s `EnvTransition` format.
|
||||
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.
|
||||
- `to_transition_teleop_action`: transforms the teleop action dict to a pipeline transition (puts keys under `action.*`, converts scalars/arrays to tensors, keeps objects like `Rotation` intact)
|
||||
- `to_output_robot_action`: transforms the pipeline transition to a robot action dict (extracts keys ending with `.pos`/`.vel` and strips `action.` prefix)
|
||||
- `to_transition_robot_observation`: transforms the robot observation dict to a pipeline transition (splits state vs images; stores state under `observation.state.*` and images under `observation.images.*`)
|
||||
|
||||
Checkout [src/lerobot/processor/converters.py](https://github.com/huggingface/lerobot/blob/main/src/lerobot/processor/converters.py) for more details.
|
||||
See `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(...)`.
|
||||
Dataset features are the keys saved in the dataset. Each step can declare what its dataset features are via `transform_features(...)`. We can then aggregate features per pipeline with `aggregate_pipeline_dataset_features()` and merge multiple groups with `merge_features(...)`.
|
||||
|
||||
Below is and example of how we declare features with the `transform_features` method in the phone to SO-100 follower examples:
|
||||
|
||||
```src/lerobot/robots/so100_follower/robot_kinematic_processor.py
|
||||
def transform_features(
|
||||
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
||||
# We only use the ee pose in the dataset, so we don't need the joint positions
|
||||
for n in self.motor_names:
|
||||
features[PipelineFeatureType.ACTION].pop(f"{n}.pos", None)
|
||||
# We specify the dataset features of this step that we want to be stored in the dataset
|
||||
for k in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]:
|
||||
features[PipelineFeatureType.ACTION][f"ee.{k}"] = PolicyFeature(
|
||||
type=FeatureType.STATE, shape=(1,)
|
||||
)
|
||||
return features
|
||||
```203:211:src/lerobot/robots/so100_follower/robot_kinematic_processor.py
|
||||
def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
|
||||
# Because this is last step we specify the dataset features of this step that we want to be stored in the dataset
|
||||
features["action.ee.x"] = float
|
||||
features["action.ee.y"] = float
|
||||
features["action.ee.z"] = float
|
||||
features["action.ee.wx"] = float
|
||||
features["action.ee.wy"] = float
|
||||
features["action.ee.wz"] = float
|
||||
return features
|
||||
```
|
||||
|
||||
Here we declare what PolicyFeatures we modify in this step, so we know what features we can expect when we run the processor. These features can then be aggregated and used to create the dataset features.
|
||||
Tip: declare features at the last step that produces them (e.g., `EEBoundsAndSafety` declares `action.ee.*`, `ForwardKinematicsJointsToEE` declares `observation.state.ee.*`).
|
||||
|
||||
Below is an example of how we aggregate and merge features in the phone to SO-100 record example:
|
||||
Below is an example of how we aggregate and merge features in the phone to SO-100 follower examples:
|
||||
|
||||
```121:145:examples/phone_so100_record.py
|
||||
features=combine_feature_dicts(
|
||||
# Run the feature contract of the pipelines
|
||||
# This tells you how the features would look like after the pipeline steps
|
||||
aggregate_pipeline_dataset_features(
|
||||
pipeline=phone_to_robot_ee_pose_processor,
|
||||
initial_features=create_initial_features(action=phone.action_features), # <- Action features we can expect, these come from our teleop device (phone) and action processor
|
||||
use_videos=True,
|
||||
),
|
||||
aggregate_pipeline_dataset_features(
|
||||
pipeline=robot_joints_to_ee_pose,
|
||||
initial_features=create_initial_features(observation=robot.observation_features), # <- Observation features we can expect, these come from our robot and observation processor
|
||||
use_videos=True,
|
||||
patterns=["observation.state.ee"], # <- Here you could optionally filter the features we want to store in the dataset, with a specific pattern
|
||||
action_ee = aggregate_pipeline_dataset_features(
|
||||
pipeline=phone_to_robot_ee_pose,
|
||||
initial_features=phone.action_features,
|
||||
use_videos=True,
|
||||
patterns=["action.ee"],
|
||||
)
|
||||
|
||||
),
|
||||
),
|
||||
gripper = aggregate_pipeline_dataset_features(
|
||||
pipeline=robot_ee_to_joints,
|
||||
initial_features={},
|
||||
use_videos=True,
|
||||
patterns=["action.gripper.pos", "observation.state.gripper.pos"],
|
||||
)
|
||||
|
||||
observation_ee = aggregate_pipeline_dataset_features(
|
||||
pipeline=robot_joints_to_ee_pose,
|
||||
initial_features=robot.observation_features,
|
||||
use_videos=True,
|
||||
patterns=["observation.state.ee"],
|
||||
)
|
||||
|
||||
dataset_features = merge_features(action_ee, gripper, observation_ee)
|
||||
```
|
||||
|
||||
How it works:
|
||||
|
||||
- `aggregate_pipeline_dataset_features(...)`: applies `transform_features` across the pipeline and filters by patterns (images included when `use_videos=True`, and state features included when `patterns` is specified).
|
||||
- `combine_feature_dicts(...)`: combine multiple feature dicts.
|
||||
- Recording with `record_loop(...)` uses `build_dataset_frame(...)` to build frames consistent with `dataset.features` before we call `add_frame(...)` to add the frame to the dataset.
|
||||
- `aggregate_pipeline_dataset_features(...)`: applies `transform_features` across the pipeline and filters by patterns (images included when `use_videos=True`).
|
||||
- `merge_features(...)`: combine multiple feature dicts.
|
||||
- Recording uses `to_dataset_frame(...)` to build frames consistent with `dataset.features` before we call `add_frame(...)` to add the frame to the dataset.
|
||||
|
||||
## Guidance when customizing robot pipelines
|
||||
|
||||
@@ -135,17 +132,17 @@ You can store any of the following features as your action/observation space:
|
||||
- Joint positions
|
||||
- Absolute EE poses
|
||||
- Relative EE deltas
|
||||
- Other features: joint velocity, torques, etc.
|
||||
- Other features: joint velocity, etc.
|
||||
|
||||
Pick what you want to use for your policy action and observation space and configure/modify the pipelines and steps accordingly.
|
||||
|
||||
### Different robots
|
||||
|
||||
- You can easily reuse pipelines, for example to use another robot with phone teleop, modify the examples and swap the robot `RobotKinematics` (URDF) and `motor_names` to use your own robot with Phone teleop. Additionally you should ensure `target_frame_name` points to your gripper/wrist.
|
||||
- Swap `RobotKinematics` URDF and `motor_names`. Ensure `target_frame_name` points to your gripper/wrist.
|
||||
|
||||
### Safety first
|
||||
|
||||
- When changing pipelines, start with tight bounds, implement safety steps when working with real robots.
|
||||
- Its advised to start with simulation first and then move to real robots.
|
||||
|
||||
Thats it! We hope this guide helps you get started with customizing your robot pipelines, If you run into any issues at any point, jump into our [Discord community](https://discord.com/invite/s3KuuzsPFb) for support.
|
||||
Hope this guide helps you get started with customizing your robot pipelines, If you run into any issues at any point, jump into our [Discord community](https://discord.com/invite/s3KuuzsPFb) for support.
|
||||
|
||||
@@ -136,7 +136,7 @@ print(f"{dataset[0]['action'].shape=}\n") # (64, c)
|
||||
# PyTorch datasets.
|
||||
dataloader = torch.utils.data.DataLoader(
|
||||
dataset,
|
||||
num_workers=4,
|
||||
num_workers=0,
|
||||
batch_size=32,
|
||||
shuffle=True,
|
||||
)
|
||||
@@ -0,0 +1,139 @@
|
||||
# Copyright 2024 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.
|
||||
|
||||
"""
|
||||
This script demonstrates how to evaluate a pretrained policy from the HuggingFace Hub or from your local
|
||||
training outputs directory. In the latter case, you might want to run examples/3_train_policy.py first.
|
||||
|
||||
It requires the installation of the 'gym_pusht' simulation environment. Install it by running:
|
||||
```bash
|
||||
pip install -e ".[pusht]"
|
||||
```
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import gym_pusht # noqa: F401
|
||||
import gymnasium as gym
|
||||
import imageio
|
||||
import numpy
|
||||
import torch
|
||||
|
||||
from lerobot.policies.diffusion.modeling_diffusion import DiffusionPolicy
|
||||
|
||||
# Create a directory to store the video of the evaluation
|
||||
output_directory = Path("outputs/eval/example_pusht_diffusion")
|
||||
output_directory.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Select your device
|
||||
device = "cuda"
|
||||
|
||||
# Provide the [hugging face repo id](https://huggingface.co/lerobot/diffusion_pusht):
|
||||
pretrained_policy_path = "lerobot/diffusion_pusht"
|
||||
# OR a path to a local outputs/train folder.
|
||||
# pretrained_policy_path = Path("outputs/train/example_pusht_diffusion")
|
||||
|
||||
policy = DiffusionPolicy.from_pretrained(pretrained_policy_path)
|
||||
|
||||
# Initialize evaluation environment to render two observation types:
|
||||
# an image of the scene and state/position of the agent. The environment
|
||||
# also automatically stops running after 300 interactions/steps.
|
||||
env = gym.make(
|
||||
"gym_pusht/PushT-v0",
|
||||
obs_type="pixels_agent_pos",
|
||||
max_episode_steps=300,
|
||||
)
|
||||
|
||||
# We can verify that the shapes of the features expected by the policy match the ones from the observations
|
||||
# produced by the environment
|
||||
print(policy.config.input_features)
|
||||
print(env.observation_space)
|
||||
|
||||
# Similarly, we can check that the actions produced by the policy will match the actions expected by the
|
||||
# environment
|
||||
print(policy.config.output_features)
|
||||
print(env.action_space)
|
||||
|
||||
# Reset the policy and environments to prepare for rollout
|
||||
policy.reset()
|
||||
numpy_observation, info = env.reset(seed=42)
|
||||
|
||||
# Prepare to collect every rewards and all the frames of the episode,
|
||||
# from initial state to final state.
|
||||
rewards = []
|
||||
frames = []
|
||||
|
||||
# Render frame of the initial state
|
||||
frames.append(env.render())
|
||||
|
||||
step = 0
|
||||
done = False
|
||||
while not done:
|
||||
# Prepare observation for the policy running in Pytorch
|
||||
state = torch.from_numpy(numpy_observation["agent_pos"])
|
||||
image = torch.from_numpy(numpy_observation["pixels"])
|
||||
|
||||
# Convert to float32 with image from channel first in [0,255]
|
||||
# to channel last in [0,1]
|
||||
state = state.to(torch.float32)
|
||||
image = image.to(torch.float32) / 255
|
||||
image = image.permute(2, 0, 1)
|
||||
|
||||
# Send data tensors from CPU to GPU
|
||||
state = state.to(device, non_blocking=True)
|
||||
image = image.to(device, non_blocking=True)
|
||||
|
||||
# Add extra (empty) batch dimension, required to forward the policy
|
||||
state = state.unsqueeze(0)
|
||||
image = image.unsqueeze(0)
|
||||
|
||||
# Create the policy input dictionary
|
||||
observation = {
|
||||
"observation.state": state,
|
||||
"observation.image": image,
|
||||
}
|
||||
|
||||
# Predict the next action with respect to the current observation
|
||||
with torch.inference_mode():
|
||||
action = policy.select_action(observation)
|
||||
|
||||
# Prepare the action for the environment
|
||||
numpy_action = action.squeeze(0).to("cpu").numpy()
|
||||
|
||||
# Step through the environment and receive a new observation
|
||||
numpy_observation, reward, terminated, truncated, info = env.step(numpy_action)
|
||||
print(f"{step=} {reward=} {terminated=}")
|
||||
|
||||
# Keep track of all the rewards and frames
|
||||
rewards.append(reward)
|
||||
frames.append(env.render())
|
||||
|
||||
# The rollout is considered done when the success state is reached (i.e. terminated is True),
|
||||
# or the maximum number of iterations is reached (i.e. truncated is True)
|
||||
done = terminated | truncated | done
|
||||
step += 1
|
||||
|
||||
if terminated:
|
||||
print("Success!")
|
||||
else:
|
||||
print("Failure!")
|
||||
|
||||
# Get the speed of environment (i.e. its number of frames per second).
|
||||
fps = env.metadata["render_fps"]
|
||||
|
||||
# Encode all frames into a mp4 video.
|
||||
video_path = output_directory / "rollout.mp4"
|
||||
imageio.mimsave(str(video_path), numpy.stack(frames), fps=fps)
|
||||
|
||||
print(f"Video of the evaluation is available in '{video_path}'.")
|
||||
@@ -12,7 +12,11 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""This script demonstrates how to train Diffusion Policy on the PushT environment."""
|
||||
"""This script demonstrates how to train Diffusion Policy on the PushT environment.
|
||||
|
||||
Once you have trained a model with this script, you can try to evaluate it on
|
||||
examples/2_evaluate_pretrained_policy.py
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
@@ -23,7 +27,6 @@ from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetad
|
||||
from lerobot.datasets.utils import dataset_to_policy_features
|
||||
from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
|
||||
from lerobot.policies.diffusion.modeling_diffusion import DiffusionPolicy
|
||||
from lerobot.policies.factory import make_pre_post_processors
|
||||
|
||||
|
||||
def main():
|
||||
@@ -53,10 +56,9 @@ def main():
|
||||
cfg = DiffusionConfig(input_features=input_features, output_features=output_features)
|
||||
|
||||
# We can now instantiate our policy with this config and the dataset stats.
|
||||
policy = DiffusionPolicy(cfg)
|
||||
policy = DiffusionPolicy(cfg, dataset_stats=dataset_metadata.stats)
|
||||
policy.train()
|
||||
policy.to(device)
|
||||
preprocessor, postprocessor = make_pre_post_processors(cfg, dataset_stats=dataset_metadata.stats)
|
||||
|
||||
# Another policy-dataset interaction is with the delta_timestamps. Each policy expects a given number frames
|
||||
# which can differ for inputs, outputs and rewards (if there are some).
|
||||
@@ -97,7 +99,7 @@ def main():
|
||||
done = False
|
||||
while not done:
|
||||
for batch in dataloader:
|
||||
batch = preprocessor(batch)
|
||||
batch = {k: (v.to(device) if isinstance(v, torch.Tensor) else v) for k, v in batch.items()}
|
||||
loss, _ = policy.forward(batch)
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
@@ -112,8 +114,6 @@ def main():
|
||||
|
||||
# Save a policy checkpoint.
|
||||
policy.save_pretrained(output_directory)
|
||||
preprocessor.save_pretrained(output_directory)
|
||||
postprocessor.save_pretrained(output_directory)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
@@ -0,0 +1,311 @@
|
||||
This tutorial will explain the training script, how to use it, and particularly how to configure everything needed for the training run.
|
||||
|
||||
> **Note:** The following assumes you're running these commands on a machine equipped with a cuda GPU. If you don't have one (or if you're using a Mac), you can add `--policy.device=cpu` (`--policy.device=mps` respectively). However, be advised that the code executes much slower on cpu.
|
||||
|
||||
## The training script
|
||||
|
||||
LeRobot offers a training script at [`lerobot/scripts/train.py`](../src/lerobot/scripts/train.py). At a high level it does the following:
|
||||
|
||||
- Initialize/load a configuration for the following steps using.
|
||||
- Instantiates a dataset.
|
||||
- (Optional) Instantiates a simulation environment corresponding to that dataset.
|
||||
- Instantiates a policy.
|
||||
- Runs a standard training loop with forward pass, backward pass, optimization step, and occasional logging, evaluation (of the policy on the environment), and checkpointing.
|
||||
|
||||
## Overview of the configuration system
|
||||
|
||||
In the training script, the main function `train` expects a `TrainPipelineConfig` object:
|
||||
|
||||
<!-- prettier-ignore-start -->
|
||||
```python
|
||||
# train.py
|
||||
@parser.wrap()
|
||||
def train(cfg: TrainPipelineConfig):
|
||||
```
|
||||
<!-- prettier-ignore-end -->
|
||||
|
||||
You can inspect the `TrainPipelineConfig` defined in [`lerobot/configs/train.py`](../src/lerobot/configs/train.py) (which is heavily commented and meant to be a reference to understand any option)
|
||||
|
||||
When running the script, inputs for the command line are parsed thanks to the `@parser.wrap()` decorator and an instance of this class is automatically generated. Under the hood, this is done with [Draccus](https://github.com/dlwh/draccus) which is a tool dedicated to this purpose. If you're familiar with Hydra, Draccus can similarly load configurations from config files (.json, .yaml) and also override their values through command line inputs. Unlike Hydra, these configurations are pre-defined in the code through dataclasses rather than being defined entirely in config files. This allows for more rigorous serialization/deserialization, typing, and to manipulate configuration as objects directly in the code and not as dictionaries or namespaces (which enables nice features in an IDE such as autocomplete, jump-to-def, etc.)
|
||||
|
||||
Let's have a look at a simplified example. Amongst other attributes, the training config has the following attributes:
|
||||
|
||||
<!-- prettier-ignore-start -->
|
||||
```python
|
||||
@dataclass
|
||||
class TrainPipelineConfig:
|
||||
dataset: DatasetConfig
|
||||
env: envs.EnvConfig | None = None
|
||||
policy: PreTrainedConfig | None = None
|
||||
```
|
||||
<!-- prettier-ignore-end -->
|
||||
|
||||
in which `DatasetConfig` for example is defined as such:
|
||||
|
||||
<!-- prettier-ignore-start -->
|
||||
```python
|
||||
@dataclass
|
||||
class DatasetConfig:
|
||||
repo_id: str
|
||||
episodes: list[int] | None = None
|
||||
video_backend: str = "pyav"
|
||||
```
|
||||
<!-- prettier-ignore-end -->
|
||||
|
||||
This creates a hierarchical relationship where, for example assuming we have a `cfg` instance of `TrainPipelineConfig`, we can access the `repo_id` value with `cfg.dataset.repo_id`.
|
||||
From the command line, we can specify this value by using a very similar syntax `--dataset.repo_id=repo/id`.
|
||||
|
||||
By default, every field takes its default value specified in the dataclass. If a field doesn't have a default value, it needs to be specified either from the command line or from a config file – which path is also given in the command line (more in this below). In the example above, the `dataset` field doesn't have a default value which means it must be specified.
|
||||
|
||||
## Specifying values from the CLI
|
||||
|
||||
Let's say that we want to train [Diffusion Policy](../src/lerobot/policies/diffusion) on the [pusht](https://huggingface.co/datasets/lerobot/pusht) dataset, using the [gym_pusht](https://github.com/huggingface/gym-pusht) environment for evaluation. The command to do so would look like this:
|
||||
|
||||
```bash
|
||||
lerobot-train \
|
||||
--dataset.repo_id=lerobot/pusht \
|
||||
--policy.type=diffusion \
|
||||
--env.type=pusht
|
||||
```
|
||||
|
||||
Let's break this down:
|
||||
|
||||
- To specify the dataset, we just need to specify its `repo_id` on the hub which is the only required argument in the `DatasetConfig`. The rest of the fields have default values and in this case we are fine with those so we can just add the option `--dataset.repo_id=lerobot/pusht`.
|
||||
- To specify the policy, we can just select diffusion policy using `--policy` appended with `.type`. Here, `.type` is a special argument which allows us to select config classes inheriting from `draccus.ChoiceRegistry` and that have been decorated with the `register_subclass()` method. To have a better explanation of this feature, have a look at this [Draccus demo](https://github.com/dlwh/draccus?tab=readme-ov-file#more-flexible-configuration-with-choice-types). In our code, we use this mechanism mainly to select policies, environments, robots, and some other components like optimizers. The policies available to select are located in [lerobot/policies](../src/lerobot/policies)
|
||||
- Similarly, we select the environment with `--env.type=pusht`. The different environment configs are available in [`lerobot/envs/configs.py`](../src/lerobot/envs/configs.py)
|
||||
|
||||
Let's see another example. Let's say you've been training [ACT](../src/lerobot/policies/act) on [lerobot/aloha_sim_insertion_human](https://huggingface.co/datasets/lerobot/aloha_sim_insertion_human) using the [gym-aloha](https://github.com/huggingface/gym-aloha) environment for evaluation with:
|
||||
|
||||
```bash
|
||||
lerobot-train \
|
||||
--policy.type=act \
|
||||
--dataset.repo_id=lerobot/aloha_sim_insertion_human \
|
||||
--env.type=aloha \
|
||||
--output_dir=outputs/train/act_aloha_insertion
|
||||
```
|
||||
|
||||
> Notice we added `--output_dir` to explicitly tell where to write outputs from this run (checkpoints, training state, configs etc.). This is not mandatory and if you don't specify it, a default directory will be created from the current date and time, env.type and policy.type. This will typically look like `outputs/train/2025-01-24/16-10-05_aloha_act`.
|
||||
|
||||
We now want to train a different policy for aloha on another task. We'll change the dataset and use [lerobot/aloha_sim_transfer_cube_human](https://huggingface.co/datasets/lerobot/aloha_sim_transfer_cube_human) instead. Of course, we also need to change the task of the environment as well to match this other task.
|
||||
Looking at the [`AlohaEnv`](../src/lerobot/envs/configs.py) config, the task is `"AlohaInsertion-v0"` by default, which corresponds to the task we trained on in the command above. The [gym-aloha](https://github.com/huggingface/gym-aloha?tab=readme-ov-file#description) environment also has the `AlohaTransferCube-v0` task which corresponds to this other task we want to train on. Putting this together, we can train this new policy on this different task using:
|
||||
|
||||
```bash
|
||||
lerobot-train \
|
||||
--policy.type=act \
|
||||
--dataset.repo_id=lerobot/aloha_sim_transfer_cube_human \
|
||||
--env.type=aloha \
|
||||
--env.task=AlohaTransferCube-v0 \
|
||||
--output_dir=outputs/train/act_aloha_transfer
|
||||
```
|
||||
|
||||
## Loading from a config file
|
||||
|
||||
Now, let's assume that we want to reproduce the run just above. That run has produced a `train_config.json` file in its checkpoints, which serializes the `TrainPipelineConfig` instance it used:
|
||||
|
||||
```json
|
||||
{
|
||||
"dataset": {
|
||||
"repo_id": "lerobot/aloha_sim_transfer_cube_human",
|
||||
"episodes": null,
|
||||
...
|
||||
},
|
||||
"env": {
|
||||
"type": "aloha",
|
||||
"task": "AlohaTransferCube-v0",
|
||||
"fps": 50,
|
||||
...
|
||||
},
|
||||
"policy": {
|
||||
"type": "act",
|
||||
"n_obs_steps": 1,
|
||||
...
|
||||
},
|
||||
...
|
||||
}
|
||||
```
|
||||
|
||||
We can then simply load the config values from this file using:
|
||||
|
||||
```bash
|
||||
lerobot-train \
|
||||
--config_path=outputs/train/act_aloha_transfer/checkpoints/last/pretrained_model/ \
|
||||
--output_dir=outputs/train/act_aloha_transfer_2
|
||||
```
|
||||
|
||||
`--config_path` is also a special argument which allows to initialize the config from a local config file. It can point to a directory that contains `train_config.json` or to the config file itself directly.
|
||||
|
||||
Similarly to Hydra, we can still override some parameters in the CLI if we want to, e.g.:
|
||||
|
||||
```bash
|
||||
lerobot-train \
|
||||
--config_path=outputs/train/act_aloha_transfer/checkpoints/last/pretrained_model/ \
|
||||
--output_dir=outputs/train/act_aloha_transfer_2
|
||||
--policy.n_action_steps=80
|
||||
```
|
||||
|
||||
> Note: While `--output_dir` is not required in general, in this case we need to specify it since it will otherwise take the value from the `train_config.json` (which is `outputs/train/act_aloha_transfer`). In order to prevent accidental deletion of previous run checkpoints, we raise an error if you're trying to write in an existing directory. This is not the case when resuming a run, which is what you'll learn next.
|
||||
|
||||
`--config_path` can also accept the repo_id of a repo on the hub that contains a `train_config.json` file, e.g. running:
|
||||
|
||||
```bash
|
||||
lerobot-train --config_path=lerobot/diffusion_pusht
|
||||
```
|
||||
|
||||
will start a training run with the same configuration used for training [lerobot/diffusion_pusht](https://huggingface.co/lerobot/diffusion_pusht)
|
||||
|
||||
## Resume training
|
||||
|
||||
Being able to resume a training run is important in case it crashed or aborted for any reason. We'll demonstrate how to do that here.
|
||||
|
||||
Let's reuse the command from the previous run and add a few more options:
|
||||
|
||||
```bash
|
||||
lerobot-train \
|
||||
--policy.type=act \
|
||||
--dataset.repo_id=lerobot/aloha_sim_transfer_cube_human \
|
||||
--env.type=aloha \
|
||||
--env.task=AlohaTransferCube-v0 \
|
||||
--log_freq=25 \
|
||||
--save_freq=100 \
|
||||
--output_dir=outputs/train/run_resumption
|
||||
```
|
||||
|
||||
Here we've taken care to set up the log frequency and checkpointing frequency to low numbers so we can showcase resumption. You should be able to see some logging and have a first checkpoint within 1 minute (depending on hardware). Wait for the first checkpoint to happen, you should see a line that looks like this in your terminal:
|
||||
|
||||
```
|
||||
INFO 2025-01-24 16:10:56 ts/train.py:263 Checkpoint policy after step 100
|
||||
```
|
||||
|
||||
Now let's simulate a crash by killing the process (hit `ctrl`+`c`). We can then simply resume this run from the last checkpoint available with:
|
||||
|
||||
```bash
|
||||
lerobot-train \
|
||||
--config_path=outputs/train/run_resumption/checkpoints/last/pretrained_model/ \
|
||||
--resume=true
|
||||
```
|
||||
|
||||
You should see from the logging that your training picks up from where it left off.
|
||||
|
||||
Another reason for which you might want to resume a run is simply to extend training and add more training steps. The number of training steps is set by the option `--steps`, which is 100 000 by default.
|
||||
You could double the number of steps of the previous run with:
|
||||
|
||||
```bash
|
||||
lerobot-train \
|
||||
--config_path=outputs/train/run_resumption/checkpoints/last/pretrained_model/ \
|
||||
--resume=true \
|
||||
--steps=200000
|
||||
```
|
||||
|
||||
## Outputs of a run
|
||||
|
||||
In the output directory, there will be a folder called `checkpoints` with the following structure:
|
||||
|
||||
```bash
|
||||
outputs/train/run_resumption/checkpoints
|
||||
├── 000100 # checkpoint_dir for training step 100
|
||||
│ ├── pretrained_model/
|
||||
│ │ ├── config.json # policy config
|
||||
│ │ ├── model.safetensors # policy weights
|
||||
│ │ └── train_config.json # train config
|
||||
│ └── training_state/
|
||||
│ ├── optimizer_param_groups.json # optimizer param groups
|
||||
│ ├── optimizer_state.safetensors # optimizer state
|
||||
│ ├── rng_state.safetensors # rng states
|
||||
│ ├── scheduler_state.json # scheduler state
|
||||
│ └── training_step.json # training step
|
||||
├── 000200
|
||||
└── last -> 000200 # symlink to the last available checkpoint
|
||||
```
|
||||
|
||||
## Fine-tuning a pre-trained policy
|
||||
|
||||
In addition to the features currently in Draccus, we've added a special `.path` argument for the policy, which allows to load a policy as you would with `PreTrainedPolicy.from_pretrained()`. In that case, `path` can be a local directory that contains a checkpoint or a repo_id pointing to a pretrained policy on the hub.
|
||||
|
||||
For example, we could fine-tune a [policy pre-trained on the aloha transfer task](https://huggingface.co/lerobot/act_aloha_sim_transfer_cube_human) on the aloha insertion task. We can achieve this with:
|
||||
|
||||
```bash
|
||||
lerobot-train \
|
||||
--policy.path=lerobot/act_aloha_sim_transfer_cube_human \
|
||||
--dataset.repo_id=lerobot/aloha_sim_insertion_human \
|
||||
--env.type=aloha \
|
||||
--env.task=AlohaInsertion-v0
|
||||
```
|
||||
|
||||
When doing so, keep in mind that the features of the fine-tuning dataset would have to match the input/output features of the pretrained policy.
|
||||
|
||||
## Typical logs and metrics
|
||||
|
||||
When you start the training process, you will first see your full configuration being printed in the terminal. You can check it to make sure that you configured your run correctly. The final configuration will also be saved with the checkpoint.
|
||||
|
||||
After that, you will see training log like this one:
|
||||
|
||||
```
|
||||
INFO 2024-08-14 13:35:12 ts/train.py:192 step:0 smpl:64 ep:1 epch:0.00 loss:1.112 grdn:15.387 lr:2.0e-07 updt_s:1.738 data_s:4.774
|
||||
```
|
||||
|
||||
or evaluation log:
|
||||
|
||||
```
|
||||
INFO 2024-08-14 13:38:45 ts/train.py:226 step:100 smpl:6K ep:52 epch:0.25 ∑rwrd:20.693 success:0.0% eval_s:120.266
|
||||
```
|
||||
|
||||
These logs will also be saved in wandb if `wandb.enable` is set to `true`. Here are the meaning of some abbreviations:
|
||||
|
||||
- `smpl`: number of samples seen during training.
|
||||
- `ep`: number of episodes seen during training. An episode contains multiple samples in a complete manipulation task.
|
||||
- `epch`: number of time all unique samples are seen (epoch).
|
||||
- `grdn`: gradient norm.
|
||||
- `∑rwrd`: compute the sum of rewards in every evaluation episode and then take an average of them.
|
||||
- `success`: average success rate of eval episodes. Reward and success are usually different except for the sparsing reward setting, where reward=1 only when the task is completed successfully.
|
||||
- `eval_s`: time to evaluate the policy in the environment, in second.
|
||||
- `updt_s`: time to update the network parameters, in second.
|
||||
- `data_s`: time to load a batch of data, in second.
|
||||
|
||||
Some metrics are useful for initial performance profiling. For example, if you find the current GPU utilization is low via the `nvidia-smi` command and `data_s` sometimes is too high, you may need to modify batch size or number of dataloading workers to accelerate dataloading. We also recommend [pytorch profiler](https://github.com/huggingface/lerobot?tab=readme-ov-file#improve-your-code-with-profiling) for detailed performance probing.
|
||||
|
||||
## In short
|
||||
|
||||
We'll summarize here the main use cases to remember from this tutorial.
|
||||
|
||||
#### Train a policy from scratch – CLI
|
||||
|
||||
```bash
|
||||
lerobot-train \
|
||||
--policy.type=act \ # <- select 'act' policy
|
||||
--env.type=pusht \ # <- select 'pusht' environment
|
||||
--dataset.repo_id=lerobot/pusht # <- train on this dataset
|
||||
```
|
||||
|
||||
#### Train a policy from scratch - config file + CLI
|
||||
|
||||
```bash
|
||||
lerobot-train \
|
||||
--config_path=path/to/pretrained_model \ # <- can also be a repo_id
|
||||
--policy.n_action_steps=80 # <- you may still override values
|
||||
```
|
||||
|
||||
#### Resume/continue a training run
|
||||
|
||||
```bash
|
||||
lerobot-train \
|
||||
--config_path=checkpoint/pretrained_model/ \
|
||||
--resume=true \
|
||||
--steps=200000 # <- you can change some training parameters
|
||||
```
|
||||
|
||||
#### Fine-tuning
|
||||
|
||||
```bash
|
||||
lerobot-train \
|
||||
--policy.path=lerobot/act_aloha_sim_transfer_cube_human \ # <- can also be a local path to a checkpoint
|
||||
--dataset.repo_id=lerobot/aloha_sim_insertion_human \
|
||||
--env.type=aloha \
|
||||
--env.task=AlohaInsertion-v0
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
Now that you know the basics of how to train a policy, you might want to know how to apply this knowledge to actual robots, or how to record your own datasets and train policies on your specific task?
|
||||
If that's the case, head over to the next tutorial [`7_get_started_with_real_robot.md`](./7_get_started_with_real_robot.md).
|
||||
|
||||
Or in the meantime, happy training! 🤗
|
||||
@@ -13,7 +13,11 @@
|
||||
# limitations under the License.
|
||||
|
||||
"""This script demonstrates how to train a Diffusion Policy on the PushT environment,
|
||||
using a dataset processed in streaming mode."""
|
||||
using a dataset processed in streaming mode.
|
||||
|
||||
Once you have trained a model with this script, you can try to evaluate it on
|
||||
examples/2_evaluate_pretrained_policy.py
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
@@ -26,7 +30,6 @@ from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset
|
||||
from lerobot.datasets.utils import dataset_to_policy_features
|
||||
from lerobot.policies.act.configuration_act import ACTConfig
|
||||
from lerobot.policies.act.modeling_act import ACTPolicy
|
||||
from lerobot.policies.factory import make_pre_post_processors
|
||||
|
||||
|
||||
def main():
|
||||
@@ -47,7 +50,9 @@ def main():
|
||||
training_steps = 10
|
||||
log_freq = 1
|
||||
|
||||
dataset_id = "lerobot/droid_1.0.1" # 26M frames! Would require 4TB of disk space if installed locally (:
|
||||
dataset_id = (
|
||||
"aractingi/droid_1.0.1" # 26M frames! Would require 4TB of disk space if installed locally (:
|
||||
)
|
||||
dataset_metadata = LeRobotDatasetMetadata(dataset_id)
|
||||
features = dataset_to_policy_features(dataset_metadata.features)
|
||||
output_features = {key: ft for key, ft in features.items() if ft.type is FeatureType.ACTION}
|
||||
@@ -55,10 +60,9 @@ def main():
|
||||
|
||||
# We can now instantiate our policy with this config and the dataset stats.
|
||||
cfg = ACTConfig(input_features=input_features, output_features=output_features)
|
||||
policy = ACTPolicy(cfg)
|
||||
policy = ACTPolicy(cfg, dataset_stats=dataset_metadata.stats)
|
||||
policy.train()
|
||||
policy.to(device)
|
||||
preprocessor, postprocessor = make_pre_post_processors(cfg, dataset_stats=dataset_metadata.stats)
|
||||
|
||||
# Delta timestamps are used to (1) augment frames used during training and (2) supervise the policy.
|
||||
# Here, we use delta-timestamps to only provide ground truth actions for supervision
|
||||
@@ -85,7 +89,13 @@ def main():
|
||||
done = False
|
||||
while not done:
|
||||
for batch in dataloader:
|
||||
batch = preprocessor(batch)
|
||||
batch = {
|
||||
k: (v.type(torch.float32) if isinstance(v, torch.Tensor) and v.dtype != torch.bool else v)
|
||||
for k, v in batch.items()
|
||||
}
|
||||
batch = {k: (v.to(device) if isinstance(v, torch.Tensor) else v) for k, v in batch.items()}
|
||||
|
||||
# batch = {k: (v.to(device) if isinstance(v, torch.Tensor) else v) for k, v in batch.items()}
|
||||
loss, _ = policy.forward(batch)
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
@@ -100,8 +110,6 @@ def main():
|
||||
|
||||
# Save a policy checkpoint.
|
||||
policy.save_pretrained(output_directory)
|
||||
preprocessor.save_pretrained(output_directory)
|
||||
postprocessor.save_pretrained(output_directory)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
@@ -1,177 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2024 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.
|
||||
|
||||
"""
|
||||
This example demonstrates how to use image transforms with LeRobot datasets for data augmentation during training.
|
||||
|
||||
Image transforms are applied to camera frames to improve model robustness and generalization. They are applied
|
||||
at training time only, not during dataset recording, allowing you to experiment with different augmentations
|
||||
without re-recording data.
|
||||
"""
|
||||
|
||||
import torch
|
||||
from torchvision.transforms import v2
|
||||
from torchvision.transforms.functional import to_pil_image
|
||||
|
||||
from lerobot.datasets.lerobot_dataset import LeRobotDataset
|
||||
from lerobot.datasets.transforms import ImageTransformConfig, ImageTransforms, ImageTransformsConfig
|
||||
|
||||
|
||||
def save_image(tensor, filename):
|
||||
"""Helper function to save a tensor as an image file."""
|
||||
if tensor.dim() == 3: # [C, H, W]
|
||||
if tensor.max() > 1.0:
|
||||
tensor = tensor / 255.0
|
||||
tensor = torch.clamp(tensor, 0.0, 1.0)
|
||||
pil_image = to_pil_image(tensor)
|
||||
pil_image.save(filename)
|
||||
print(f"Saved: {filename}")
|
||||
else:
|
||||
print(f"Skipped {filename}: unexpected tensor shape {tensor.shape}")
|
||||
|
||||
|
||||
def example_1_default_transforms():
|
||||
"""Example 1: Use default transform configuration and save original vs transformed images"""
|
||||
print("\n Example 1: Default Transform Configuration with Image Saving")
|
||||
|
||||
repo_id = "pepijn223/record_main_0" # Example dataset
|
||||
|
||||
try:
|
||||
# Load dataset without transforms (original)
|
||||
dataset_original = LeRobotDataset(repo_id=repo_id)
|
||||
|
||||
# Load dataset with transforms enabled
|
||||
transforms_config = ImageTransformsConfig(
|
||||
enable=True, # Enable transforms (disabled by default)
|
||||
max_num_transforms=2, # Apply up to 2 transforms per frame
|
||||
random_order=False, # Apply in standard order
|
||||
)
|
||||
dataset_with_transforms = LeRobotDataset(
|
||||
repo_id=repo_id, image_transforms=ImageTransforms(transforms_config)
|
||||
)
|
||||
|
||||
# Save original and transformed images for comparison
|
||||
if len(dataset_original) > 0:
|
||||
frame_idx = 0 # Use first frame
|
||||
original_sample = dataset_original[frame_idx]
|
||||
transformed_sample = dataset_with_transforms[frame_idx]
|
||||
|
||||
print(f"Saving comparison images (frame {frame_idx}):")
|
||||
|
||||
for cam_key in dataset_original.meta.camera_keys:
|
||||
if cam_key in original_sample and cam_key in transformed_sample:
|
||||
cam_name = cam_key.replace(".", "_").replace("/", "_")
|
||||
|
||||
# Save original and transformed images
|
||||
save_image(original_sample[cam_key], f"{cam_name}_original.png")
|
||||
save_image(transformed_sample[cam_key], f"{cam_name}_transformed.png")
|
||||
|
||||
except Exception as e:
|
||||
print(f"Could not load dataset '{repo_id}': {e}")
|
||||
|
||||
|
||||
def example_2_custom_transforms():
|
||||
"""Example 2: Create custom transform configuration and save examples"""
|
||||
print("\n Example 2: Custom Transform Configuration")
|
||||
|
||||
repo_id = "pepijn223/record_main_0" # Example dataset
|
||||
|
||||
try:
|
||||
# Create custom transform configuration with strong effects
|
||||
custom_transforms_config = ImageTransformsConfig(
|
||||
enable=True,
|
||||
max_num_transforms=2, # Apply up to 2 transforms per frame
|
||||
random_order=True, # Apply transforms in random order
|
||||
tfs={
|
||||
"brightness": ImageTransformConfig(
|
||||
weight=1.0,
|
||||
type="ColorJitter",
|
||||
kwargs={"brightness": (0.5, 1.5)}, # Strong brightness range
|
||||
),
|
||||
"contrast": ImageTransformConfig(
|
||||
weight=1.0, # Higher weight = more likely to be selected
|
||||
type="ColorJitter",
|
||||
kwargs={"contrast": (0.6, 1.4)}, # Strong contrast
|
||||
),
|
||||
"sharpness": ImageTransformConfig(
|
||||
weight=0.5, # Lower weight = less likely to be selected
|
||||
type="SharpnessJitter",
|
||||
kwargs={"sharpness": (0.2, 2.0)}, # Strong sharpness variation
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
dataset_with_custom_transforms = LeRobotDataset(
|
||||
repo_id=repo_id, image_transforms=ImageTransforms(custom_transforms_config)
|
||||
)
|
||||
|
||||
# Save examples with strong transforms
|
||||
if len(dataset_with_custom_transforms) > 0:
|
||||
sample = dataset_with_custom_transforms[0]
|
||||
print("Saving custom transform examples:")
|
||||
|
||||
for cam_key in dataset_with_custom_transforms.meta.camera_keys:
|
||||
if cam_key in sample:
|
||||
cam_name = cam_key.replace(".", "_").replace("/", "_")
|
||||
save_image(sample[cam_key], f"{cam_name}_custom_transforms.png")
|
||||
|
||||
except Exception as e:
|
||||
print(f"Could not load dataset '{repo_id}': {e}")
|
||||
|
||||
|
||||
def example_3_torchvision_transforms():
|
||||
"""Example 3: Use pure torchvision transforms and save examples"""
|
||||
print("\n Example 3: Pure Torchvision Transforms")
|
||||
|
||||
repo_id = "pepijn223/record_main_0" # Example dataset
|
||||
|
||||
try:
|
||||
# Create torchvision transform pipeline
|
||||
torchvision_transforms = v2.Compose(
|
||||
[
|
||||
v2.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3, hue=0.1),
|
||||
v2.GaussianBlur(kernel_size=3, sigma=(0.1, 2.0)),
|
||||
v2.RandomRotation(degrees=10), # Small rotation
|
||||
]
|
||||
)
|
||||
|
||||
dataset_with_torchvision = LeRobotDataset(repo_id=repo_id, image_transforms=torchvision_transforms)
|
||||
|
||||
# Save examples with torchvision transforms
|
||||
if len(dataset_with_torchvision) > 0:
|
||||
sample = dataset_with_torchvision[0]
|
||||
print("Saving torchvision transform examples:")
|
||||
|
||||
for cam_key in dataset_with_torchvision.meta.camera_keys:
|
||||
if cam_key in sample:
|
||||
cam_name = cam_key.replace(".", "_").replace("/", "_")
|
||||
save_image(sample[cam_key], f"{cam_name}_torchvision.png")
|
||||
|
||||
except Exception as e:
|
||||
print(f"Could not load dataset '{repo_id}': {e}")
|
||||
|
||||
|
||||
def main():
|
||||
"""Run all examples"""
|
||||
print("LeRobot Dataset Image Transforms Examples")
|
||||
|
||||
example_1_default_transforms()
|
||||
example_2_custom_transforms()
|
||||
example_3_torchvision_transforms()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -60,8 +60,6 @@ preprocessor, postprocessor = make_pre_post_processors(
|
||||
policy_cfg=policy,
|
||||
pretrained_path=HF_MODEL_ID,
|
||||
dataset_stats=dataset.meta.stats,
|
||||
# The inference device is automatically set to match the detected hardware, overriding any previous device settings from training to ensure compatibility.
|
||||
preprocessor_overrides={"device_processor": {"device": str(policy.config.device)}},
|
||||
)
|
||||
|
||||
# Connect the robot
|
||||
|
||||
@@ -26,7 +26,7 @@ from lerobot.utils.control_utils import init_keyboard_listener
|
||||
from lerobot.utils.utils import log_say
|
||||
from lerobot.utils.visualization_utils import _init_rerun
|
||||
|
||||
NUM_EPISODES = 2
|
||||
NUM_EPISODES = 3
|
||||
FPS = 30
|
||||
EPISODE_TIME_SEC = 30
|
||||
RESET_TIME_SEC = 10
|
||||
|
||||
@@ -32,9 +32,7 @@ robot = LeKiwiClient(robot_config)
|
||||
|
||||
# Fetch the dataset to replay
|
||||
dataset = LeRobotDataset("<hf_username>/<dataset_repo_id>", episodes=[EPISODE_IDX])
|
||||
# Filter dataset to only include frames from the specified episode since episodes are chunked in dataset V3.0
|
||||
episode_frames = dataset.hf_dataset.filter(lambda x: x["episode_index"] == EPISODE_IDX)
|
||||
actions = episode_frames.select_columns("action")
|
||||
actions = dataset.hf_dataset.select_columns("action")
|
||||
|
||||
# Connect to the robot
|
||||
robot.connect()
|
||||
@@ -44,7 +42,7 @@ if not robot.is_connected:
|
||||
|
||||
print("Starting replay loop...")
|
||||
log_say(f"Replaying episode {EPISODE_IDX}")
|
||||
for idx in range(len(episode_frames)):
|
||||
for idx in range(dataset.num_frames):
|
||||
t0 = time.perf_counter()
|
||||
|
||||
# Get recorded action from dataset
|
||||
|
||||
@@ -30,13 +30,14 @@ from lerobot.processor import (
|
||||
)
|
||||
from lerobot.processor.converters import (
|
||||
observation_to_transition,
|
||||
robot_action_observation_to_transition,
|
||||
robot_action_to_transition,
|
||||
transition_to_observation,
|
||||
transition_to_robot_action,
|
||||
)
|
||||
from lerobot.record import record_loop
|
||||
from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
|
||||
from lerobot.robots.so100_follower.robot_kinematic_processor import (
|
||||
AddRobotObservationAsComplimentaryData,
|
||||
ForwardKinematicsJointsToEE,
|
||||
InverseKinematicsEEToJoints,
|
||||
)
|
||||
@@ -68,21 +69,22 @@ policy = ACTPolicy.from_pretrained(HF_MODEL_ID)
|
||||
|
||||
# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
|
||||
kinematics_solver = RobotKinematics(
|
||||
urdf_path="./SO101/so101_new_calib.urdf",
|
||||
urdf_path="./src/lerobot/teleoperators/sim/so101_new_calib.urdf",
|
||||
target_frame_name="gripper_frame_link",
|
||||
joint_names=list(robot.bus.motors.keys()),
|
||||
)
|
||||
|
||||
# Build pipeline to convert EE action to joints action
|
||||
robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
|
||||
robot_ee_to_joints_processor = RobotProcessorPipeline[RobotAction, RobotAction](
|
||||
steps=[
|
||||
AddRobotObservationAsComplimentaryData(robot=robot),
|
||||
InverseKinematicsEEToJoints(
|
||||
kinematics=kinematics_solver,
|
||||
motor_names=list(robot.bus.motors.keys()),
|
||||
initial_guess_current_joints=True,
|
||||
),
|
||||
],
|
||||
to_transition=robot_action_observation_to_transition,
|
||||
to_transition=robot_action_to_transition,
|
||||
to_output=transition_to_robot_action,
|
||||
)
|
||||
|
||||
@@ -128,8 +130,6 @@ preprocessor, postprocessor = make_pre_post_processors(
|
||||
policy_cfg=policy,
|
||||
pretrained_path=HF_MODEL_ID,
|
||||
dataset_stats=dataset.meta.stats,
|
||||
# The inference device is automatically set to match the detected hardware, overriding any previous device settings from training to ensure compatibility.
|
||||
preprocessor_overrides={"device_processor": {"device": str(policy.config.device)}},
|
||||
)
|
||||
|
||||
# Connect the robot
|
||||
|
||||
@@ -22,13 +22,14 @@ from lerobot.model.kinematics import RobotKinematics
|
||||
from lerobot.processor import RobotAction, RobotObservation, RobotProcessorPipeline
|
||||
from lerobot.processor.converters import (
|
||||
observation_to_transition,
|
||||
robot_action_observation_to_transition,
|
||||
robot_action_to_transition,
|
||||
transition_to_observation,
|
||||
transition_to_robot_action,
|
||||
)
|
||||
from lerobot.record import record_loop
|
||||
from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
|
||||
from lerobot.robots.so100_follower.robot_kinematic_processor import (
|
||||
AddRobotObservationAsComplimentaryData,
|
||||
EEBoundsAndSafety,
|
||||
EEReferenceAndDelta,
|
||||
ForwardKinematicsJointsToEE,
|
||||
@@ -43,7 +44,7 @@ from lerobot.utils.control_utils import init_keyboard_listener
|
||||
from lerobot.utils.utils import log_say
|
||||
from lerobot.utils.visualization_utils import _init_rerun
|
||||
|
||||
NUM_EPISODES = 2
|
||||
NUM_EPISODES = 10
|
||||
FPS = 30
|
||||
EPISODE_TIME_SEC = 60
|
||||
RESET_TIME_SEC = 30
|
||||
@@ -53,7 +54,7 @@ HF_REPO_ID = "<hf_username>/<dataset_repo_id>"
|
||||
# Create the robot and teleoperator configurations
|
||||
camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)}
|
||||
robot_config = SO100FollowerConfig(
|
||||
port="/dev/tty.usbmodem5A460814411",
|
||||
port="/dev/tty.usbmodem58760434471",
|
||||
id="my_awesome_follower_arm",
|
||||
cameras=camera_config,
|
||||
use_degrees=True,
|
||||
@@ -66,34 +67,34 @@ phone = Phone(teleop_config)
|
||||
|
||||
# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
|
||||
kinematics_solver = RobotKinematics(
|
||||
urdf_path="./SO101/so101_new_calib.urdf",
|
||||
urdf_path="./src/lerobot/teleoperators/sim/so101_new_calib.urdf",
|
||||
target_frame_name="gripper_frame_link",
|
||||
joint_names=list(robot.bus.motors.keys()),
|
||||
)
|
||||
|
||||
# Build pipeline to convert phone action to EE action
|
||||
phone_to_robot_ee_pose_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
|
||||
phone_to_robot_ee_pose_processor = RobotProcessorPipeline[RobotAction, RobotAction](
|
||||
steps=[
|
||||
MapPhoneActionToRobotAction(platform=teleop_config.phone_os),
|
||||
AddRobotObservationAsComplimentaryData(robot=robot),
|
||||
EEReferenceAndDelta(
|
||||
kinematics=kinematics_solver,
|
||||
end_effector_step_sizes={"x": 0.5, "y": 0.5, "z": 0.5},
|
||||
motor_names=list(robot.bus.motors.keys()),
|
||||
use_latched_reference=True,
|
||||
),
|
||||
EEBoundsAndSafety(
|
||||
end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
|
||||
max_ee_step_m=0.20,
|
||||
max_ee_twist_step_rad=0.50,
|
||||
),
|
||||
GripperVelocityToJoint(speed_factor=20.0),
|
||||
GripperVelocityToJoint(),
|
||||
],
|
||||
to_transition=robot_action_observation_to_transition,
|
||||
to_transition=robot_action_to_transition,
|
||||
to_output=transition_to_robot_action,
|
||||
)
|
||||
|
||||
# Build pipeline to convert EE action to joints action
|
||||
robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
|
||||
robot_ee_to_joints_processor = RobotProcessorPipeline[RobotAction, RobotAction](
|
||||
steps=[
|
||||
InverseKinematicsEEToJoints(
|
||||
kinematics=kinematics_solver,
|
||||
@@ -101,7 +102,7 @@ robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotOb
|
||||
initial_guess_current_joints=True,
|
||||
),
|
||||
],
|
||||
to_transition=robot_action_observation_to_transition,
|
||||
to_transition=robot_action_to_transition,
|
||||
to_output=transition_to_robot_action,
|
||||
)
|
||||
|
||||
|
||||
@@ -18,13 +18,11 @@ import time
|
||||
|
||||
from lerobot.datasets.lerobot_dataset import LeRobotDataset
|
||||
from lerobot.model.kinematics import RobotKinematics
|
||||
from lerobot.processor import RobotAction, RobotObservation, RobotProcessorPipeline
|
||||
from lerobot.processor.converters import (
|
||||
robot_action_observation_to_transition,
|
||||
transition_to_robot_action,
|
||||
)
|
||||
from lerobot.processor import RobotAction, RobotProcessorPipeline
|
||||
from lerobot.processor.converters import robot_action_to_transition, transition_to_robot_action
|
||||
from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
|
||||
from lerobot.robots.so100_follower.robot_kinematic_processor import (
|
||||
AddRobotObservationAsComplimentaryData,
|
||||
InverseKinematicsEEToJoints,
|
||||
)
|
||||
from lerobot.robots.so100_follower.so100_follower import SO100Follower
|
||||
@@ -36,7 +34,7 @@ HF_REPO_ID = "<hf_username>/<dataset_repo_id>"
|
||||
|
||||
# Initialize the robot config
|
||||
robot_config = SO100FollowerConfig(
|
||||
port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True
|
||||
port="/dev/tty.usbmodem58760434471", id="my_awesome_follower_arm", use_degrees=True
|
||||
)
|
||||
|
||||
# Initialize the robot
|
||||
@@ -44,29 +42,28 @@ robot = SO100Follower(robot_config)
|
||||
|
||||
# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
|
||||
kinematics_solver = RobotKinematics(
|
||||
urdf_path="./SO101/so101_new_calib.urdf",
|
||||
urdf_path="./src/lerobot/teleoperators/sim/so101_new_calib.urdf",
|
||||
target_frame_name="gripper_frame_link",
|
||||
joint_names=list(robot.bus.motors.keys()),
|
||||
)
|
||||
|
||||
# Build pipeline to convert EE action to joints action
|
||||
robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
|
||||
robot_ee_to_joints_processor = RobotProcessorPipeline[RobotAction, RobotAction](
|
||||
steps=[
|
||||
AddRobotObservationAsComplimentaryData(robot=robot),
|
||||
InverseKinematicsEEToJoints(
|
||||
kinematics=kinematics_solver,
|
||||
motor_names=list(robot.bus.motors.keys()),
|
||||
initial_guess_current_joints=False, # Because replay is open loop
|
||||
),
|
||||
],
|
||||
to_transition=robot_action_observation_to_transition,
|
||||
to_transition=robot_action_to_transition,
|
||||
to_output=transition_to_robot_action,
|
||||
)
|
||||
|
||||
# Fetch the dataset to replay
|
||||
dataset = LeRobotDataset(HF_REPO_ID, episodes=[EPISODE_IDX])
|
||||
# Filter dataset to only include frames from the specified episode since episodes are chunked in dataset V3.0
|
||||
episode_frames = dataset.hf_dataset.filter(lambda x: x["episode_index"] == EPISODE_IDX)
|
||||
actions = episode_frames.select_columns("action")
|
||||
actions = dataset.hf_dataset.select_columns("action")
|
||||
|
||||
# Connect to the robot
|
||||
robot.connect()
|
||||
@@ -76,7 +73,7 @@ if not robot.is_connected:
|
||||
|
||||
print("Starting replay loop...")
|
||||
log_say(f"Replaying episode {EPISODE_IDX}")
|
||||
for idx in range(len(episode_frames)):
|
||||
for idx in range(dataset.num_frames):
|
||||
t0 = time.perf_counter()
|
||||
|
||||
# Get recorded action from dataset
|
||||
@@ -84,11 +81,8 @@ for idx in range(len(episode_frames)):
|
||||
name: float(actions[idx]["action"][i]) for i, name in enumerate(dataset.features["action"]["names"])
|
||||
}
|
||||
|
||||
# Get robot observation
|
||||
robot_obs = robot.get_observation()
|
||||
|
||||
# Dataset EE -> robot joints
|
||||
joint_action = robot_ee_to_joints_processor((ee_action, robot_obs))
|
||||
joint_action = robot_ee_to_joints_processor(ee_action)
|
||||
|
||||
# Send action to robot
|
||||
_ = robot.send_action(joint_action)
|
||||
|
||||
@@ -16,13 +16,11 @@
|
||||
import time
|
||||
|
||||
from lerobot.model.kinematics import RobotKinematics
|
||||
from lerobot.processor import RobotAction, RobotObservation, RobotProcessorPipeline
|
||||
from lerobot.processor.converters import (
|
||||
robot_action_observation_to_transition,
|
||||
transition_to_robot_action,
|
||||
)
|
||||
from lerobot.processor import RobotAction, RobotProcessorPipeline
|
||||
from lerobot.processor.converters import robot_action_to_transition, transition_to_robot_action
|
||||
from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
|
||||
from lerobot.robots.so100_follower.robot_kinematic_processor import (
|
||||
AddRobotObservationAsComplimentaryData,
|
||||
EEBoundsAndSafety,
|
||||
EEReferenceAndDelta,
|
||||
GripperVelocityToJoint,
|
||||
@@ -39,7 +37,7 @@ FPS = 30
|
||||
|
||||
# Initialize the robot and teleoperator
|
||||
robot_config = SO100FollowerConfig(
|
||||
port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True
|
||||
port="/dev/tty.usbmodem58760434471", id="my_awesome_follower_arm", use_degrees=True
|
||||
)
|
||||
teleop_config = PhoneConfig(phone_os=PhoneOS.IOS) # or PhoneOS.ANDROID
|
||||
|
||||
@@ -49,20 +47,20 @@ teleop_device = Phone(teleop_config)
|
||||
|
||||
# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
|
||||
kinematics_solver = RobotKinematics(
|
||||
urdf_path="./SO101/so101_new_calib.urdf",
|
||||
urdf_path="./src/lerobot/teleoperators/sim/so101_new_calib.urdf",
|
||||
target_frame_name="gripper_frame_link",
|
||||
joint_names=list(robot.bus.motors.keys()),
|
||||
)
|
||||
|
||||
# Build pipeline to convert phone action to ee pose action to joint action
|
||||
phone_to_robot_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
|
||||
phone_to_robot_joints_processor = RobotProcessorPipeline[RobotAction, RobotAction](
|
||||
steps=[
|
||||
MapPhoneActionToRobotAction(platform=teleop_config.phone_os),
|
||||
AddRobotObservationAsComplimentaryData(robot=robot),
|
||||
EEReferenceAndDelta(
|
||||
kinematics=kinematics_solver,
|
||||
end_effector_step_sizes={"x": 0.5, "y": 0.5, "z": 0.5},
|
||||
motor_names=list(robot.bus.motors.keys()),
|
||||
use_latched_reference=True,
|
||||
),
|
||||
EEBoundsAndSafety(
|
||||
end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
|
||||
@@ -75,10 +73,9 @@ phone_to_robot_joints_processor = RobotProcessorPipeline[tuple[RobotAction, Robo
|
||||
InverseKinematicsEEToJoints(
|
||||
kinematics=kinematics_solver,
|
||||
motor_names=list(robot.bus.motors.keys()),
|
||||
initial_guess_current_joints=True,
|
||||
),
|
||||
],
|
||||
to_transition=robot_action_observation_to_transition,
|
||||
to_transition=robot_action_to_transition,
|
||||
to_output=transition_to_robot_action,
|
||||
)
|
||||
|
||||
@@ -96,14 +93,11 @@ print("Starting teleop loop. Move your phone to teleoperate the robot...")
|
||||
while True:
|
||||
t0 = time.perf_counter()
|
||||
|
||||
# Get robot observation
|
||||
robot_obs = robot.get_observation()
|
||||
|
||||
# Get teleop action
|
||||
phone_obs = teleop_device.get_action()
|
||||
|
||||
# Phone -> EE pose -> Joints transition
|
||||
joint_action = phone_to_robot_joints_processor((phone_obs, robot_obs))
|
||||
joint_action = phone_to_robot_joints_processor(phone_obs)
|
||||
|
||||
# Send action to robot
|
||||
_ = robot.send_action(joint_action)
|
||||
|
||||
@@ -30,13 +30,14 @@ from lerobot.processor import (
|
||||
)
|
||||
from lerobot.processor.converters import (
|
||||
observation_to_transition,
|
||||
robot_action_observation_to_transition,
|
||||
robot_action_to_transition,
|
||||
transition_to_observation,
|
||||
transition_to_robot_action,
|
||||
)
|
||||
from lerobot.record import record_loop
|
||||
from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
|
||||
from lerobot.robots.so100_follower.robot_kinematic_processor import (
|
||||
AddRobotObservationAsComplimentaryData,
|
||||
ForwardKinematicsJointsToEE,
|
||||
InverseKinematicsEEToJoints,
|
||||
)
|
||||
@@ -68,21 +69,22 @@ policy = ACTPolicy.from_pretrained(HF_MODEL_ID)
|
||||
|
||||
# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
|
||||
kinematics_solver = RobotKinematics(
|
||||
urdf_path="./SO101/so101_new_calib.urdf",
|
||||
urdf_path="./src/lerobot/teleoperators/sim/so101_new_calib.urdf",
|
||||
target_frame_name="gripper_frame_link",
|
||||
joint_names=list(robot.bus.motors.keys()),
|
||||
)
|
||||
|
||||
# Build pipeline to convert EE action to joints action
|
||||
robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
|
||||
robot_ee_to_joints_processor = RobotProcessorPipeline[RobotAction, RobotAction](
|
||||
steps=[
|
||||
AddRobotObservationAsComplimentaryData(robot=robot),
|
||||
InverseKinematicsEEToJoints(
|
||||
kinematics=kinematics_solver,
|
||||
motor_names=list(robot.bus.motors.keys()),
|
||||
initial_guess_current_joints=True,
|
||||
),
|
||||
],
|
||||
to_transition=robot_action_observation_to_transition,
|
||||
to_transition=robot_action_to_transition,
|
||||
to_output=transition_to_robot_action,
|
||||
)
|
||||
|
||||
@@ -129,8 +131,6 @@ preprocessor, postprocessor = make_pre_post_processors(
|
||||
policy_cfg=policy,
|
||||
pretrained_path=HF_MODEL_ID,
|
||||
dataset_stats=dataset.meta.stats,
|
||||
# The inference device is automatically set to match the detected hardware, overriding any previous device settings from training to ensure compatibility.
|
||||
preprocessor_overrides={"device_processor": {"device": str(policy.config.device)}},
|
||||
)
|
||||
|
||||
# Connect the robot and teleoperator
|
||||
|
||||
@@ -23,13 +23,14 @@ from lerobot.model.kinematics import RobotKinematics
|
||||
from lerobot.processor import RobotAction, RobotObservation, RobotProcessorPipeline
|
||||
from lerobot.processor.converters import (
|
||||
observation_to_transition,
|
||||
robot_action_observation_to_transition,
|
||||
robot_action_to_transition,
|
||||
transition_to_observation,
|
||||
transition_to_robot_action,
|
||||
)
|
||||
from lerobot.record import record_loop
|
||||
from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
|
||||
from lerobot.robots.so100_follower.robot_kinematic_processor import (
|
||||
AddRobotObservationAsComplimentaryData,
|
||||
EEBoundsAndSafety,
|
||||
ForwardKinematicsJointsToEE,
|
||||
InverseKinematicsEEToJoints,
|
||||
@@ -41,7 +42,7 @@ from lerobot.utils.control_utils import init_keyboard_listener
|
||||
from lerobot.utils.utils import log_say
|
||||
from lerobot.utils.visualization_utils import _init_rerun
|
||||
|
||||
NUM_EPISODES = 2
|
||||
NUM_EPISODES = 10
|
||||
FPS = 30
|
||||
EPISODE_TIME_SEC = 60
|
||||
RESET_TIME_SEC = 30
|
||||
@@ -61,14 +62,14 @@ leader = SO100Leader(leader_config)
|
||||
|
||||
# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
|
||||
follower_kinematics_solver = RobotKinematics(
|
||||
urdf_path="./SO101/so101_new_calib.urdf",
|
||||
urdf_path="./examples/phone_to_so100/SO101/so101_new_calib.urdf",
|
||||
target_frame_name="gripper_frame_link",
|
||||
joint_names=list(follower.bus.motors.keys()),
|
||||
)
|
||||
|
||||
# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
|
||||
leader_kinematics_solver = RobotKinematics(
|
||||
urdf_path="./SO101/so101_new_calib.urdf",
|
||||
urdf_path="./examples/phone_to_so100/SO101/so101_new_calib.urdf",
|
||||
target_frame_name="gripper_frame_link",
|
||||
joint_names=list(leader.bus.motors.keys()),
|
||||
)
|
||||
@@ -85,19 +86,20 @@ follower_joints_to_ee = RobotProcessorPipeline[RobotObservation, RobotObservatio
|
||||
)
|
||||
|
||||
# Build pipeline to convert leader joints to EE action
|
||||
leader_joints_to_ee = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
|
||||
leader_joints_to_ee = RobotProcessorPipeline[RobotAction, RobotAction](
|
||||
steps=[
|
||||
ForwardKinematicsJointsToEE(
|
||||
kinematics=leader_kinematics_solver, motor_names=list(leader.bus.motors.keys())
|
||||
),
|
||||
],
|
||||
to_transition=robot_action_observation_to_transition,
|
||||
to_transition=robot_action_to_transition,
|
||||
to_output=transition_to_robot_action,
|
||||
)
|
||||
|
||||
# Build pipeline to convert EE action to follower joints
|
||||
ee_to_follower_joints = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
|
||||
ee_to_follower_joints = RobotProcessorPipeline[RobotAction, RobotAction](
|
||||
[
|
||||
AddRobotObservationAsComplimentaryData(robot=follower),
|
||||
EEBoundsAndSafety(
|
||||
end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
|
||||
max_ee_step_m=0.10,
|
||||
@@ -106,10 +108,9 @@ ee_to_follower_joints = RobotProcessorPipeline[tuple[RobotAction, RobotObservati
|
||||
InverseKinematicsEEToJoints(
|
||||
kinematics=follower_kinematics_solver,
|
||||
motor_names=list(follower.bus.motors.keys()),
|
||||
initial_guess_current_joints=True,
|
||||
),
|
||||
],
|
||||
to_transition=robot_action_observation_to_transition,
|
||||
to_transition=robot_action_to_transition,
|
||||
to_output=transition_to_robot_action,
|
||||
)
|
||||
|
||||
|
||||
@@ -19,13 +19,11 @@ import time
|
||||
|
||||
from lerobot.datasets.lerobot_dataset import LeRobotDataset
|
||||
from lerobot.model.kinematics import RobotKinematics
|
||||
from lerobot.processor import RobotAction, RobotObservation, RobotProcessorPipeline
|
||||
from lerobot.processor.converters import (
|
||||
robot_action_observation_to_transition,
|
||||
transition_to_robot_action,
|
||||
)
|
||||
from lerobot.processor import RobotAction, RobotProcessorPipeline
|
||||
from lerobot.processor.converters import robot_action_to_transition, transition_to_robot_action
|
||||
from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
|
||||
from lerobot.robots.so100_follower.robot_kinematic_processor import (
|
||||
AddRobotObservationAsComplimentaryData,
|
||||
InverseKinematicsEEToJoints,
|
||||
)
|
||||
from lerobot.robots.so100_follower.so100_follower import SO100Follower
|
||||
@@ -45,29 +43,28 @@ robot = SO100Follower(robot_config)
|
||||
|
||||
# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
|
||||
kinematics_solver = RobotKinematics(
|
||||
urdf_path="./SO101/so101_new_calib.urdf",
|
||||
urdf_path="./src/lerobot/teleoperators/sim/so101_new_calib.urdf",
|
||||
target_frame_name="gripper_frame_link",
|
||||
joint_names=list(robot.bus.motors.keys()),
|
||||
)
|
||||
|
||||
# Build pipeline to convert EE action to joints action
|
||||
robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
|
||||
robot_ee_to_joints_processor = RobotProcessorPipeline[RobotAction, RobotAction](
|
||||
steps=[
|
||||
AddRobotObservationAsComplimentaryData(robot=robot),
|
||||
InverseKinematicsEEToJoints(
|
||||
kinematics=kinematics_solver,
|
||||
motor_names=list(robot.bus.motors.keys()),
|
||||
initial_guess_current_joints=False, # Because replay is open loop
|
||||
),
|
||||
],
|
||||
to_transition=robot_action_observation_to_transition,
|
||||
to_transition=robot_action_to_transition,
|
||||
to_output=transition_to_robot_action,
|
||||
)
|
||||
|
||||
# Fetch the dataset to replay
|
||||
dataset = LeRobotDataset(HF_REPO_ID, episodes=[EPISODE_IDX])
|
||||
# Filter dataset to only include frames from the specified episode since episodes are chunked in dataset V3.0
|
||||
episode_frames = dataset.hf_dataset.filter(lambda x: x["episode_index"] == EPISODE_IDX)
|
||||
actions = episode_frames.select_columns("action")
|
||||
actions = dataset.hf_dataset.select_columns("action")
|
||||
|
||||
# Connect to the robot
|
||||
robot.connect()
|
||||
@@ -77,7 +74,7 @@ if not robot.is_connected:
|
||||
|
||||
print("Starting replay loop...")
|
||||
log_say(f"Replaying episode {EPISODE_IDX}")
|
||||
for idx in range(len(episode_frames)):
|
||||
for idx in range(dataset.num_frames):
|
||||
t0 = time.perf_counter()
|
||||
|
||||
# Get recorded action from dataset
|
||||
@@ -85,11 +82,8 @@ for idx in range(len(episode_frames)):
|
||||
name: float(actions[idx]["action"][i]) for i, name in enumerate(dataset.features["action"]["names"])
|
||||
}
|
||||
|
||||
# Get robot observation
|
||||
robot_obs = robot.get_observation()
|
||||
|
||||
# Dataset EE -> robot joints
|
||||
joint_action = robot_ee_to_joints_processor((ee_action, robot_obs))
|
||||
joint_action = robot_ee_to_joints_processor(ee_action)
|
||||
|
||||
# Send action to robot
|
||||
_ = robot.send_action(joint_action)
|
||||
|
||||
@@ -17,14 +17,14 @@
|
||||
import time
|
||||
|
||||
from lerobot.model.kinematics import RobotKinematics
|
||||
from lerobot.processor import RobotAction, RobotObservation, RobotProcessorPipeline
|
||||
from lerobot.processor import RobotAction, RobotProcessorPipeline
|
||||
from lerobot.processor.converters import (
|
||||
robot_action_observation_to_transition,
|
||||
robot_action_to_transition,
|
||||
transition_to_robot_action,
|
||||
)
|
||||
from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
|
||||
from lerobot.robots.so100_follower.robot_kinematic_processor import (
|
||||
AddRobotObservationAsComplimentaryData,
|
||||
EEBoundsAndSafety,
|
||||
ForwardKinematicsJointsToEE,
|
||||
InverseKinematicsEEToJoints,
|
||||
@@ -49,14 +49,14 @@ leader = SO100Leader(leader_config)
|
||||
|
||||
# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
|
||||
follower_kinematics_solver = RobotKinematics(
|
||||
urdf_path="./SO101/so101_new_calib.urdf",
|
||||
urdf_path="./examples/phone_to_so100/SO101/so101_new_calib.urdf",
|
||||
target_frame_name="gripper_frame_link",
|
||||
joint_names=list(follower.bus.motors.keys()),
|
||||
)
|
||||
|
||||
# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
|
||||
leader_kinematics_solver = RobotKinematics(
|
||||
urdf_path="./SO101/so101_new_calib.urdf",
|
||||
urdf_path="./examples/phone_to_so100/SO101/so101_new_calib.urdf",
|
||||
target_frame_name="gripper_frame_link",
|
||||
joint_names=list(leader.bus.motors.keys()),
|
||||
)
|
||||
@@ -73,8 +73,9 @@ leader_to_ee = RobotProcessorPipeline[RobotAction, RobotAction](
|
||||
)
|
||||
|
||||
# build pipeline to convert EE action to robot joints
|
||||
ee_to_follower_joints = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
|
||||
ee_to_follower_joints = RobotProcessorPipeline[RobotAction, RobotAction](
|
||||
[
|
||||
AddRobotObservationAsComplimentaryData(robot=follower),
|
||||
EEBoundsAndSafety(
|
||||
end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
|
||||
max_ee_step_m=0.10,
|
||||
@@ -83,10 +84,9 @@ ee_to_follower_joints = RobotProcessorPipeline[tuple[RobotAction, RobotObservati
|
||||
InverseKinematicsEEToJoints(
|
||||
kinematics=follower_kinematics_solver,
|
||||
motor_names=list(follower.bus.motors.keys()),
|
||||
initial_guess_current_joints=False,
|
||||
),
|
||||
],
|
||||
to_transition=robot_action_observation_to_transition,
|
||||
to_transition=robot_action_to_transition,
|
||||
to_output=transition_to_robot_action,
|
||||
)
|
||||
|
||||
@@ -101,9 +101,6 @@ print("Starting teleop loop...")
|
||||
while True:
|
||||
t0 = time.perf_counter()
|
||||
|
||||
# Get robot observation
|
||||
robot_obs = follower.get_observation()
|
||||
|
||||
# Get teleop observation
|
||||
leader_joints_obs = leader.get_action()
|
||||
|
||||
@@ -111,7 +108,7 @@ while True:
|
||||
leader_ee_act = leader_to_ee(leader_joints_obs)
|
||||
|
||||
# teleop EE -> robot joints
|
||||
follower_joints_act = ee_to_follower_joints((leader_ee_act, robot_obs))
|
||||
follower_joints_act = ee_to_follower_joints(leader_ee_act)
|
||||
|
||||
# Send action to robot
|
||||
_ = follower.send_action(follower_joints_act)
|
||||
|
||||
@@ -0,0 +1,58 @@
|
||||
#!/bin/bash
|
||||
|
||||
# storage / caches
|
||||
RAID=/raid/jade
|
||||
export TRANSFORMERS_CACHE=$RAID/.cache/huggingface/transformers
|
||||
export HF_HOME=$RAID/.cache/huggingface
|
||||
export HF_DATASETS_CACHE=$RAID/.cache/huggingface/datasets
|
||||
export HF_LEROBOT_HOME=$RAID/.cache/huggingface/lerobot
|
||||
export WANDB_CACHE_DIR=$RAID/.cache/wandb
|
||||
export TMPDIR=$RAID/.cache/tmp
|
||||
mkdir -p $TMPDIR
|
||||
export WANDB_MODE=offline
|
||||
export HF_DATASETS_OFFLINE=1
|
||||
export HF_HUB_OFFLINE=1
|
||||
export TOKENIZERS_PARALLELISM=false
|
||||
export MUJOCO_GL=egl
|
||||
export CUDA_VISIBLE_DEVICES=2
|
||||
|
||||
# CONFIGURATION
|
||||
POLICY_PATH="/raid/jade/logs/lerobot/lerobot_2_HuggingFaceVLA_libero_smolvla_lr1e-4bs32steps100000/checkpoints/100000/pretrained_model"
|
||||
POLICY_PATH="/raid/jade/models/smolvla_pipe"
|
||||
TASK=libero_spatial
|
||||
ENV_TYPE="libero"
|
||||
BATCH_SIZE=1
|
||||
N_EPISODES=1
|
||||
# storage / caches
|
||||
RAID=/raid/jade
|
||||
N_ACTION_STEPS=1
|
||||
export TRANSFORMERS_CACHE=$RAID/.cache/huggingface/transformers
|
||||
export HF_HOME=$RAID/.cache/huggingface
|
||||
export HF_DATASETS_CACHE=$RAID/.cache/huggingface/datasets
|
||||
export HF_LEROBOT_HOME=$RAID/.cache/huggingface/lerobot
|
||||
export WANDB_CACHE_DIR=$RAID/.cache/wandb
|
||||
export TMPDIR=$RAID/.cache/tmp
|
||||
mkdir -p $TMPDIR
|
||||
export WANDB_MODE=offline
|
||||
# export HF_DATASETS_OFFLINE=1
|
||||
# export HF_HUB_OFFLINE=1
|
||||
export TOKENIZERS_PARALLELISM=false
|
||||
export MUJOCO_GL=egl
|
||||
export MUJOCO_GL=egl
|
||||
unset HF_HUB_OFFLINE
|
||||
# RUN EVALUATION
|
||||
python src/lerobot/scripts/eval.py \
|
||||
--policy.path="$POLICY_PATH" \
|
||||
--env.type="$ENV_TYPE" \
|
||||
--eval.batch_size="$BATCH_SIZE" \
|
||||
--eval.n_episodes="$N_EPISODES" \
|
||||
--env.task=$TASK \
|
||||
--env.max_parallel_tasks=10 \
|
||||
# python examples/evaluate_libero.py \
|
||||
# --policy_path "$POLICY_PATH" \
|
||||
# --task_suite_name "$TASK" \
|
||||
# --num_steps_wait 10 \
|
||||
# --num_trials_per_task 10 \
|
||||
# --video_out_path "data/libero/videos" \
|
||||
# --device "cuda" \
|
||||
# --seed 7
|
||||
+18
-9
@@ -59,7 +59,7 @@ keywords = ["lerobot", "huggingface", "robotics", "machine learning", "artifici
|
||||
dependencies = [
|
||||
|
||||
# Hugging Face dependencies
|
||||
"datasets>=4.0.0",
|
||||
"datasets>=2.19.0,<=3.6.0", # TODO: Bumb dependency
|
||||
"diffusers>=0.27.2",
|
||||
"huggingface-hub[hf-transfer,cli]>=0.34.2",
|
||||
|
||||
@@ -121,7 +121,7 @@ phone = ["hebi-py>=2.8.0", "teleop>=0.1.0"]
|
||||
# Policies
|
||||
pi0 = ["lerobot[transformers-dep]"]
|
||||
smolvla = ["lerobot[transformers-dep]", "num2words>=0.5.14", "accelerate>=1.7.0", "safetensors>=0.4.3"]
|
||||
hilserl = ["lerobot[transformers-dep]", "gym-hil>=0.1.11", "lerobot[grpcio-dep]", "lerobot[placo-dep]"]
|
||||
hilserl = ["lerobot[transformers-dep]", "gym-hil>=0.1.9", "lerobot[grpcio-dep]", "lerobot[placo-dep]"]
|
||||
|
||||
# Features
|
||||
async = ["lerobot[grpcio-dep]", "matplotlib>=3.10.3"]
|
||||
@@ -135,9 +135,21 @@ video_benchmark = ["scikit-image>=0.23.2", "pandas>=2.2.2"]
|
||||
aloha = ["gym-aloha>=0.1.1"]
|
||||
pusht = ["gym-pusht>=0.1.5", "pymunk>=6.6.0,<7.0.0"] # TODO: Fix pymunk version in gym-pusht instead
|
||||
xarm = ["gym-xarm>=0.1.1"]
|
||||
libero = ["lerobot[transformers-dep]", "libero @ git+https://github.com/huggingface/lerobot-libero.git@main#egg=libero"]
|
||||
|
||||
|
||||
libero = [
|
||||
"hydra-core>=1.2,<1.4",
|
||||
"easydict>=1.9",
|
||||
"lerobot[transformers-dep]",
|
||||
"robomimic==0.2.0",
|
||||
"thop>=0.1.0.post2206102148",
|
||||
"robosuite==1.4.0",
|
||||
"bddl==1.0.1",
|
||||
"matplotlib>=3.5.3",
|
||||
"cloudpickle>=2.0.0",
|
||||
"gym>=0.25,<0.27",
|
||||
"future>=0.18.3",
|
||||
"egl_probe @ git+https://github.com/jadechoghari/egl_probe.git#egg=egl_probe",
|
||||
"libero @ git+https://github.com/jadechoghari/LIBERO.git@main#egg=libero",
|
||||
]
|
||||
# All
|
||||
all = [
|
||||
"lerobot[dynamixel]",
|
||||
@@ -158,7 +170,7 @@ all = [
|
||||
"lerobot[pusht]",
|
||||
"lerobot[xarm]",
|
||||
"lerobot[phone]",
|
||||
"lerobot[libero]",
|
||||
"lerobot[libero]"
|
||||
]
|
||||
|
||||
[project.scripts]
|
||||
@@ -171,9 +183,6 @@ lerobot-setup-motors="lerobot.setup_motors:main"
|
||||
lerobot-teleoperate="lerobot.teleoperate:main"
|
||||
lerobot-eval="lerobot.scripts.eval:main"
|
||||
lerobot-train="lerobot.scripts.train:main"
|
||||
lerobot-dataset-viz="lerobot.scripts.lerobot_dataset_viz:main"
|
||||
lerobot-info="lerobot.scripts.lerobot_info:main"
|
||||
lerobot-imgtransform-viz="lerobot.scripts.lerobot_imgtransform_viz:main"
|
||||
|
||||
# ---------------- Tool Configurations ----------------
|
||||
[tool.setuptools.packages.find]
|
||||
|
||||
@@ -196,10 +196,11 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
|
||||
config = json.load(f)
|
||||
|
||||
config.pop("type")
|
||||
with tempfile.NamedTemporaryFile("w+", delete=False, suffix=".json") as f:
|
||||
with tempfile.NamedTemporaryFile("w+") as f:
|
||||
json.dump(config, f)
|
||||
config_file = f.name
|
||||
f.flush()
|
||||
|
||||
cli_overrides = policy_kwargs.pop("cli_overrides", [])
|
||||
with draccus.config_type("json"):
|
||||
return draccus.parse(orig_config.__class__, config_file, args=cli_overrides)
|
||||
cli_overrides = policy_kwargs.pop("cli_overrides", [])
|
||||
with draccus.config_type("json"):
|
||||
return draccus.parse(orig_config.__class__, config_file, args=cli_overrides)
|
||||
|
||||
@@ -404,7 +404,7 @@ def convert_info(root, new_root, data_file_size_in_mb, video_file_size_in_mb):
|
||||
info["video_files_size_in_mb"] = video_file_size_in_mb
|
||||
info["data_path"] = DEFAULT_DATA_PATH
|
||||
info["video_path"] = DEFAULT_VIDEO_PATH
|
||||
info["fps"] = int(info["fps"])
|
||||
info["fps"] = float(info["fps"])
|
||||
for key in info["features"]:
|
||||
if info["features"][key]["dtype"] == "video":
|
||||
# already has fps in video_info
|
||||
|
||||
+87
-59
@@ -31,7 +31,6 @@ class EnvConfig(draccus.ChoiceRegistry, abc.ABC):
|
||||
features: dict[str, PolicyFeature] = field(default_factory=dict)
|
||||
features_map: dict[str, str] = field(default_factory=dict)
|
||||
max_parallel_tasks: int = 1
|
||||
disable_env_checker: bool = True
|
||||
|
||||
@property
|
||||
def type(self) -> str:
|
||||
@@ -163,71 +162,33 @@ class XarmEnv(EnvConfig):
|
||||
|
||||
|
||||
@dataclass
|
||||
class ImagePreprocessingConfig:
|
||||
crop_params_dict: dict[str, tuple[int, int, int, int]] | None = None
|
||||
resize_size: tuple[int, int] | None = None
|
||||
class VideoRecordConfig:
|
||||
"""Configuration for video recording in ManiSkill environments."""
|
||||
|
||||
enabled: bool = False
|
||||
record_dir: str = "videos"
|
||||
trajectory_name: str = "trajectory"
|
||||
|
||||
|
||||
@dataclass
|
||||
class RewardClassifierConfig:
|
||||
"""Configuration for reward classification."""
|
||||
|
||||
pretrained_path: str | None = None
|
||||
success_threshold: float = 0.5
|
||||
success_reward: float = 1.0
|
||||
|
||||
|
||||
@dataclass
|
||||
class InverseKinematicsConfig:
|
||||
"""Configuration for inverse kinematics processing."""
|
||||
|
||||
urdf_path: str | None = None
|
||||
target_frame_name: str | None = None
|
||||
end_effector_bounds: dict[str, list[float]] | None = None
|
||||
end_effector_step_sizes: dict[str, float] | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class ObservationConfig:
|
||||
"""Configuration for observation processing."""
|
||||
class EnvTransformConfig:
|
||||
"""Configuration for environment wrappers."""
|
||||
|
||||
# ee_action_space_params: EEActionSpaceConfig = field(default_factory=EEActionSpaceConfig)
|
||||
control_mode: str = "gamepad"
|
||||
display_cameras: bool = False
|
||||
add_joint_velocity_to_observation: bool = False
|
||||
add_current_to_observation: bool = False
|
||||
add_ee_pose_to_observation: bool = False
|
||||
display_cameras: bool = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class GripperConfig:
|
||||
"""Configuration for gripper control and penalties."""
|
||||
|
||||
use_gripper: bool = True
|
||||
gripper_penalty: float = 0.0
|
||||
gripper_penalty_in_reward: bool = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class ResetConfig:
|
||||
"""Configuration for environment reset behavior."""
|
||||
|
||||
crop_params_dict: dict[str, tuple[int, int, int, int]] | None = None
|
||||
resize_size: tuple[int, int] | None = None
|
||||
control_time_s: float = 20.0
|
||||
fixed_reset_joint_positions: Any | None = None
|
||||
reset_time_s: float = 5.0
|
||||
control_time_s: float = 20.0
|
||||
terminate_on_success: bool = True
|
||||
|
||||
|
||||
@dataclass
|
||||
class HILSerlProcessorConfig:
|
||||
"""Configuration for environment processing pipeline."""
|
||||
|
||||
control_mode: str = "gamepad"
|
||||
observation: ObservationConfig | None = None
|
||||
image_preprocessing: ImagePreprocessingConfig | None = None
|
||||
gripper: GripperConfig | None = None
|
||||
reset: ResetConfig | None = None
|
||||
inverse_kinematics: InverseKinematicsConfig | None = None
|
||||
reward_classifier: RewardClassifierConfig | None = None
|
||||
max_gripper_pos: float | None = 100.0
|
||||
use_gripper: bool = True
|
||||
gripper_quantization_threshold: float | None = 0.8
|
||||
gripper_penalty: float = 0.0
|
||||
gripper_penalty_in_reward: bool = False
|
||||
|
||||
|
||||
@EnvConfig.register_subclass(name="gym_manipulator")
|
||||
@@ -237,15 +198,82 @@ class HILSerlRobotEnvConfig(EnvConfig):
|
||||
|
||||
robot: RobotConfig | None = None
|
||||
teleop: TeleoperatorConfig | None = None
|
||||
processor: HILSerlProcessorConfig = field(default_factory=HILSerlProcessorConfig)
|
||||
|
||||
wrapper: EnvTransformConfig | None = None
|
||||
fps: int = 10
|
||||
name: str = "real_robot"
|
||||
mode: str | None = None # Either "record", "replay", None
|
||||
repo_id: str | None = None
|
||||
dataset_root: str | None = None
|
||||
task: str | None = ""
|
||||
num_episodes: int = 10 # only for record mode
|
||||
episode: int = 0
|
||||
device: str = "cuda"
|
||||
push_to_hub: bool = True
|
||||
pretrained_policy_name_or_path: str | None = None
|
||||
reward_classifier_pretrained_path: str | None = None
|
||||
# For the reward classifier, to record more positive examples after a success
|
||||
number_of_steps_after_success: int = 0
|
||||
|
||||
@property
|
||||
def gym_kwargs(self) -> dict:
|
||||
return {}
|
||||
|
||||
|
||||
@EnvConfig.register_subclass("hil")
|
||||
@dataclass
|
||||
class HILEnvConfig(EnvConfig):
|
||||
"""Configuration for the HIL environment."""
|
||||
|
||||
name: str = "PandaPickCube"
|
||||
task: str | None = "PandaPickCubeKeyboard-v0"
|
||||
use_viewer: bool = True
|
||||
gripper_penalty: float = 0.0
|
||||
use_gamepad: bool = True
|
||||
state_dim: int = 18
|
||||
action_dim: int = 4
|
||||
fps: int = 100
|
||||
episode_length: int = 100
|
||||
video_record: VideoRecordConfig = field(default_factory=VideoRecordConfig)
|
||||
features: dict[str, PolicyFeature] = field(
|
||||
default_factory=lambda: {
|
||||
"action": PolicyFeature(type=FeatureType.ACTION, shape=(4,)),
|
||||
"observation.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 128, 128)),
|
||||
"observation.state": PolicyFeature(type=FeatureType.STATE, shape=(18,)),
|
||||
}
|
||||
)
|
||||
features_map: dict[str, str] = field(
|
||||
default_factory=lambda: {
|
||||
"action": ACTION,
|
||||
"observation.image": OBS_IMAGE,
|
||||
"observation.state": OBS_STATE,
|
||||
}
|
||||
)
|
||||
################# args from hilserlrobotenv
|
||||
reward_classifier_pretrained_path: str | None = None
|
||||
robot_config: RobotConfig | None = None
|
||||
teleop_config: TeleoperatorConfig | None = None
|
||||
wrapper: EnvTransformConfig | None = None
|
||||
mode: str | None = None # Either "record", "replay", None
|
||||
repo_id: str | None = None
|
||||
dataset_root: str | None = None
|
||||
num_episodes: int = 10 # only for record mode
|
||||
episode: int = 0
|
||||
device: str = "cuda"
|
||||
push_to_hub: bool = True
|
||||
pretrained_policy_name_or_path: str | None = None
|
||||
# For the reward classifier, to record more positive examples after a success
|
||||
number_of_steps_after_success: int = 0
|
||||
############################
|
||||
|
||||
@property
|
||||
def gym_kwargs(self) -> dict:
|
||||
return {
|
||||
"use_viewer": self.use_viewer,
|
||||
"use_gamepad": self.use_gamepad,
|
||||
"gripper_penalty": self.gripper_penalty,
|
||||
}
|
||||
|
||||
|
||||
@EnvConfig.register_subclass("libero")
|
||||
@dataclass
|
||||
class LiberoEnv(EnvConfig):
|
||||
|
||||
@@ -17,7 +17,7 @@ import importlib
|
||||
|
||||
import gymnasium as gym
|
||||
|
||||
from lerobot.envs.configs import AlohaEnv, EnvConfig, LiberoEnv, PushtEnv, XarmEnv
|
||||
from lerobot.envs.configs import AlohaEnv, EnvConfig, HILEnvConfig, LiberoEnv, PushtEnv, XarmEnv
|
||||
|
||||
|
||||
def make_env_config(env_type: str, **kwargs) -> EnvConfig:
|
||||
@@ -27,6 +27,8 @@ def make_env_config(env_type: str, **kwargs) -> EnvConfig:
|
||||
return PushtEnv(**kwargs)
|
||||
elif env_type == "xarm":
|
||||
return XarmEnv(**kwargs)
|
||||
elif env_type == "hil":
|
||||
return HILEnvConfig(**kwargs)
|
||||
elif env_type == "libero":
|
||||
return LiberoEnv(**kwargs)
|
||||
else:
|
||||
@@ -76,13 +78,14 @@ def make_env(
|
||||
try:
|
||||
importlib.import_module(package_name)
|
||||
except ModuleNotFoundError as e:
|
||||
print(f"{package_name} is not installed. Please install it with `pip install 'lerobot[{cfg.type}]'`")
|
||||
raise e
|
||||
raise ModuleNotFoundError(
|
||||
f'{package_name} is not installed. Install with: pip install "lerobot[{cfg.type}]"'
|
||||
) from e
|
||||
|
||||
gym_handle = f"{package_name}/{cfg.task}"
|
||||
|
||||
def _make_one():
|
||||
return gym.make(gym_handle, disable_env_checker=cfg.disable_env_checker, **(cfg.gym_kwargs or {}))
|
||||
return gym.make(gym_handle, disable_env_checker=True, **(cfg.gym_kwargs or {}))
|
||||
|
||||
vec = env_cls([_make_one for _ in range(n_envs)])
|
||||
|
||||
|
||||
+36
-14
@@ -15,6 +15,7 @@
|
||||
# limitations under the License.
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import os
|
||||
from collections import defaultdict
|
||||
from collections.abc import Callable, Iterable, Mapping, Sequence
|
||||
@@ -28,7 +29,6 @@ import torch
|
||||
from gymnasium import spaces
|
||||
from libero.libero import benchmark, get_libero_path
|
||||
from libero.libero.envs import OffScreenRenderEnv
|
||||
from robosuite.utils.transform_utils import quat2axisangle
|
||||
|
||||
|
||||
def _parse_camera_names(camera_name: str | Sequence[str]) -> list[str]:
|
||||
@@ -44,7 +44,7 @@ def _parse_camera_names(camera_name: str | Sequence[str]) -> list[str]:
|
||||
return cams
|
||||
|
||||
|
||||
def _get_suite(name: str) -> benchmark.Benchmark:
|
||||
def _get_suite(name: str) -> Any:
|
||||
"""Instantiate a LIBERO suite by name with clear validation."""
|
||||
bench = benchmark.get_benchmark_dict()
|
||||
if name not in bench:
|
||||
@@ -66,6 +66,33 @@ def _select_task_ids(total_tasks: int, task_ids: Iterable[int] | None) -> list[i
|
||||
return ids
|
||||
|
||||
|
||||
def quat2axisangle(quat: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
Copied from robosuite: https://github.com/ARISE-Initiative/robosuite/blob/eafb81f54ffc104f905ee48a16bb15f059176ad3/robosuite/utils/transform_utils.py#L490C1-L512C55
|
||||
|
||||
Converts quaternion to axis-angle format.
|
||||
Returns a unit vector direction scaled by its angle in radians.
|
||||
|
||||
Args:
|
||||
quat (np.array): (x,y,z,w) vec4 float angles
|
||||
|
||||
Returns:
|
||||
np.array: (ax,ay,az) axis-angle exponential coordinates
|
||||
"""
|
||||
# clip quaternion
|
||||
if quat[3] > 1.0:
|
||||
quat[3] = 1.0
|
||||
elif quat[3] < -1.0:
|
||||
quat[3] = -1.0
|
||||
|
||||
den = np.sqrt(1.0 - quat[3] * quat[3])
|
||||
if math.isclose(den, 0.0):
|
||||
# This is (close to) a zero degree rotation, immediately return
|
||||
return np.zeros(3)
|
||||
|
||||
return (quat[:3] * 2.0 * math.acos(quat[3])) / den
|
||||
|
||||
|
||||
def get_task_init_states(task_suite: Any, i: int) -> np.ndarray:
|
||||
init_states_path = (
|
||||
Path(get_libero_path("init_states"))
|
||||
@@ -83,10 +110,6 @@ def get_libero_dummy_action():
|
||||
|
||||
OBS_STATE_DIM = 8
|
||||
ACTION_DIM = 7
|
||||
AGENT_POS_LOW = -1000.0
|
||||
AGENT_POS_HIGH = 1000.0
|
||||
ACTION_LOW = -1.0
|
||||
ACTION_HIGH = 1.0
|
||||
TASK_SUITE_MAX_STEPS: dict[str, int] = {
|
||||
"libero_spatial": 280, # longest training demo has 193 steps
|
||||
"libero_object": 280, # longest training demo has 254 steps
|
||||
@@ -125,8 +148,8 @@ class LiberoEnv(gym.Env):
|
||||
self.visualization_width = visualization_width
|
||||
self.visualization_height = visualization_height
|
||||
self.init_states = init_states
|
||||
self.camera_name = _parse_camera_names(
|
||||
camera_name
|
||||
self.camera_name = camera_name.split(
|
||||
","
|
||||
) # agentview_image (main) or robot0_eye_in_hand_image (wrist)
|
||||
|
||||
# Map raw camera names to "image1" and "image2".
|
||||
@@ -176,17 +199,15 @@ class LiberoEnv(gym.Env):
|
||||
{
|
||||
"pixels": spaces.Dict(images),
|
||||
"agent_pos": spaces.Box(
|
||||
low=AGENT_POS_LOW,
|
||||
high=AGENT_POS_HIGH,
|
||||
low=-1000.0,
|
||||
high=1000.0,
|
||||
shape=(OBS_STATE_DIM,),
|
||||
dtype=np.float64,
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
self.action_space = spaces.Box(
|
||||
low=ACTION_LOW, high=ACTION_HIGH, shape=(ACTION_DIM,), dtype=np.float32
|
||||
)
|
||||
self.action_space = spaces.Box(low=-1, high=1, shape=(ACTION_DIM,), dtype=np.float32)
|
||||
|
||||
def render(self):
|
||||
raw_obs = self._env.env._get_observations()
|
||||
@@ -291,6 +312,7 @@ def _make_env_fns(
|
||||
gym_kwargs: Mapping[str, Any],
|
||||
) -> list[Callable[[], LiberoEnv]]:
|
||||
"""Build n_envs factory callables for a single (suite, task_id)."""
|
||||
joined_cams = ",".join(camera_names) # keep backward-compat: downstream expects a string
|
||||
|
||||
def _make_env(episode_index: int, **kwargs) -> LiberoEnv:
|
||||
local_kwargs = dict(kwargs)
|
||||
@@ -298,7 +320,7 @@ def _make_env_fns(
|
||||
task_suite=suite,
|
||||
task_id=task_id,
|
||||
task_suite_name=suite_name,
|
||||
camera_name=camera_names,
|
||||
camera_name=joined_cams,
|
||||
init_states=init_states,
|
||||
episode_index=episode_index,
|
||||
**local_kwargs,
|
||||
|
||||
@@ -99,7 +99,6 @@ def env_to_policy_features(env_cfg: EnvConfig) -> dict[str, PolicyFeature]:
|
||||
|
||||
policy_key = env_cfg.features_map[key]
|
||||
policy_features[policy_key] = feature
|
||||
|
||||
return policy_features
|
||||
|
||||
|
||||
|
||||
@@ -65,10 +65,10 @@ def make_act_pre_post_processors(
|
||||
),
|
||||
]
|
||||
output_steps = [
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
]
|
||||
|
||||
return (
|
||||
|
||||
@@ -73,10 +73,10 @@ def make_diffusion_pre_post_processors(
|
||||
),
|
||||
]
|
||||
output_steps = [
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
]
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
|
||||
@@ -359,7 +359,6 @@ def make_policy(
|
||||
if env_cfg is None:
|
||||
raise ValueError("env_cfg cannot be None when ds_meta is not provided")
|
||||
features = env_to_policy_features(env_cfg)
|
||||
|
||||
cfg.output_features = {key: ft for key, ft in features.items() if ft.type is FeatureType.ACTION}
|
||||
cfg.input_features = {key: ft for key, ft in features.items() if key not in cfg.output_features}
|
||||
kwargs["config"] = cfg
|
||||
@@ -372,10 +371,7 @@ def make_policy(
|
||||
else:
|
||||
# Make a fresh policy.
|
||||
policy = policy_cls(**kwargs)
|
||||
|
||||
policy.to(cfg.device)
|
||||
assert isinstance(policy, torch.nn.Module)
|
||||
|
||||
# policy = torch.compile(policy, mode="reduce-overhead")
|
||||
|
||||
return policy
|
||||
|
||||
@@ -146,10 +146,10 @@ def make_pi0_pre_post_processors(
|
||||
]
|
||||
|
||||
output_steps: list[ProcessorStep] = [
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
]
|
||||
|
||||
return (
|
||||
|
||||
@@ -73,10 +73,10 @@ def make_pi0fast_pre_post_processors(
|
||||
),
|
||||
]
|
||||
output_steps = [
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
]
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
|
||||
@@ -246,9 +246,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
|
||||
base_model=base_model,
|
||||
)
|
||||
|
||||
template_card = (
|
||||
files("lerobot.templates").joinpath("lerobot_modelcard_template.md").read_text(encoding="utf-8")
|
||||
)
|
||||
template_card = files("lerobot.templates").joinpath("lerobot_modelcard_template.md").read_text()
|
||||
card = ModelCard.from_template(card_data, template_str=template_card)
|
||||
card.validate()
|
||||
return card
|
||||
|
||||
@@ -73,10 +73,10 @@ def make_sac_pre_post_processors(
|
||||
),
|
||||
]
|
||||
output_steps = [
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
]
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
|
||||
@@ -84,10 +84,10 @@ def make_smolvla_pre_post_processors(
|
||||
),
|
||||
]
|
||||
output_steps = [
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
]
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
|
||||
@@ -71,10 +71,10 @@ def make_tdmpc_pre_post_processors(
|
||||
),
|
||||
]
|
||||
output_steps = [
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
]
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
|
||||
@@ -72,10 +72,10 @@ def make_vqbet_pre_post_processors(
|
||||
),
|
||||
]
|
||||
output_steps = [
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
]
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
|
||||
@@ -36,10 +36,7 @@ from .factory import (
|
||||
make_default_robot_observation_processor,
|
||||
make_default_teleop_action_processor,
|
||||
)
|
||||
from .gym_action_processor import (
|
||||
Numpy2TorchActionProcessorStep,
|
||||
Torch2NumpyActionProcessorStep,
|
||||
)
|
||||
from .gym_action_processor import Numpy2TorchActionProcessorStep, Torch2NumpyActionProcessorStep
|
||||
from .hil_processor import (
|
||||
AddTeleopActionAsComplimentaryDataStep,
|
||||
AddTeleopEventsAsInfoStep,
|
||||
@@ -70,10 +67,6 @@ from .pipeline import (
|
||||
RobotProcessorPipeline,
|
||||
TruncatedProcessorStep,
|
||||
)
|
||||
from .policy_robot_bridge import (
|
||||
PolicyActionToRobotActionProcessorStep,
|
||||
RobotActionToPolicyActionProcessorStep,
|
||||
)
|
||||
from .rename_processor import RenameObservationsProcessorStep
|
||||
from .tokenizer_processor import TokenizerProcessorStep
|
||||
|
||||
@@ -123,8 +116,6 @@ __all__ = [
|
||||
"RobotProcessorPipeline",
|
||||
"TokenizerProcessorStep",
|
||||
"Torch2NumpyActionProcessorStep",
|
||||
"RobotActionToPolicyActionProcessorStep",
|
||||
"PolicyActionToRobotActionProcessorStep",
|
||||
"transition_to_batch",
|
||||
"TransitionKey",
|
||||
"TruncatedProcessorStep",
|
||||
|
||||
@@ -207,37 +207,14 @@ def create_transition(
|
||||
}
|
||||
|
||||
|
||||
def robot_action_observation_to_transition(
|
||||
action_observation: tuple[RobotAction, RobotObservation],
|
||||
) -> EnvTransition:
|
||||
"""
|
||||
Convert a raw robot action and observation dictionary into a standardized `EnvTransition`.
|
||||
|
||||
Args:
|
||||
action: The raw action dictionary from a teleoperation device or controller.
|
||||
observation: The raw observation dictionary from the environment.
|
||||
|
||||
Returns:
|
||||
An `EnvTransition` containing the formatted observation.
|
||||
"""
|
||||
if not isinstance(action_observation, tuple):
|
||||
raise ValueError("action_observation should be a tuple type with an action and observation")
|
||||
|
||||
action, observation = action_observation
|
||||
|
||||
if action is not None and not isinstance(action, dict):
|
||||
raise ValueError(f"Action should be a RobotAction type got {type(action)}")
|
||||
|
||||
if observation is not None and not isinstance(observation, dict):
|
||||
raise ValueError(f"Observation should be a RobotObservation type got {type(observation)}")
|
||||
|
||||
return create_transition(action=action, observation=observation)
|
||||
|
||||
|
||||
def robot_action_to_transition(action: RobotAction) -> EnvTransition:
|
||||
"""
|
||||
Convert a raw robot action dictionary into a standardized `EnvTransition`.
|
||||
|
||||
The keys in the action dictionary are prefixed with "action." and stored under
|
||||
the `ACTION` key in the transition. Values are converted to tensors, except for
|
||||
special types like `Rotation`.
|
||||
|
||||
Args:
|
||||
action: The raw action dictionary from a teleoperation device or controller.
|
||||
|
||||
@@ -253,6 +230,10 @@ def observation_to_transition(observation: RobotObservation) -> EnvTransition:
|
||||
"""
|
||||
Convert a raw robot observation dictionary into a standardized `EnvTransition`.
|
||||
|
||||
The observation is split into state and image components. State keys are prefixed
|
||||
with "observation.state." and image keys with "observation.images.". The result is
|
||||
stored under the `OBSERVATION` key in the transition.
|
||||
|
||||
Args:
|
||||
observation: The raw observation dictionary from the environment.
|
||||
|
||||
|
||||
@@ -48,12 +48,12 @@ class MapTensorToDeltaActionDictStep(ActionProcessorStep):
|
||||
|
||||
# TODO (maractingi): add rotation
|
||||
delta_action = {
|
||||
"delta_x": action[0].item(),
|
||||
"delta_y": action[1].item(),
|
||||
"delta_z": action[2].item(),
|
||||
"delta_x": action[0],
|
||||
"delta_y": action[1],
|
||||
"delta_z": action[2],
|
||||
}
|
||||
if self.use_gripper:
|
||||
delta_action["gripper"] = action[3].item()
|
||||
delta_action["gripper"] = action[3]
|
||||
return delta_action
|
||||
|
||||
def transform_features(
|
||||
@@ -126,7 +126,7 @@ class MapDeltaActionToRobotActionStep(RobotActionProcessorStep):
|
||||
"target_wx": target_wx,
|
||||
"target_wy": target_wy,
|
||||
"target_wz": target_wz,
|
||||
"gripper_vel": float(gripper),
|
||||
"gripper": float(gripper),
|
||||
}
|
||||
|
||||
return action
|
||||
|
||||
@@ -105,10 +105,6 @@ class DeviceProcessorStep(ProcessorStep):
|
||||
# In both cases, use the configured device.
|
||||
target_device = self.tensor_device
|
||||
|
||||
# MPS workaround: Convert float64 to float32 since MPS doesn't support float64
|
||||
if target_device.type == "mps" and tensor.dtype == torch.float64:
|
||||
tensor = tensor.to(dtype=torch.float32)
|
||||
|
||||
# Only move if necessary
|
||||
if tensor.device != target_device:
|
||||
tensor = tensor.to(target_device, non_blocking=self.non_blocking)
|
||||
|
||||
@@ -16,7 +16,7 @@
|
||||
|
||||
from .converters import (
|
||||
observation_to_transition,
|
||||
robot_action_observation_to_transition,
|
||||
robot_action_to_transition,
|
||||
transition_to_observation,
|
||||
transition_to_robot_action,
|
||||
)
|
||||
@@ -24,23 +24,19 @@ from .core import RobotAction, RobotObservation
|
||||
from .pipeline import IdentityProcessorStep, RobotProcessorPipeline
|
||||
|
||||
|
||||
def make_default_teleop_action_processor() -> RobotProcessorPipeline[
|
||||
tuple[RobotAction, RobotObservation], RobotAction
|
||||
]:
|
||||
teleop_action_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
|
||||
def make_default_teleop_action_processor() -> RobotProcessorPipeline[RobotAction, RobotAction]:
|
||||
teleop_action_processor = RobotProcessorPipeline[RobotAction, RobotAction](
|
||||
steps=[IdentityProcessorStep()],
|
||||
to_transition=robot_action_observation_to_transition,
|
||||
to_transition=robot_action_to_transition,
|
||||
to_output=transition_to_robot_action,
|
||||
)
|
||||
return teleop_action_processor
|
||||
|
||||
|
||||
def make_default_robot_action_processor() -> RobotProcessorPipeline[
|
||||
tuple[RobotAction, RobotObservation], RobotAction
|
||||
]:
|
||||
robot_action_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
|
||||
def make_default_robot_action_processor() -> RobotProcessorPipeline[RobotAction, RobotAction]:
|
||||
robot_action_processor = RobotProcessorPipeline[RobotAction, RobotAction](
|
||||
steps=[IdentityProcessorStep()],
|
||||
to_transition=robot_action_observation_to_transition,
|
||||
to_transition=robot_action_to_transition,
|
||||
to_output=transition_to_robot_action,
|
||||
)
|
||||
return robot_action_processor
|
||||
|
||||
@@ -19,8 +19,8 @@ from dataclasses import dataclass
|
||||
from lerobot.configs.types import PipelineFeatureType, PolicyFeature
|
||||
|
||||
from .converters import to_tensor
|
||||
from .core import EnvAction, EnvTransition, PolicyAction
|
||||
from .pipeline import ActionProcessorStep, ProcessorStep, ProcessorStepRegistry
|
||||
from .core import EnvAction, PolicyAction
|
||||
from .pipeline import ActionProcessorStep, ProcessorStepRegistry
|
||||
|
||||
|
||||
@ProcessorStepRegistry.register("torch2numpy_action_processor")
|
||||
@@ -69,27 +69,23 @@ class Torch2NumpyActionProcessorStep(ActionProcessorStep):
|
||||
|
||||
@ProcessorStepRegistry.register("numpy2torch_action_processor")
|
||||
@dataclass
|
||||
class Numpy2TorchActionProcessorStep(ProcessorStep):
|
||||
"""Converts a NumPy array action to a PyTorch tensor when action is present."""
|
||||
class Numpy2TorchActionProcessorStep(ActionProcessorStep):
|
||||
"""
|
||||
Converts a NumPy array action to a PyTorch tensor.
|
||||
|
||||
def __call__(self, transition: EnvTransition) -> EnvTransition:
|
||||
"""Converts numpy action to torch tensor if action exists, otherwise passes through."""
|
||||
from .core import TransitionKey
|
||||
This step is useful for converting actions from environments or hardware,
|
||||
which are often NumPy arrays, into PyTorch tensors that can be processed
|
||||
by a policy or model.
|
||||
"""
|
||||
|
||||
self._current_transition = transition.copy()
|
||||
new_transition = self._current_transition
|
||||
|
||||
action = new_transition.get(TransitionKey.ACTION)
|
||||
if action is not None:
|
||||
if not isinstance(action, EnvAction):
|
||||
raise TypeError(
|
||||
f"Expected np.ndarray or None, got {type(action).__name__}. "
|
||||
"Use appropriate processor for non-tensor actions."
|
||||
)
|
||||
torch_action = to_tensor(action, dtype=None) # Preserve original dtype
|
||||
new_transition[TransitionKey.ACTION] = torch_action
|
||||
|
||||
return new_transition
|
||||
def action(self, action: EnvAction) -> PolicyAction:
|
||||
if not isinstance(action, EnvAction):
|
||||
raise TypeError(
|
||||
f"Expected np.ndarray or None, got {type(action).__name__}. "
|
||||
"Use appropriate processor for non-tensor actions."
|
||||
)
|
||||
torch_action = to_tensor(action, dtype=None) # Preserve original dtype
|
||||
return torch_action
|
||||
|
||||
def transform_features(
|
||||
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||
|
||||
@@ -340,16 +340,11 @@ class GripperPenaltyProcessorStep(ComplementaryDataProcessorStep):
|
||||
"""
|
||||
action = self.transition.get(TransitionKey.ACTION)
|
||||
|
||||
raw_joint_positions = complementary_data.get("raw_joint_positions", None)
|
||||
if raw_joint_positions is None:
|
||||
return complementary_data
|
||||
|
||||
current_gripper_pos = raw_joint_positions.get(GRIPPER_KEY, None)
|
||||
current_gripper_pos = complementary_data.get("raw_joint_positions", None).get(GRIPPER_KEY, None)
|
||||
if current_gripper_pos is None:
|
||||
return complementary_data
|
||||
|
||||
# Gripper action is a PolicyAction at this stage
|
||||
gripper_action = action[-1].item()
|
||||
gripper_action = action[f"{GRIPPER_KEY}.pos"]
|
||||
gripper_action_normalized = gripper_action / self.max_gripper_pos
|
||||
|
||||
# Normalize gripper state and action
|
||||
|
||||
+145
-693
@@ -237,19 +237,6 @@ class ProcessorKwargs(TypedDict, total=False):
|
||||
after_step_hooks: list[Callable[[int, EnvTransition], None]] | None
|
||||
|
||||
|
||||
class ProcessorMigrationError(Exception):
|
||||
"""Raised when a model needs migration to the processor format"""
|
||||
|
||||
def __init__(self, model_path: str | Path, migration_command: str, original_error: str):
|
||||
self.model_path = model_path
|
||||
self.migration_command = migration_command
|
||||
self.original_error = original_error
|
||||
super().__init__(
|
||||
f"Model '{model_path}' requires migration to processor format. "
|
||||
f"Run: {migration_command}\n\nOriginal error: {original_error}"
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class DataProcessorPipeline(HubMixin, Generic[TInput, TOutput]):
|
||||
"""A sequential pipeline for processing data, integrated with the Hugging Face Hub.
|
||||
@@ -452,7 +439,6 @@ class DataProcessorPipeline(HubMixin, Generic[TInput, TOutput]):
|
||||
def from_pretrained(
|
||||
cls,
|
||||
pretrained_model_name_or_path: str | Path,
|
||||
config_filename: str,
|
||||
*,
|
||||
force_download: bool = False,
|
||||
resume_download: bool | None = None,
|
||||
@@ -461,74 +447,24 @@ class DataProcessorPipeline(HubMixin, Generic[TInput, TOutput]):
|
||||
cache_dir: str | Path | None = None,
|
||||
local_files_only: bool = False,
|
||||
revision: str | None = None,
|
||||
config_filename: str | None = None,
|
||||
overrides: dict[str, Any] | None = None,
|
||||
to_transition: Callable[[TInput], EnvTransition] | None = None,
|
||||
to_output: Callable[[EnvTransition], TOutput] | None = None,
|
||||
**kwargs,
|
||||
) -> DataProcessorPipeline[TInput, TOutput]:
|
||||
"""Loads a pipeline from a local directory, single file, or Hugging Face Hub repository.
|
||||
"""Loads a pipeline from a local directory or a Hugging Face Hub repository.
|
||||
|
||||
This method implements a simplified loading pipeline with intelligent migration detection:
|
||||
|
||||
**Simplified Loading Strategy**:
|
||||
1. **Config Loading** (_load_config):
|
||||
- **Directory**: Load specified config_filename from directory
|
||||
- **Single file**: Load file directly (config_filename ignored)
|
||||
- **Hub repository**: Download specified config_filename from Hub
|
||||
|
||||
2. **Config Validation** (_validate_loaded_config):
|
||||
- Format validation: Ensure config is valid processor format
|
||||
- Migration detection: Guide users to migrate old LeRobot models
|
||||
- Clear errors: Provide actionable error messages
|
||||
|
||||
3. **Step Construction** (_build_steps_with_overrides):
|
||||
- Class resolution: Registry lookup or dynamic imports
|
||||
- Override merging: User parameters override saved config
|
||||
- State loading: Load .safetensors files for stateful steps
|
||||
|
||||
4. **Override Validation** (_validate_overrides_used):
|
||||
- Ensure all user overrides were applied (catch typos)
|
||||
- Provide helpful error messages with available keys
|
||||
|
||||
**Migration Detection**:
|
||||
- **Smart detection**: Analyzes JSON files to detect old LeRobot models
|
||||
- **Precise targeting**: Avoids false positives on other HuggingFace models
|
||||
- **Clear guidance**: Provides exact migration command to run
|
||||
- **Error mode**: Always raises ProcessorMigrationError for clear user action
|
||||
|
||||
**Loading Examples**:
|
||||
```python
|
||||
# Directory loading
|
||||
pipeline = DataProcessorPipeline.from_pretrained("/models/my_model", config_filename="processor.json")
|
||||
|
||||
# Single file loading
|
||||
pipeline = DataProcessorPipeline.from_pretrained(
|
||||
"/models/my_model/processor.json", config_filename="processor.json"
|
||||
)
|
||||
|
||||
# Hub loading
|
||||
pipeline = DataProcessorPipeline.from_pretrained("user/repo", config_filename="processor.json")
|
||||
|
||||
# Multiple configs (preprocessor/postprocessor)
|
||||
preprocessor = DataProcessorPipeline.from_pretrained(
|
||||
"model", config_filename="policy_preprocessor.json"
|
||||
)
|
||||
postprocessor = DataProcessorPipeline.from_pretrained(
|
||||
"model", config_filename="policy_postprocessor.json"
|
||||
)
|
||||
```
|
||||
|
||||
**Override System**:
|
||||
- **Key matching**: Use registry names or class names as override keys
|
||||
- **Config merging**: User overrides take precedence over saved config
|
||||
- **Validation**: Ensure all override keys match actual steps (catch typos)
|
||||
- **Example**: overrides={"NormalizeStep": {"device": "cuda"}}
|
||||
This method reconstructs a `DataProcessorPipeline` by:
|
||||
1. Loading the main JSON configuration file.
|
||||
2. Iterating through the steps defined in the config.
|
||||
3. Dynamically importing or looking up each step's class.
|
||||
4. Instantiating each step with its saved configuration, potentially with overrides.
|
||||
5. Loading the step's state from its `.safetensors` file, if it exists.
|
||||
|
||||
Args:
|
||||
pretrained_model_name_or_path: The identifier of the repository on the Hugging Face Hub,
|
||||
a path to a local directory, or a path to a single config file.
|
||||
config_filename: The name of the pipeline's JSON configuration file. Always required
|
||||
to prevent ambiguity when multiple configs exist (e.g., preprocessor vs postprocessor).
|
||||
pretrained_model_name_or_path: The identifier of the repository on the Hugging Face Hub
|
||||
or a path to a local directory.
|
||||
force_download: Whether to force (re)downloading the files.
|
||||
resume_download: Whether to resume a previously interrupted download.
|
||||
proxies: A dictionary of proxy servers to use.
|
||||
@@ -536,6 +472,9 @@ class DataProcessorPipeline(HubMixin, Generic[TInput, TOutput]):
|
||||
cache_dir: The path to a specific cache folder to store downloaded files.
|
||||
local_files_only: If True, avoid downloading files from the Hub.
|
||||
revision: The specific model version to use (e.g., a branch name, tag name, or commit id).
|
||||
config_filename: The name of the pipeline's JSON configuration file. If not provided,
|
||||
it's auto-detected in local directories (if only one .json file exists). This parameter
|
||||
is mandatory when loading from Hugging Face Hub repositories.
|
||||
overrides: A dictionary to override the configuration of specific steps. Keys should
|
||||
match the step's class name or registry name.
|
||||
to_transition: A custom function to convert input data to `EnvTransition`.
|
||||
@@ -550,667 +489,180 @@ class DataProcessorPipeline(HubMixin, Generic[TInput, TOutput]):
|
||||
ValueError: If configuration is ambiguous or instantiation fails.
|
||||
ImportError: If a step's class cannot be imported.
|
||||
KeyError: If an override key doesn't match any step in the pipeline.
|
||||
ProcessorMigrationError: If the model requires migration to processor format.
|
||||
"""
|
||||
model_id = str(pretrained_model_name_or_path)
|
||||
hub_download_kwargs = {
|
||||
"force_download": force_download,
|
||||
"resume_download": resume_download,
|
||||
"proxies": proxies,
|
||||
"token": token,
|
||||
"cache_dir": cache_dir,
|
||||
"local_files_only": local_files_only,
|
||||
"revision": revision,
|
||||
}
|
||||
loaded_config: dict[str, Any] | None = None
|
||||
base_path: Path | None = None
|
||||
|
||||
# 1. Load configuration using simplified 3-way logic
|
||||
loaded_config, base_path = cls._load_config(model_id, config_filename, hub_download_kwargs)
|
||||
# Standard pattern: try local directory first
|
||||
if Path(model_id).is_dir():
|
||||
base_path = Path(model_id)
|
||||
|
||||
# 2. Validate configuration and handle migration
|
||||
cls._validate_loaded_config(model_id, loaded_config, config_filename)
|
||||
# Handle config filename
|
||||
if config_filename is None:
|
||||
json_files = list(base_path.glob("*.json"))
|
||||
if len(json_files) == 0:
|
||||
# No config files found locally, will try Hub next
|
||||
pass
|
||||
elif len(json_files) == 1:
|
||||
config_filename = json_files[0].name
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Multiple .json files found in {model_id}: {[f.name for f in json_files]}. "
|
||||
f"Please specify which one to load using the config_filename parameter."
|
||||
)
|
||||
|
||||
# 3. Build steps with overrides
|
||||
steps, validated_overrides = cls._build_steps_with_overrides(
|
||||
loaded_config, overrides or {}, model_id, base_path, hub_download_kwargs
|
||||
)
|
||||
# Try to load config from local directory
|
||||
if config_filename and (base_path / config_filename).exists():
|
||||
with open(base_path / config_filename) as f:
|
||||
loaded_config = json.load(f)
|
||||
|
||||
# 4. Validate that all overrides were used
|
||||
cls._validate_overrides_used(validated_overrides, loaded_config)
|
||||
# If not found locally, try Hub
|
||||
if loaded_config is None:
|
||||
# Check if this looks like a local path that doesn't exist
|
||||
# Hub repo IDs have format "user/repo" with exactly one slash
|
||||
# Local paths typically have multiple slashes, backslashes, or start with ./ or ../
|
||||
looks_like_local_path = (
|
||||
model_id.count("/") > 1 # Multiple slashes suggest local path
|
||||
or "\\" in model_id # Backslashes are only in local paths
|
||||
or Path(model_id).is_absolute() # Absolute paths are local
|
||||
or model_id.startswith("./")
|
||||
or model_id.startswith("../") # Relative path indicators
|
||||
)
|
||||
|
||||
# 5. Construct and return the final pipeline instance
|
||||
return cls(
|
||||
steps=steps,
|
||||
name=loaded_config.get("name", "DataProcessorPipeline"),
|
||||
to_transition=to_transition or cast(Callable[[TInput], EnvTransition], batch_to_transition),
|
||||
to_output=to_output or cast(Callable[[EnvTransition], TOutput], transition_to_batch),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _load_config(
|
||||
cls,
|
||||
model_id: str,
|
||||
config_filename: str,
|
||||
hub_download_kwargs: dict[str, Any],
|
||||
) -> tuple[dict[str, Any], Path]:
|
||||
"""Load configuration from local file or Hugging Face Hub.
|
||||
|
||||
This method implements a super-simplified 3-way loading strategy:
|
||||
|
||||
1. **Local directory**: Load config_filename from directory
|
||||
- Example: model_id="/models/my_model", config_filename="processor.json"
|
||||
- Loads: "/models/my_model/processor.json"
|
||||
|
||||
2. **Single file**: Load file directly (ignore config_filename)
|
||||
- Example: model_id="/models/my_model/processor.json"
|
||||
- Loads: "/models/my_model/processor.json" (config_filename ignored)
|
||||
|
||||
3. **Hub repository**: Download config_filename from Hub
|
||||
- Example: model_id="user/repo", config_filename="processor.json"
|
||||
- Downloads and loads: config_filename from Hub repo
|
||||
|
||||
**Benefits of Explicit config_filename**:
|
||||
- No auto-detection complexity or edge cases
|
||||
- No risk of loading wrong config (preprocessor vs postprocessor)
|
||||
- Consistent behavior across local and Hub usage
|
||||
- Clear, predictable errors
|
||||
|
||||
Args:
|
||||
model_id: The model identifier (Hub repo ID, local directory, or file path)
|
||||
config_filename: The explicit config filename to load (always required)
|
||||
hub_download_kwargs: Parameters for hf_hub_download (tokens, cache, etc.)
|
||||
|
||||
Returns:
|
||||
Tuple of (loaded_config, base_path)
|
||||
- loaded_config: Parsed JSON config dict (always loaded, never None)
|
||||
- base_path: Directory containing config file (for state file resolution)
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If config file cannot be found locally or on Hub
|
||||
"""
|
||||
model_path = Path(model_id)
|
||||
|
||||
if model_path.is_dir():
|
||||
# Directory: load specified config from directory
|
||||
config_path = model_path / config_filename
|
||||
if not config_path.exists():
|
||||
# Check for migration before giving clear error
|
||||
if cls._should_suggest_migration(model_path):
|
||||
cls._suggest_processor_migration(model_id, f"Config file '{config_filename}' not found")
|
||||
raise FileNotFoundError(
|
||||
f"Config file '{config_filename}' not found in directory '{model_id}'"
|
||||
if looks_like_local_path:
|
||||
# This appears to be a local path that doesn't exist
|
||||
raise FileNotFoundError(f"Local path '{model_id}' does not exist")
|
||||
# For Hub repositories, config_filename is mandatory
|
||||
if config_filename is None:
|
||||
raise ValueError(
|
||||
f"When loading from Hugging Face Hub, 'config_filename' must be specified. "
|
||||
f"Example: DataProcessorPipeline.from_pretrained('{model_id}', config_filename='processor.json')"
|
||||
)
|
||||
|
||||
with open(config_path) as f:
|
||||
return json.load(f), model_path
|
||||
|
||||
elif model_path.is_file():
|
||||
# File: load file directly (config_filename is ignored for single files)
|
||||
with open(model_path) as f:
|
||||
return json.load(f), model_path.parent
|
||||
|
||||
else:
|
||||
# Hub: download specified config
|
||||
try:
|
||||
# Download the configuration file from the Hub
|
||||
config_path = hf_hub_download(
|
||||
repo_id=model_id,
|
||||
filename=config_filename,
|
||||
repo_type="model",
|
||||
**hub_download_kwargs,
|
||||
force_download=force_download,
|
||||
resume_download=resume_download,
|
||||
proxies=proxies,
|
||||
token=token,
|
||||
cache_dir=cache_dir,
|
||||
local_files_only=local_files_only,
|
||||
revision=revision,
|
||||
)
|
||||
|
||||
with open(config_path) as f:
|
||||
return json.load(f), Path(config_path).parent
|
||||
loaded_config = json.load(f)
|
||||
|
||||
# The base path for other files (like state tensors) is the directory of the config file
|
||||
base_path = Path(config_path).parent
|
||||
|
||||
except Exception as e:
|
||||
raise FileNotFoundError(
|
||||
f"Could not find '{config_filename}' on the HuggingFace Hub at '{model_id}'"
|
||||
f"Could not find {config_filename} on the HuggingFace Hub at {model_id}"
|
||||
) from e
|
||||
|
||||
@classmethod
|
||||
def _validate_loaded_config(
|
||||
cls, model_id: str, loaded_config: dict[str, Any], config_filename: str
|
||||
) -> None:
|
||||
"""Validate that a config was loaded and is a valid processor config.
|
||||
# At this point, loaded_config must be loaded successfully
|
||||
if loaded_config is None:
|
||||
raise RuntimeError("Failed to load configuration from local directory or Hub")
|
||||
|
||||
This method validates processor config format with intelligent migration detection:
|
||||
if overrides is None:
|
||||
overrides = {}
|
||||
|
||||
**Config Format Validation**:
|
||||
- Use _is_processor_config() to validate structure
|
||||
- Must have "steps" field with list of step configurations
|
||||
- Each step needs "class" or "registry_name"
|
||||
- If validation fails AND local directory: Check for migration need
|
||||
- If migration needed: Raise ProcessorMigrationError with command
|
||||
- If no migration: Raise ValueError with helpful error message
|
||||
|
||||
**Migration Detection Logic**:
|
||||
- Only triggered for local directories (not Hub repos)
|
||||
- Analyzes all JSON files in directory to detect old LeRobot models
|
||||
- Provides exact migration command with model path
|
||||
|
||||
Args:
|
||||
model_id: The model identifier (used for migration detection)
|
||||
loaded_config: The loaded config dictionary (guaranteed non-None)
|
||||
config_filename: The config filename that was loaded (for error messages)
|
||||
|
||||
Raises:
|
||||
ValueError: If config format is invalid
|
||||
ProcessorMigrationError: If model needs migration to processor format
|
||||
"""
|
||||
# Validate that this is actually a processor config
|
||||
if not cls._is_processor_config(loaded_config):
|
||||
if Path(model_id).is_dir() and cls._should_suggest_migration(Path(model_id)):
|
||||
cls._suggest_processor_migration(
|
||||
model_id,
|
||||
f"Config file '{config_filename}' is not a valid processor configuration",
|
||||
)
|
||||
raise ValueError(
|
||||
f"Config file '{config_filename}' is not a valid processor configuration. "
|
||||
f"Expected a config with 'steps' field, but got: {list(loaded_config.keys())}"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _build_steps_with_overrides(
|
||||
cls,
|
||||
loaded_config: dict[str, Any],
|
||||
overrides: dict[str, Any],
|
||||
model_id: str,
|
||||
base_path: Path | None,
|
||||
hub_download_kwargs: dict[str, Any],
|
||||
) -> tuple[list[ProcessorStep], set[str]]:
|
||||
"""Build all processor steps with overrides and state loading.
|
||||
|
||||
This method orchestrates the complete step construction pipeline:
|
||||
|
||||
**For each step in loaded_config["steps"]**:
|
||||
|
||||
1. **Class Resolution** (via _resolve_step_class):
|
||||
- **If "registry_name" exists**: Look up in ProcessorStepRegistry
|
||||
Example: {"registry_name": "normalize_step"} -> Get registered class
|
||||
- **Else use "class" field**: Dynamic import from full module path
|
||||
Example: {"class": "lerobot.processor.normalize.NormalizeStep"}
|
||||
- **Result**: (step_class, step_key) where step_key is used for overrides
|
||||
|
||||
2. **Step Instantiation** (via _instantiate_step):
|
||||
- **Merge configs**: saved_config + user_overrides
|
||||
- **Override priority**: User overrides take precedence over saved config
|
||||
- **Example**: saved={"mean": 0.0}, override={"mean": 1.0} -> final={"mean": 1.0}
|
||||
- **Result**: Instantiated ProcessorStep object
|
||||
|
||||
3. **State Loading** (via _load_step_state):
|
||||
- **If step has "state_file"**: Load tensor state from .safetensors
|
||||
- **Local first**: Check base_path/state_file.safetensors
|
||||
- **Hub fallback**: Download state file if not found locally
|
||||
- **Optional**: Only load if step has load_state_dict method
|
||||
|
||||
4. **Override Tracking**:
|
||||
- **Track used overrides**: Remove step_key from remaining set
|
||||
- **Purpose**: Validate all user overrides were applied (detect typos)
|
||||
|
||||
**Error Handling**:
|
||||
- Class resolution errors -> ImportError with helpful message
|
||||
- Instantiation errors -> ValueError with config details
|
||||
- State loading errors -> Propagated from load_state_dict
|
||||
|
||||
Args:
|
||||
loaded_config: The loaded processor configuration (must have "steps" field)
|
||||
overrides: User-provided parameter overrides (keyed by class/registry name)
|
||||
model_id: The model identifier (needed for Hub state file downloads)
|
||||
base_path: Local directory path for finding state files
|
||||
hub_download_kwargs: Parameters for hf_hub_download (tokens, cache, etc.)
|
||||
|
||||
Returns:
|
||||
Tuple of (instantiated_steps_list, unused_override_keys)
|
||||
- instantiated_steps_list: List of ready-to-use ProcessorStep instances
|
||||
- unused_override_keys: Override keys that didn't match any step (for validation)
|
||||
|
||||
Raises:
|
||||
ImportError: If a step class cannot be imported or found in registry
|
||||
ValueError: If a step cannot be instantiated with its configuration
|
||||
"""
|
||||
steps: list[ProcessorStep] = []
|
||||
override_keys = set(overrides.keys())
|
||||
|
||||
steps: list[ProcessorStep] = []
|
||||
for step_entry in loaded_config["steps"]:
|
||||
# 1. Get step class and key
|
||||
step_class, step_key = cls._resolve_step_class(step_entry)
|
||||
# Determine the step class, prioritizing the registry.
|
||||
if "registry_name" in step_entry:
|
||||
try:
|
||||
step_class = ProcessorStepRegistry.get(step_entry["registry_name"])
|
||||
step_key = step_entry["registry_name"]
|
||||
except KeyError as e:
|
||||
raise ImportError(f"Failed to load processor step from registry. {str(e)}") from e
|
||||
else:
|
||||
# Fallback to dynamic import using the full class path.
|
||||
full_class_path = step_entry["class"]
|
||||
module_path, class_name = full_class_path.rsplit(".", 1)
|
||||
|
||||
# 2. Instantiate step with overrides
|
||||
step_instance = cls._instantiate_step(step_entry, step_class, step_key, overrides)
|
||||
try:
|
||||
module = importlib.import_module(module_path)
|
||||
step_class = getattr(module, class_name)
|
||||
step_key = class_name
|
||||
except (ImportError, AttributeError) as e:
|
||||
raise ImportError(
|
||||
f"Failed to load processor step '{full_class_path}'. "
|
||||
f"Make sure the module '{module_path}' is installed and contains class '{class_name}'. "
|
||||
f"Consider registering the step using @ProcessorStepRegistry.register() for better portability. "
|
||||
f"Error: {str(e)}"
|
||||
) from e
|
||||
|
||||
# 3. Load step state if available
|
||||
cls._load_step_state(step_instance, step_entry, model_id, base_path, hub_download_kwargs)
|
||||
|
||||
# 4. Track used overrides
|
||||
if step_key in override_keys:
|
||||
override_keys.discard(step_key)
|
||||
|
||||
steps.append(step_instance)
|
||||
|
||||
return steps, override_keys
|
||||
|
||||
@classmethod
|
||||
def _resolve_step_class(cls, step_entry: dict[str, Any]) -> tuple[type[ProcessorStep], str]:
|
||||
"""Resolve step class from registry or import path.
|
||||
|
||||
This method implements a two-tier resolution strategy:
|
||||
|
||||
**Tier 1: Registry-based resolution** (preferred):
|
||||
- **If "registry_name" in step_entry**: Look up in ProcessorStepRegistry
|
||||
- **Advantage**: Faster, no imports needed, guaranteed compatibility
|
||||
- **Example**: {"registry_name": "normalize_step"} -> Get pre-registered class
|
||||
- **Error**: KeyError if registry_name not found -> Convert to ImportError
|
||||
|
||||
**Tier 2: Dynamic import fallback**:
|
||||
- **Else use "class" field**: Full module.ClassName import path
|
||||
- **Process**: Split "module.path.ClassName" into module + class parts
|
||||
- **Import**: Use importlib.import_module() + getattr()
|
||||
- **Example**: "lerobot.processor.normalize.NormalizeStep"
|
||||
a. Import module: "lerobot.processor.normalize"
|
||||
b. Get class: getattr(module, "NormalizeStep")
|
||||
- **step_key**: Use class_name ("NormalizeStep") for overrides
|
||||
|
||||
**Override Key Strategy**:
|
||||
- Registry steps: Use registry_name ("normalize_step")
|
||||
- Import steps: Use class_name ("NormalizeStep")
|
||||
- This allows users to override with: {"normalize_step": {...}} or {"NormalizeStep": {...}}
|
||||
|
||||
**Error Handling**:
|
||||
- Registry KeyError -> ImportError with registry context
|
||||
- Import/Attribute errors -> ImportError with helpful suggestions
|
||||
- All errors include troubleshooting guidance
|
||||
|
||||
Args:
|
||||
step_entry: The step configuration dictionary (must have "registry_name" or "class")
|
||||
|
||||
Returns:
|
||||
Tuple of (step_class, step_key)
|
||||
- step_class: The resolved ProcessorStep class (ready for instantiation)
|
||||
- step_key: The key used for user overrides (registry_name or class_name)
|
||||
|
||||
Raises:
|
||||
ImportError: If step class cannot be loaded from registry or import path
|
||||
"""
|
||||
if "registry_name" in step_entry:
|
||||
# Instantiate the step, merging saved config with user-provided overrides.
|
||||
try:
|
||||
step_class = ProcessorStepRegistry.get(step_entry["registry_name"])
|
||||
return step_class, step_entry["registry_name"]
|
||||
except KeyError as e:
|
||||
raise ImportError(f"Failed to load processor step from registry. {str(e)}") from e
|
||||
else:
|
||||
# Fallback to dynamic import using the full class path
|
||||
full_class_path = step_entry["class"]
|
||||
module_path, class_name = full_class_path.rsplit(".", 1)
|
||||
saved_cfg = step_entry.get("config", {})
|
||||
step_overrides = overrides.get(step_key, {})
|
||||
merged_cfg = {**saved_cfg, **step_overrides}
|
||||
step_instance: ProcessorStep = step_class(**merged_cfg)
|
||||
|
||||
try:
|
||||
module = importlib.import_module(module_path)
|
||||
step_class = getattr(module, class_name)
|
||||
return step_class, class_name
|
||||
except (ImportError, AttributeError) as e:
|
||||
raise ImportError(
|
||||
f"Failed to load processor step '{full_class_path}'. "
|
||||
f"Make sure the module '{module_path}' is installed and contains class '{class_name}'. "
|
||||
f"Consider registering the step using @ProcessorStepRegistry.register() for better portability. "
|
||||
if step_key in override_keys:
|
||||
override_keys.discard(step_key)
|
||||
|
||||
except Exception as e:
|
||||
step_name = step_entry.get("registry_name", step_entry.get("class", "Unknown"))
|
||||
raise ValueError(
|
||||
f"Failed to instantiate processor step '{step_name}' with config: {step_entry.get('config', {})}. "
|
||||
f"Error: {str(e)}"
|
||||
) from e
|
||||
|
||||
@classmethod
|
||||
def _instantiate_step(
|
||||
cls,
|
||||
step_entry: dict[str, Any],
|
||||
step_class: type[ProcessorStep],
|
||||
step_key: str,
|
||||
overrides: dict[str, Any],
|
||||
) -> ProcessorStep:
|
||||
"""Instantiate a single processor step with config overrides.
|
||||
# Load the step's state if a state file is specified.
|
||||
if "state_file" in step_entry and hasattr(step_instance, "load_state_dict"):
|
||||
# Check if state file exists locally first
|
||||
if base_path and (base_path / step_entry["state_file"]).exists():
|
||||
state_path = str(base_path / step_entry["state_file"])
|
||||
else:
|
||||
# Download the state file from the Hub.
|
||||
state_path = hf_hub_download(
|
||||
repo_id=model_id,
|
||||
filename=step_entry["state_file"],
|
||||
repo_type="model",
|
||||
force_download=force_download,
|
||||
resume_download=resume_download,
|
||||
proxies=proxies,
|
||||
token=token,
|
||||
cache_dir=cache_dir,
|
||||
local_files_only=local_files_only,
|
||||
revision=revision,
|
||||
)
|
||||
|
||||
This method handles the configuration merging and instantiation logic:
|
||||
step_instance.load_state_dict(load_file(state_path))
|
||||
|
||||
**Configuration Merging Strategy**:
|
||||
1. **Extract saved config**: Get step_entry.get("config", {}) from saved pipeline
|
||||
- Example: {"config": {"mean": 0.0, "std": 1.0}}
|
||||
2. **Extract user overrides**: Get overrides.get(step_key, {}) for this step
|
||||
- Example: overrides = {"NormalizeStep": {"mean": 2.0, "device": "cuda"}}
|
||||
3. **Merge with priority**: {**saved_cfg, **step_overrides}
|
||||
- **Override priority**: User values override saved values
|
||||
- **Result**: {"mean": 2.0, "std": 1.0, "device": "cuda"}
|
||||
steps.append(step_instance)
|
||||
|
||||
**Instantiation Process**:
|
||||
- **Call constructor**: step_class(**merged_cfg)
|
||||
- **Example**: NormalizeStep(mean=2.0, std=1.0, device="cuda")
|
||||
# Check for any unused override keys, which likely indicates a typo by the user.
|
||||
if override_keys:
|
||||
available_keys = [
|
||||
step.get("registry_name") or step["class"].rsplit(".", 1)[1]
|
||||
for step in loaded_config["steps"]
|
||||
]
|
||||
|
||||
**Error Handling**:
|
||||
- **Any exception during instantiation**: Convert to ValueError
|
||||
- **Include context**: step name, attempted config, original error
|
||||
- **Purpose**: Help users debug configuration issues
|
||||
- **Common causes**:
|
||||
a. Invalid parameter types (str instead of float)
|
||||
b. Missing required parameters
|
||||
c. Incompatible parameter combinations
|
||||
|
||||
Args:
|
||||
step_entry: The step configuration from saved config (contains "config" dict)
|
||||
step_class: The step class to instantiate (already resolved)
|
||||
step_key: The key used for overrides ("registry_name" or class name)
|
||||
overrides: User-provided parameter overrides (keyed by step_key)
|
||||
|
||||
Returns:
|
||||
The instantiated processor step (ready for use)
|
||||
|
||||
Raises:
|
||||
ValueError: If step cannot be instantiated, with detailed error context
|
||||
"""
|
||||
try:
|
||||
saved_cfg = step_entry.get("config", {})
|
||||
step_overrides = overrides.get(step_key, {})
|
||||
merged_cfg = {**saved_cfg, **step_overrides}
|
||||
return step_class(**merged_cfg)
|
||||
except Exception as e:
|
||||
step_name = step_entry.get("registry_name", step_entry.get("class", "Unknown"))
|
||||
raise ValueError(
|
||||
f"Failed to instantiate processor step '{step_name}' with config: {step_entry.get('config', {})}. "
|
||||
f"Error: {str(e)}"
|
||||
) from e
|
||||
|
||||
@classmethod
|
||||
def _load_step_state(
|
||||
cls,
|
||||
step_instance: ProcessorStep,
|
||||
step_entry: dict[str, Any],
|
||||
model_id: str,
|
||||
base_path: Path | None,
|
||||
hub_download_kwargs: dict[str, Any],
|
||||
) -> None:
|
||||
"""Load state dictionary for a processor step if available.
|
||||
|
||||
This method implements conditional state loading with local/Hub fallback:
|
||||
|
||||
**Precondition Checks** (early return if not met):
|
||||
1. **"state_file" in step_entry**: Step config specifies a state file
|
||||
- **If missing**: Step has no saved state (e.g., stateless transforms)
|
||||
2. **hasattr(step_instance, "load_state_dict")**: Step supports state loading
|
||||
- **If missing**: Step doesn't implement state loading (rare)
|
||||
|
||||
**State File Resolution Strategy**:
|
||||
1. **Local file priority**: Check base_path/state_filename exists
|
||||
- **Advantage**: Faster, no network calls
|
||||
- **Example**: "/models/my_model/normalize_step_0.safetensors"
|
||||
- **Use case**: Loading from local saved model directory
|
||||
|
||||
2. **Hub download fallback**: Download state file from repository
|
||||
- **When triggered**: Local file not found or base_path is None
|
||||
- **Process**: Use hf_hub_download with same parameters as config
|
||||
- **Example**: Download "normalize_step_0.safetensors" from "user/repo"
|
||||
- **Result**: Downloaded to local cache, path returned
|
||||
|
||||
**State Loading Process**:
|
||||
- **Load tensors**: Use safetensors.torch.load_file()
|
||||
- **Apply to step**: Call step_instance.load_state_dict(tensor_dict)
|
||||
- **In-place modification**: Updates step's internal tensor state
|
||||
|
||||
**Common state file examples**:
|
||||
- "normalize_step_0.safetensors" - normalization statistics
|
||||
- "custom_step_1.safetensors" - learned parameters
|
||||
- "tokenizer_step_2.safetensors" - vocabulary embeddings
|
||||
|
||||
Args:
|
||||
step_instance: The step instance to load state into (must have load_state_dict)
|
||||
step_entry: The step configuration dictionary (may contain "state_file")
|
||||
model_id: The model identifier (used for Hub downloads if needed)
|
||||
base_path: Local directory path for finding state files (None for Hub-only)
|
||||
hub_download_kwargs: Parameters for hf_hub_download (tokens, cache, etc.)
|
||||
|
||||
Note:
|
||||
This method modifies step_instance in-place and returns None.
|
||||
If state loading fails, exceptions from load_state_dict propagate.
|
||||
"""
|
||||
if "state_file" not in step_entry or not hasattr(step_instance, "load_state_dict"):
|
||||
return
|
||||
|
||||
state_filename = step_entry["state_file"]
|
||||
|
||||
# Try local file first
|
||||
if base_path and (base_path / state_filename).exists():
|
||||
state_path = str(base_path / state_filename)
|
||||
else:
|
||||
# Download from Hub
|
||||
state_path = hf_hub_download(
|
||||
repo_id=model_id,
|
||||
filename=state_filename,
|
||||
repo_type="model",
|
||||
**hub_download_kwargs,
|
||||
raise KeyError(
|
||||
f"Override keys {list(override_keys)} do not match any step in the saved configuration. "
|
||||
f"Available step keys: {available_keys}. "
|
||||
f"Make sure override keys match exact step class names or registry names."
|
||||
)
|
||||
|
||||
step_instance.load_state_dict(load_file(state_path))
|
||||
|
||||
@classmethod
|
||||
def _validate_overrides_used(
|
||||
cls, remaining_override_keys: set[str], loaded_config: dict[str, Any]
|
||||
) -> None:
|
||||
"""Validate that all provided overrides were used.
|
||||
|
||||
This method ensures user overrides are valid to catch typos and configuration errors:
|
||||
|
||||
**Validation Logic**:
|
||||
1. **If remaining_override_keys is empty**: All overrides were used -> Success
|
||||
- **Early return**: No validation needed
|
||||
- **Normal case**: User provided correct override keys
|
||||
|
||||
2. **If remaining_override_keys has entries**: Some overrides unused -> Error
|
||||
- **Root cause**: User provided keys that don't match any step
|
||||
- **Common issues**:
|
||||
a. Typos in step names ("NormalizStep" vs "NormalizeStep")
|
||||
b. Using wrong key type (class name vs registry name)
|
||||
c. Step doesn't exist in saved pipeline
|
||||
|
||||
**Helpful Error Generation**:
|
||||
- **Extract available keys**: Build list of valid override keys from config
|
||||
a. **Registry steps**: Use "registry_name" directly
|
||||
b. **Import steps**: Extract class name from "class" field
|
||||
- Example: "lerobot.processor.normalize.NormalizeStep" -> "NormalizeStep"
|
||||
- **Error message includes**:
|
||||
a. Invalid keys provided by user
|
||||
b. List of valid keys they can use
|
||||
c. Guidance about registry vs class names
|
||||
|
||||
**Override Key Resolution Rules**:
|
||||
- Steps with "registry_name": Use registry_name for overrides
|
||||
- Steps with "class": Use final class name for overrides
|
||||
- Users must match these exact keys in their overrides dict
|
||||
|
||||
Args:
|
||||
remaining_override_keys: Override keys that weren't matched to any step
|
||||
loaded_config: The loaded processor configuration (contains "steps" list)
|
||||
|
||||
Raises:
|
||||
KeyError: If any override keys were not used, with helpful error message
|
||||
"""
|
||||
if not remaining_override_keys:
|
||||
return
|
||||
|
||||
available_keys = [
|
||||
step.get("registry_name") or step["class"].rsplit(".", 1)[1] for step in loaded_config["steps"]
|
||||
]
|
||||
|
||||
raise KeyError(
|
||||
f"Override keys {list(remaining_override_keys)} do not match any step in the saved configuration. "
|
||||
f"Available step keys: {available_keys}. "
|
||||
f"Make sure override keys match exact step class names or registry names."
|
||||
# Construct and return the final pipeline instance.
|
||||
return cls(
|
||||
steps=steps,
|
||||
name=loaded_config.get("name", "DataProcessorPipeline"),
|
||||
to_transition=to_transition or batch_to_transition,
|
||||
to_output=to_output or cast(Callable[[EnvTransition], TOutput], transition_to_batch),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _should_suggest_migration(cls, model_path: Path) -> bool:
|
||||
"""Check if directory has JSON files but no processor configs.
|
||||
|
||||
This method implements smart migration detection to avoid false positives:
|
||||
|
||||
**Decision Logic**:
|
||||
1. **No JSON files found**: Return False
|
||||
- **Reason**: Empty directory or only non-config files
|
||||
- **Example**: Directory with only .safetensors, .md files
|
||||
- **Action**: No migration needed
|
||||
|
||||
2. **JSON files exist**: Analyze each file
|
||||
- **Goal**: Determine if ANY file is a valid processor config
|
||||
- **Process**:
|
||||
a. Try to parse each .json file
|
||||
b. Skip files with JSON parse errors (malformed)
|
||||
c. Check if parsed config passes _is_processor_config()
|
||||
- **If ANY valid processor found**: Return False (no migration)
|
||||
- **If NO valid processors found**: Return True (migration needed)
|
||||
|
||||
**Examples**:
|
||||
- **No migration**: ["processor.json", "config.json"] where processor.json is valid
|
||||
- **Migration needed**: ["config.json", "train.json"] where both are model configs
|
||||
- **No migration**: [] (empty directory)
|
||||
- **Migration needed**: ["old_model_config.json"] with old LeRobot format
|
||||
|
||||
**Why this works**:
|
||||
- **Precise detection**: Only suggests migration for actual old LeRobot models
|
||||
- **Avoids false positives**: Won't trigger on other HuggingFace model types
|
||||
- **Graceful handling**: Ignores malformed JSON files
|
||||
|
||||
Args:
|
||||
model_path: Path to local directory to analyze
|
||||
|
||||
Returns:
|
||||
True if directory has JSON configs but none are processor configs (migration needed)
|
||||
False if no JSON files or at least one valid processor config exists
|
||||
"""
|
||||
json_files = list(model_path.glob("*.json"))
|
||||
if len(json_files) == 0:
|
||||
return False
|
||||
|
||||
# Check if any JSON file is a processor config
|
||||
for json_file in json_files:
|
||||
try:
|
||||
with open(json_file) as f:
|
||||
config = json.load(f)
|
||||
|
||||
if cls._is_processor_config(config):
|
||||
return False # Found at least one processor config, no migration needed
|
||||
|
||||
except (json.JSONDecodeError, OSError):
|
||||
# Skip files that can't be parsed as JSON
|
||||
continue
|
||||
|
||||
# Have JSON files but no processor configs - suggest migration
|
||||
return True
|
||||
|
||||
@classmethod
|
||||
def _is_processor_config(cls, config: dict) -> bool:
|
||||
"""Check if config follows DataProcessorPipeline format.
|
||||
|
||||
This method validates the processor configuration structure:
|
||||
|
||||
**Required Structure Validation**:
|
||||
1. **"steps" field existence**: Must have top-level "steps" key
|
||||
- **If missing**: Not a processor config (e.g., model config, train config)
|
||||
- **Example invalid**: {"type": "act", "hidden_dim": 256}
|
||||
|
||||
2. **"steps" field type**: Must be a list, not other types
|
||||
- **If not list**: Invalid format
|
||||
- **Example invalid**: {"steps": "some_string"} or {"steps": {"key": "value"}}
|
||||
|
||||
3. **Empty steps validation**: Empty list is valid
|
||||
- **If len(steps) == 0**: Return True immediately
|
||||
- **Use case**: Empty processor pipeline (no-op)
|
||||
- **Example valid**: {"name": "EmptyProcessor", "steps": []}
|
||||
|
||||
**Individual Step Validation** (for non-empty steps):
|
||||
For each step in the steps list:
|
||||
1. **Step type**: Must be a dictionary
|
||||
- **If not dict**: Invalid step format
|
||||
- **Example invalid**: ["string_step", 123, true]
|
||||
|
||||
2. **Step identifier**: Must have either "class" OR "registry_name"
|
||||
- **"registry_name"**: Registered step (preferred)
|
||||
Example: {"registry_name": "normalize_step", "config": {...}}
|
||||
- **"class"**: Full import path
|
||||
Example: {"class": "lerobot.processor.normalize.NormalizeStep"}
|
||||
- **If neither**: Invalid step (can't resolve class)
|
||||
- **If both**: Also valid (registry_name takes precedence)
|
||||
|
||||
**Valid Processor Config Examples**:
|
||||
- {"steps": []} - Empty processor
|
||||
- {"steps": [{"registry_name": "normalize"}]} - Registry step
|
||||
- {"steps": [{"class": "my.module.Step"}]} - Import step
|
||||
- {"name": "MyProcessor", "steps": [...]} - With name
|
||||
|
||||
**Invalid Config Examples**:
|
||||
- {"type": "act"} - Missing "steps"
|
||||
- {"steps": "normalize"} - Steps not a list
|
||||
- {"steps": [{}]} - Step missing class/registry_name
|
||||
- {"steps": ["string"]} - Step not a dict
|
||||
|
||||
Args:
|
||||
config: The configuration dictionary to validate
|
||||
|
||||
Returns:
|
||||
True if config follows valid DataProcessorPipeline format, False otherwise
|
||||
"""
|
||||
# Must have a "steps" field with a list of step configurations
|
||||
if not isinstance(config.get("steps"), list):
|
||||
return False
|
||||
|
||||
steps = config["steps"]
|
||||
if len(steps) == 0:
|
||||
return True # Empty processor is valid
|
||||
|
||||
# Each step must be a dict with either "class" or "registry_name"
|
||||
for step in steps:
|
||||
if not isinstance(step, dict):
|
||||
return False
|
||||
if not ("class" in step or "registry_name" in step):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
@classmethod
|
||||
def _suggest_processor_migration(cls, model_path: str | Path, original_error: str) -> None:
|
||||
"""Raise migration error when we detect JSON files but no processor configs.
|
||||
|
||||
This method is called when migration detection determines that a model
|
||||
directory contains configuration files but none are valid processor configs.
|
||||
This typically indicates an old LeRobot model that needs migration.
|
||||
|
||||
**When this is called**:
|
||||
- User tries to load DataProcessorPipeline from local directory
|
||||
- Directory contains JSON configuration files
|
||||
- None of the JSON files follow processor config format
|
||||
- _should_suggest_migration() returned True
|
||||
|
||||
**Migration Command Generation**:
|
||||
- Constructs exact command user needs to run
|
||||
- Uses the migration script: migrate_policy_normalization.py
|
||||
- Includes the model path automatically
|
||||
- Example: "python src/lerobot/processor/migrate_policy_normalization.py --pretrained-path /models/old_model"
|
||||
|
||||
**Error Structure**:
|
||||
- **Always raises**: ProcessorMigrationError (never returns)
|
||||
- **Includes**: model_path, migration_command, original_error
|
||||
- **Purpose**: Force user attention to migration need
|
||||
- **User experience**: Clear actionable error with exact command to run
|
||||
|
||||
**Migration Process**:
|
||||
The suggested command will:
|
||||
1. Extract normalization stats from old model
|
||||
2. Create new processor configs (preprocessor + postprocessor)
|
||||
3. Remove normalization layers from model
|
||||
4. Save migrated model with processor pipeline
|
||||
|
||||
Args:
|
||||
model_path: Path to the model directory needing migration
|
||||
original_error: The error that triggered migration detection (for context)
|
||||
|
||||
Raises:
|
||||
ProcessorMigrationError: Always raised (this method never returns normally)
|
||||
"""
|
||||
migration_command = (
|
||||
f"python src/lerobot/processor/migrate_policy_normalization.py --pretrained-path {model_path}"
|
||||
)
|
||||
|
||||
raise ProcessorMigrationError(model_path, migration_command, original_error)
|
||||
|
||||
def __len__(self) -> int:
|
||||
"""Returns the number of steps in the pipeline."""
|
||||
return len(self.steps)
|
||||
|
||||
@@ -1,52 +0,0 @@
|
||||
from dataclasses import asdict, dataclass
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
|
||||
from lerobot.processor import ActionProcessorStep, PolicyAction, ProcessorStepRegistry, RobotAction
|
||||
|
||||
|
||||
@dataclass
|
||||
@ProcessorStepRegistry.register("robot_action_to_policy_action_processor")
|
||||
class RobotActionToPolicyActionProcessorStep(ActionProcessorStep):
|
||||
"""Processor step to map a dictionary to a tensor action."""
|
||||
|
||||
motor_names: list[str]
|
||||
|
||||
def action(self, action: RobotAction) -> PolicyAction:
|
||||
if len(self.motor_names) != len(action):
|
||||
raise ValueError(f"Action must have {len(self.motor_names)} elements, got {len(action)}")
|
||||
return torch.tensor([action[f"{name}.pos"] for name in self.motor_names])
|
||||
|
||||
def get_config(self) -> dict[str, Any]:
|
||||
return asdict(self)
|
||||
|
||||
def transform_features(self, features):
|
||||
features[PipelineFeatureType.ACTION]["action"] = PolicyFeature(
|
||||
type=FeatureType.ACTION, shape=(len(self.motor_names),)
|
||||
)
|
||||
return features
|
||||
|
||||
|
||||
@dataclass
|
||||
@ProcessorStepRegistry.register("policy_action_to_robot_action_processor")
|
||||
class PolicyActionToRobotActionProcessorStep(ActionProcessorStep):
|
||||
"""Processor step to map a policy action to a robot action."""
|
||||
|
||||
motor_names: list[str]
|
||||
|
||||
def action(self, action: PolicyAction) -> RobotAction:
|
||||
if len(self.motor_names) != len(action):
|
||||
raise ValueError(f"Action must have {len(self.motor_names)} elements, got {len(action)}")
|
||||
return {f"{name}.pos": action[i] for i, name in enumerate(self.motor_names)}
|
||||
|
||||
def get_config(self) -> dict[str, Any]:
|
||||
return asdict(self)
|
||||
|
||||
def transform_features(self, features):
|
||||
for name in self.motor_names:
|
||||
features[PipelineFeatureType.ACTION][f"{name}.pos"] = PolicyFeature(
|
||||
type=FeatureType.ACTION, shape=(1,)
|
||||
)
|
||||
return features
|
||||
+8
-13
@@ -21,12 +21,11 @@ Example:
|
||||
lerobot-record \
|
||||
--robot.type=so100_follower \
|
||||
--robot.port=/dev/tty.usbmodem58760431541 \
|
||||
--robot.cameras="{laptop: {type: opencv, index_or_path: 0, width: 640, height: 480, fps: 30}}" \
|
||||
--robot.cameras="{laptop: {type: opencv, camera_index: 0, width: 640, height: 480}}" \
|
||||
--robot.id=black \
|
||||
--dataset.repo_id=<my_username>/<my_dataset_name> \
|
||||
--dataset.repo_id=aliberts/record-test \
|
||||
--dataset.num_episodes=2 \
|
||||
--dataset.single_task="Grab the cube" \
|
||||
--display_data=true
|
||||
# <- Teleop optional if you want to teleoperate to record or in between episodes with a policy \
|
||||
# --teleop.type=so100_leader \
|
||||
# --teleop.port=/dev/tty.usbmodem58760431551 \
|
||||
@@ -236,12 +235,8 @@ def record_loop(
|
||||
robot: Robot,
|
||||
events: dict,
|
||||
fps: int,
|
||||
teleop_action_processor: RobotProcessorPipeline[
|
||||
tuple[RobotAction, RobotObservation], RobotAction
|
||||
], # runs after teleop
|
||||
robot_action_processor: RobotProcessorPipeline[
|
||||
tuple[RobotAction, RobotObservation], RobotAction
|
||||
], # runs before robot
|
||||
teleop_action_processor: RobotProcessorPipeline[RobotAction, RobotAction], # runs after teleop
|
||||
robot_action_processor: RobotProcessorPipeline[RobotAction, RobotAction], # runs before robot
|
||||
robot_observation_processor: RobotProcessorPipeline[
|
||||
RobotObservation, RobotObservation
|
||||
], # runs after robot
|
||||
@@ -327,7 +322,7 @@ def record_loop(
|
||||
act = teleop.get_action()
|
||||
|
||||
# Applies a pipeline to the raw teleop action, default is IdentityProcessor
|
||||
act_processed_teleop = teleop_action_processor((act, obs))
|
||||
act_processed_teleop = teleop_action_processor(act)
|
||||
|
||||
elif policy is None and isinstance(teleop, list):
|
||||
arm_action = teleop_arm.get_action()
|
||||
@@ -335,7 +330,7 @@ def record_loop(
|
||||
keyboard_action = teleop_keyboard.get_action()
|
||||
base_action = robot._from_keyboard_to_base_action(keyboard_action)
|
||||
act = {**arm_action, **base_action} if len(base_action) > 0 else arm_action
|
||||
act_processed_teleop = teleop_action_processor((act, obs))
|
||||
act_processed_teleop = teleop_action_processor(act)
|
||||
else:
|
||||
logging.info(
|
||||
"No policy or teleoperator provided, skipping action generation."
|
||||
@@ -347,10 +342,10 @@ def record_loop(
|
||||
# Applies a pipeline to the action, default is IdentityProcessor
|
||||
if policy is not None and act_processed_policy is not None:
|
||||
action_values = act_processed_policy
|
||||
robot_action_to_send = robot_action_processor((act_processed_policy, obs))
|
||||
robot_action_to_send = robot_action_processor(act_processed_policy)
|
||||
else:
|
||||
action_values = act_processed_teleop
|
||||
robot_action_to_send = robot_action_processor((act_processed_teleop, obs))
|
||||
robot_action_to_send = robot_action_processor(act_processed_teleop)
|
||||
|
||||
# Send action to robot
|
||||
# Action can eventually be clipped using `max_relative_target`,
|
||||
|
||||
@@ -23,7 +23,7 @@ lerobot-replay \
|
||||
--robot.port=/dev/tty.usbmodem58760431541 \
|
||||
--robot.id=black \
|
||||
--dataset.repo_id=aliberts/record-test \
|
||||
--dataset.episode=0
|
||||
--dataset.episode=2
|
||||
```
|
||||
|
||||
Example replay with bimanual so100:
|
||||
@@ -57,6 +57,7 @@ from lerobot.robots import ( # noqa: F401
|
||||
hope_jr,
|
||||
koch_follower,
|
||||
make_robot_from_config,
|
||||
reachy2,
|
||||
so100_follower,
|
||||
so101_follower,
|
||||
)
|
||||
@@ -112,9 +113,7 @@ def replay(cfg: ReplayConfig):
|
||||
for i, name in enumerate(dataset.features["action"]["names"]):
|
||||
action[name] = action_array[i]
|
||||
|
||||
robot_obs = robot.get_observation()
|
||||
|
||||
processed_action = robot_action_processor((action, robot_obs))
|
||||
processed_action = robot_action_processor(action)
|
||||
|
||||
_ = robot.send_action(processed_action)
|
||||
|
||||
|
||||
@@ -22,6 +22,7 @@ import numpy as np
|
||||
from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
|
||||
from lerobot.model.kinematics import RobotKinematics
|
||||
from lerobot.processor import (
|
||||
ComplementaryDataProcessorStep,
|
||||
EnvTransition,
|
||||
ObservationProcessorStep,
|
||||
ProcessorStep,
|
||||
@@ -30,6 +31,7 @@ from lerobot.processor import (
|
||||
RobotActionProcessorStep,
|
||||
TransitionKey,
|
||||
)
|
||||
from lerobot.robots.robot import Robot
|
||||
from lerobot.utils.rotation import Rotation
|
||||
|
||||
|
||||
@@ -66,46 +68,38 @@ class EEReferenceAndDelta(RobotActionProcessorStep):
|
||||
use_latched_reference: bool = (
|
||||
True # If True, latch reference on enable; if False, always use current pose
|
||||
)
|
||||
use_ik_solution: bool = False
|
||||
|
||||
reference_ee_pose: np.ndarray | None = field(default=None, init=False, repr=False)
|
||||
_prev_enabled: bool = field(default=False, init=False, repr=False)
|
||||
_command_when_disabled: np.ndarray | None = field(default=None, init=False, repr=False)
|
||||
|
||||
def action(self, action: RobotAction) -> RobotAction:
|
||||
observation = self.transition.get(TransitionKey.OBSERVATION).copy()
|
||||
new_action = action.copy()
|
||||
comp = self.transition.get(TransitionKey.COMPLEMENTARY_DATA)
|
||||
|
||||
if observation is None:
|
||||
raise ValueError("Joints observation is require for computing robot kinematics")
|
||||
|
||||
if self.use_ik_solution and "IK_solution" in self.transition.get(TransitionKey.COMPLEMENTARY_DATA):
|
||||
q_raw = self.transition.get(TransitionKey.COMPLEMENTARY_DATA)["IK_solution"]
|
||||
else:
|
||||
q_raw = np.array(
|
||||
[
|
||||
float(v)
|
||||
for k, v in observation.items()
|
||||
if isinstance(k, str)
|
||||
and k.endswith(".pos")
|
||||
and k.removesuffix(".pos") in self.motor_names
|
||||
],
|
||||
dtype=float,
|
||||
# Get joint positions from complimentary data
|
||||
raw = comp["raw_joint_positions"]
|
||||
if raw is None:
|
||||
raise ValueError(
|
||||
"raw_joint_positions is not in complementary data and is required for EEReferenceAndDelta"
|
||||
)
|
||||
|
||||
if q_raw is None:
|
||||
raise ValueError("Joints observation is require for computing robot kinematics")
|
||||
if "reference_joint_positions" in comp:
|
||||
q = comp["reference_joint_positions"]
|
||||
else:
|
||||
q = np.array([float(raw[n]) for n in self.motor_names], dtype=float)
|
||||
|
||||
# Current pose from FK on measured joints
|
||||
t_curr = self.kinematics.forward_kinematics(q_raw)
|
||||
t_curr = self.kinematics.forward_kinematics(q)
|
||||
|
||||
enabled = bool(action.pop("enabled"))
|
||||
tx = float(action.pop("target_x"))
|
||||
ty = float(action.pop("target_y"))
|
||||
tz = float(action.pop("target_z"))
|
||||
wx = float(action.pop("target_wx"))
|
||||
wy = float(action.pop("target_wy"))
|
||||
wz = float(action.pop("target_wz"))
|
||||
gripper_vel = float(action.pop("gripper_vel"))
|
||||
enabled = bool(new_action.pop("enabled"))
|
||||
tx = float(new_action.pop("target_x"))
|
||||
ty = float(new_action.pop("target_y"))
|
||||
tz = float(new_action.pop("target_z"))
|
||||
wx = float(new_action.pop("target_wx"))
|
||||
wy = float(new_action.pop("target_wy"))
|
||||
wz = float(new_action.pop("target_wz"))
|
||||
gripper_vel = float(new_action.pop("gripper_vel"))
|
||||
|
||||
desired = None
|
||||
|
||||
@@ -141,16 +135,16 @@ class EEReferenceAndDelta(RobotActionProcessorStep):
|
||||
# Write action fields
|
||||
pos = desired[:3, 3]
|
||||
tw = Rotation.from_matrix(desired[:3, :3]).as_rotvec()
|
||||
action["ee.x"] = float(pos[0])
|
||||
action["ee.y"] = float(pos[1])
|
||||
action["ee.z"] = float(pos[2])
|
||||
action["ee.wx"] = float(tw[0])
|
||||
action["ee.wy"] = float(tw[1])
|
||||
action["ee.wz"] = float(tw[2])
|
||||
action["ee.gripper_vel"] = gripper_vel
|
||||
new_action["ee.x"] = float(pos[0])
|
||||
new_action["ee.y"] = float(pos[1])
|
||||
new_action["ee.z"] = float(pos[2])
|
||||
new_action["ee.wx"] = float(tw[0])
|
||||
new_action["ee.wy"] = float(tw[1])
|
||||
new_action["ee.wz"] = float(tw[2])
|
||||
new_action["ee.gripper_vel"] = gripper_vel
|
||||
|
||||
self._prev_enabled = enabled
|
||||
return action
|
||||
return new_action
|
||||
|
||||
def reset(self):
|
||||
"""Resets the internal state of the processor."""
|
||||
@@ -256,7 +250,7 @@ class EEBoundsAndSafety(RobotActionProcessorStep):
|
||||
|
||||
@ProcessorStepRegistry.register("inverse_kinematics_ee_to_joints")
|
||||
@dataclass
|
||||
class InverseKinematicsEEToJoints(RobotActionProcessorStep):
|
||||
class InverseKinematicsEEToJoints(ProcessorStep):
|
||||
"""
|
||||
Computes desired joint positions from a target end-effector pose using inverse kinematics (IK).
|
||||
|
||||
@@ -276,36 +270,40 @@ class InverseKinematicsEEToJoints(RobotActionProcessorStep):
|
||||
q_curr: np.ndarray | None = field(default=None, init=False, repr=False)
|
||||
initial_guess_current_joints: bool = True
|
||||
|
||||
def action(self, action: RobotAction) -> RobotAction:
|
||||
x = action.pop("ee.x")
|
||||
y = action.pop("ee.y")
|
||||
z = action.pop("ee.z")
|
||||
wx = action.pop("ee.wx")
|
||||
wy = action.pop("ee.wy")
|
||||
wz = action.pop("ee.wz")
|
||||
gripper_pos = action.pop("ee.gripper_pos")
|
||||
def __call__(self, transition: EnvTransition) -> EnvTransition:
|
||||
new_transition = transition.copy()
|
||||
act = new_transition.get(TransitionKey.ACTION).copy()
|
||||
|
||||
if not isinstance(act, dict):
|
||||
raise ValueError(f"Action should be a RobotAction type got {type(act)}")
|
||||
|
||||
comp = new_transition.get(TransitionKey.COMPLEMENTARY_DATA) or {}
|
||||
|
||||
x = act.pop("ee.x")
|
||||
y = act.pop("ee.y")
|
||||
z = act.pop("ee.z")
|
||||
wx = act.pop("ee.wx")
|
||||
wy = act.pop("ee.wy")
|
||||
wz = act.pop("ee.wz")
|
||||
gripper_pos = act.pop("ee.gripper_pos")
|
||||
|
||||
if None in (x, y, z, wx, wy, wz, gripper_pos):
|
||||
raise ValueError(
|
||||
"Missing required end-effector pose components: ee.x, ee.y, ee.z, ee.wx, ee.wy, ee.wz, ee.gripper_pos must all be present in action"
|
||||
)
|
||||
|
||||
observation = self.transition.get(TransitionKey.OBSERVATION).copy()
|
||||
if observation is None:
|
||||
raise ValueError("Joints observation is require for computing robot kinematics")
|
||||
|
||||
q_raw = np.array(
|
||||
[float(v) for k, v in observation.items() if isinstance(k, str) and k.endswith(".pos")],
|
||||
dtype=float,
|
||||
)
|
||||
if q_raw is None:
|
||||
raise ValueError("Joints observation is require for computing robot kinematics")
|
||||
# Get joint positions from complimentary data
|
||||
raw = comp["raw_joint_positions"]
|
||||
if raw is None:
|
||||
raise ValueError(
|
||||
"raw_joint_positions is not in complementary data and is required for EEReferenceAndDelta"
|
||||
)
|
||||
|
||||
if self.initial_guess_current_joints: # Use current joints as initial guess
|
||||
self.q_curr = q_raw
|
||||
self.q_curr = np.array([float(raw[n]) for n in self.motor_names], dtype=float)
|
||||
else: # Use previous ik solution as initial guess
|
||||
if self.q_curr is None:
|
||||
self.q_curr = q_raw
|
||||
self.q_curr = np.array([float(raw[n]) for n in self.motor_names], dtype=float)
|
||||
|
||||
# Build desired 4x4 transform from pos + rotvec (twist)
|
||||
t_des = np.eye(4, dtype=float)
|
||||
@@ -316,14 +314,17 @@ class InverseKinematicsEEToJoints(RobotActionProcessorStep):
|
||||
q_target = self.kinematics.inverse_kinematics(self.q_curr, t_des)
|
||||
self.q_curr = q_target
|
||||
|
||||
new_act = dict(act)
|
||||
# TODO: This is sentitive to order of motor_names = q_target mapping
|
||||
for i, name in enumerate(self.motor_names):
|
||||
if name != "gripper":
|
||||
action[f"{name}.pos"] = float(q_target[i])
|
||||
new_act[f"{name}.pos"] = float(q_target[i])
|
||||
else:
|
||||
action["gripper.pos"] = float(gripper_pos)
|
||||
|
||||
return action
|
||||
new_act["gripper.pos"] = float(gripper_pos)
|
||||
new_transition[TransitionKey.ACTION] = new_act
|
||||
if not self.initial_guess_current_joints:
|
||||
new_transition[TransitionKey.COMPLEMENTARY_DATA]["reference_joint_positions"] = q_target
|
||||
return new_transition
|
||||
|
||||
def transform_features(
|
||||
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||
@@ -367,30 +368,29 @@ class GripperVelocityToJoint(RobotActionProcessorStep):
|
||||
discrete_gripper: bool = False
|
||||
|
||||
def action(self, action: RobotAction) -> RobotAction:
|
||||
observation = self.transition.get(TransitionKey.OBSERVATION).copy()
|
||||
complementary_data = self.transition.get(TransitionKey.COMPLEMENTARY_DATA)
|
||||
|
||||
gripper_vel = action.pop("ee.gripper_vel")
|
||||
|
||||
if observation is None:
|
||||
raise ValueError("Joints observation is require for computing robot kinematics")
|
||||
if "raw_joint_positions" not in complementary_data:
|
||||
raise ValueError(
|
||||
"raw_joint_positions is not in complementary data and is required for GripperVelocityToJoint"
|
||||
)
|
||||
|
||||
q_raw = np.array(
|
||||
[float(v) for k, v in observation.items() if isinstance(k, str) and k.endswith(".pos")],
|
||||
dtype=float,
|
||||
)
|
||||
if q_raw is None:
|
||||
raise ValueError("Joints observation is require for computing robot kinematics")
|
||||
curr_gripper_pos = complementary_data["raw_joint_positions"]["gripper"]
|
||||
|
||||
if self.discrete_gripper:
|
||||
# Discrete gripper actions are in [0, 1, 2]
|
||||
# 0: open, 1: close, 2: stay
|
||||
# We need to shift them to [-1, 0, 1] and then scale them to clip_max
|
||||
gripper_vel = (gripper_vel - 1) * self.clip_max
|
||||
# TODO(Michel,Adil): Fix this logic
|
||||
# if self.discrete_gripper:
|
||||
# # Discrete gripper actions are in [0, 1, 2]
|
||||
# # 0: open, 1: close, 2: stay
|
||||
# # We need to shift them to [-1, 0, 1] and then scale them to clip_max
|
||||
# gripper_action = gripper_vel
|
||||
# gripper_action *= self.clip_max
|
||||
# action["ee.gripper_pos"] = gripper_action
|
||||
|
||||
# Compute desired gripper position
|
||||
delta = gripper_vel * float(self.speed_factor)
|
||||
# TODO: This assumes gripper is the last specified joint in the robot
|
||||
gripper_pos = float(np.clip(q_raw[-1] + delta, self.clip_min, self.clip_max))
|
||||
gripper_pos = float(np.clip(curr_gripper_pos + delta, self.clip_min, self.clip_max))
|
||||
action["ee.gripper_pos"] = gripper_pos
|
||||
|
||||
return action
|
||||
@@ -494,6 +494,41 @@ class ForwardKinematicsJointsToEEAction(RobotActionProcessorStep):
|
||||
return features
|
||||
|
||||
|
||||
@ProcessorStepRegistry.register("add_robot_observation")
|
||||
@dataclass
|
||||
class AddRobotObservationAsComplimentaryData(ComplementaryDataProcessorStep):
|
||||
"""
|
||||
Reads the robot's current observation and adds it to the transition's complementary data.
|
||||
|
||||
This step acts as a bridge to the physical robot, injecting its real-time sensor readings
|
||||
(like raw joint positions) into the data processing pipeline. This data is then available
|
||||
for other processing steps.
|
||||
|
||||
Attributes:
|
||||
robot: An instance of a `Robot` class used to get observations from hardware.
|
||||
"""
|
||||
|
||||
robot: Robot
|
||||
|
||||
def complementary_data(self, comp: dict | None) -> dict:
|
||||
new_comp = dict(comp)
|
||||
obs = (
|
||||
self.robot.get_observation()
|
||||
) # todo(steven): why not self.trtansition.get(TransitionKey.OBSERVATION)?
|
||||
|
||||
new_comp["raw_joint_positions"] = {
|
||||
k.removesuffix(".pos"): float(v)
|
||||
for k, v in obs.items()
|
||||
if isinstance(k, str) and k.endswith(".pos")
|
||||
}
|
||||
return new_comp
|
||||
|
||||
def transform_features(
|
||||
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
||||
return features
|
||||
|
||||
|
||||
@ProcessorStepRegistry.register(name="forward_kinematics_joints_to_ee")
|
||||
@dataclass
|
||||
class ForwardKinematicsJointsToEE(ProcessorStep):
|
||||
@@ -523,94 +558,3 @@ class ForwardKinematicsJointsToEE(ProcessorStep):
|
||||
if features[PipelineFeatureType.OBSERVATION] is not None:
|
||||
features = self.joints_to_ee_observation_processor.transform_features(features)
|
||||
return features
|
||||
|
||||
|
||||
@ProcessorStepRegistry.register("inverse_kinematics_rl_step")
|
||||
@dataclass
|
||||
class InverseKinematicsRLStep(ProcessorStep):
|
||||
"""
|
||||
Computes desired joint positions from a target end-effector pose using inverse kinematics (IK).
|
||||
|
||||
This is modified from the InverseKinematicsEEToJoints step to be used in the RL pipeline.
|
||||
"""
|
||||
|
||||
kinematics: RobotKinematics
|
||||
motor_names: list[str]
|
||||
q_curr: np.ndarray | None = field(default=None, init=False, repr=False)
|
||||
initial_guess_current_joints: bool = True
|
||||
|
||||
def __call__(self, transition: EnvTransition) -> EnvTransition:
|
||||
new_transition = dict(transition)
|
||||
action = new_transition.get(TransitionKey.ACTION)
|
||||
if action is None:
|
||||
raise ValueError("Action is required for InverseKinematicsEEToJoints")
|
||||
action = dict(action)
|
||||
|
||||
x = action.pop("ee.x")
|
||||
y = action.pop("ee.y")
|
||||
z = action.pop("ee.z")
|
||||
wx = action.pop("ee.wx")
|
||||
wy = action.pop("ee.wy")
|
||||
wz = action.pop("ee.wz")
|
||||
gripper_pos = action.pop("ee.gripper_pos")
|
||||
|
||||
if None in (x, y, z, wx, wy, wz, gripper_pos):
|
||||
raise ValueError(
|
||||
"Missing required end-effector pose components: ee.x, ee.y, ee.z, ee.wx, ee.wy, ee.wz, ee.gripper_pos must all be present in action"
|
||||
)
|
||||
|
||||
observation = new_transition.get(TransitionKey.OBSERVATION).copy()
|
||||
if observation is None:
|
||||
raise ValueError("Joints observation is require for computing robot kinematics")
|
||||
|
||||
q_raw = np.array(
|
||||
[float(v) for k, v in observation.items() if isinstance(k, str) and k.endswith(".pos")],
|
||||
dtype=float,
|
||||
)
|
||||
if q_raw is None:
|
||||
raise ValueError("Joints observation is require for computing robot kinematics")
|
||||
|
||||
if self.initial_guess_current_joints: # Use current joints as initial guess
|
||||
self.q_curr = q_raw
|
||||
else: # Use previous ik solution as initial guess
|
||||
if self.q_curr is None:
|
||||
self.q_curr = q_raw
|
||||
|
||||
# Build desired 4x4 transform from pos + rotvec (twist)
|
||||
t_des = np.eye(4, dtype=float)
|
||||
t_des[:3, :3] = Rotation.from_rotvec([wx, wy, wz]).as_matrix()
|
||||
t_des[:3, 3] = [x, y, z]
|
||||
|
||||
# Compute inverse kinematics
|
||||
q_target = self.kinematics.inverse_kinematics(self.q_curr, t_des)
|
||||
self.q_curr = q_target
|
||||
|
||||
# TODO: This is sentitive to order of motor_names = q_target mapping
|
||||
for i, name in enumerate(self.motor_names):
|
||||
if name != "gripper":
|
||||
action[f"{name}.pos"] = float(q_target[i])
|
||||
else:
|
||||
action["gripper.pos"] = float(gripper_pos)
|
||||
|
||||
new_transition[TransitionKey.ACTION] = action
|
||||
complementary_data = new_transition.get(TransitionKey.COMPLEMENTARY_DATA, {})
|
||||
complementary_data["IK_solution"] = q_target
|
||||
new_transition[TransitionKey.COMPLEMENTARY_DATA] = complementary_data
|
||||
return new_transition
|
||||
|
||||
def transform_features(
|
||||
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
||||
for feat in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]:
|
||||
features[PipelineFeatureType.ACTION].pop(f"ee.{feat}", None)
|
||||
|
||||
for name in self.motor_names:
|
||||
features[PipelineFeatureType.ACTION][f"{name}.pos"] = PolicyFeature(
|
||||
type=FeatureType.ACTION, shape=(1,)
|
||||
)
|
||||
|
||||
return features
|
||||
|
||||
def reset(self):
|
||||
"""Resets the initial guess for the IK solver."""
|
||||
self.q_curr = None
|
||||
|
||||
@@ -170,4 +170,8 @@ python lerobot/scripts/control_robot.py \
|
||||
--control.episode=0
|
||||
```
|
||||
|
||||
Follow [previous tutorial](https://github.com/huggingface/lerobot/blob/main/examples/7_get_started_with_real_robot.md#4-train-a-policy-on-your-data) to train a policy on your data and run inference on your robot. You will need to adapt the code for Stretch.
|
||||
|
||||
> TODO(rcadene, aliberts): Add already setup environment and policy yaml configuration files
|
||||
|
||||
If you need help, please reach out on Discord in the channel `#stretch3-mobile-arm`.
|
||||
|
||||
@@ -118,7 +118,7 @@ echo ${HF_USER}/aloha_test
|
||||
If you didn't upload with `--control.push_to_hub=false`, you can also visualize it locally with [Rerun](https://github.com/rerun-io/rerun):
|
||||
|
||||
```bash
|
||||
lerobot-dataset-viz \
|
||||
python -m lerobot.scripts.visualize_dataset \
|
||||
--repo-id ${HF_USER}/aloha_test --episode 0
|
||||
```
|
||||
|
||||
@@ -193,4 +193,6 @@ As you can see, it's almost the same command as previously used to record your t
|
||||
|
||||
## More
|
||||
|
||||
Follow this [previous tutorial](https://github.com/huggingface/lerobot/blob/main/examples/7_get_started_with_real_robot.md#4-train-a-policy-on-your-data) for a more in-depth explanation.
|
||||
|
||||
If you have any question or need help, please reach out on Discord in the channel `#aloha-arm`.
|
||||
|
||||
@@ -0,0 +1,90 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2024 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.
|
||||
|
||||
"""Use this script to get a quick summary of your system config.
|
||||
It should be able to run without any of LeRobot's dependencies or LeRobot itself installed.
|
||||
"""
|
||||
|
||||
import platform
|
||||
|
||||
HAS_HF_HUB = True
|
||||
HAS_HF_DATASETS = True
|
||||
HAS_NP = True
|
||||
HAS_TORCH = True
|
||||
HAS_LEROBOT = True
|
||||
|
||||
try:
|
||||
import huggingface_hub
|
||||
except ImportError:
|
||||
HAS_HF_HUB = False
|
||||
|
||||
try:
|
||||
import datasets
|
||||
except ImportError:
|
||||
HAS_HF_DATASETS = False
|
||||
|
||||
try:
|
||||
import numpy as np
|
||||
except ImportError:
|
||||
HAS_NP = False
|
||||
|
||||
try:
|
||||
import torch
|
||||
except ImportError:
|
||||
HAS_TORCH = False
|
||||
|
||||
try:
|
||||
import lerobot
|
||||
except ImportError:
|
||||
HAS_LEROBOT = False
|
||||
|
||||
|
||||
lerobot_version = lerobot.__version__ if HAS_LEROBOT else "N/A"
|
||||
hf_hub_version = huggingface_hub.__version__ if HAS_HF_HUB else "N/A"
|
||||
hf_datasets_version = datasets.__version__ if HAS_HF_DATASETS else "N/A"
|
||||
np_version = np.__version__ if HAS_NP else "N/A"
|
||||
|
||||
torch_version = torch.__version__ if HAS_TORCH else "N/A"
|
||||
torch_cuda_available = torch.cuda.is_available() if HAS_TORCH else "N/A"
|
||||
cuda_version = torch._C._cuda_getCompiledVersion() if HAS_TORCH and torch.version.cuda is not None else "N/A"
|
||||
|
||||
|
||||
# TODO(aliberts): refactor into an actual command `lerobot env`
|
||||
def display_sys_info() -> dict:
|
||||
"""Run this to get basic system info to help for tracking issues & bugs."""
|
||||
info = {
|
||||
"`lerobot` version": lerobot_version,
|
||||
"Platform": platform.platform(),
|
||||
"Python version": platform.python_version(),
|
||||
"Huggingface_hub version": hf_hub_version,
|
||||
"Dataset version": hf_datasets_version,
|
||||
"Numpy version": np_version,
|
||||
"PyTorch version (GPU?)": f"{torch_version} ({torch_cuda_available})",
|
||||
"Cuda version": cuda_version,
|
||||
"Using GPU in script?": "<fill in>",
|
||||
# "Using distributed or parallel set-up in script?": "<fill in>",
|
||||
}
|
||||
print("\nCopy-and-paste the text below in your GitHub issue and FILL OUT the last point.\n")
|
||||
print(format_dict(info))
|
||||
return info
|
||||
|
||||
|
||||
def format_dict(d: dict) -> str:
|
||||
return "\n".join([f"- {prop}: {val}" for prop, val in d.items()]) + "\n"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
display_sys_info()
|
||||
+197
-190
@@ -52,11 +52,10 @@ import logging
|
||||
import threading
|
||||
import time
|
||||
from collections import defaultdict
|
||||
from collections.abc import Callable
|
||||
from collections.abc import Callable, Iterator
|
||||
from contextlib import nullcontext
|
||||
from copy import deepcopy
|
||||
from dataclasses import asdict
|
||||
from functools import partial
|
||||
from pathlib import Path
|
||||
from pprint import pformat
|
||||
from typing import Any, TypedDict
|
||||
@@ -164,11 +163,11 @@ def rollout(
|
||||
# Infer "task" from attributes of environments.
|
||||
# TODO: works with SyncVectorEnv but not AsyncVectorEnv
|
||||
observation = add_envs_task(env, observation)
|
||||
|
||||
observation = preprocessor(observation)
|
||||
with torch.inference_mode():
|
||||
action = policy.select_action(observation)
|
||||
action = postprocessor(action)
|
||||
|
||||
# Convert to CPU / numpy.
|
||||
action_numpy: np.ndarray = action.to("cpu").numpy()
|
||||
assert action_numpy.ndim == 2, "Action dimensions should be (batch, action_dim)"
|
||||
@@ -232,9 +231,9 @@ def rollout(
|
||||
def eval_policy(
|
||||
env: gym.vector.VectorEnv,
|
||||
policy: PreTrainedPolicy,
|
||||
preprocessor: PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
|
||||
postprocessor: PolicyProcessorPipeline[PolicyAction, PolicyAction],
|
||||
n_episodes: int,
|
||||
preprocessor: PolicyProcessorPipeline[dict[str, Any], dict[str, Any]] | None = None,
|
||||
postprocessor: PolicyProcessorPipeline[PolicyAction, PolicyAction] | None = None,
|
||||
max_episodes_rendered: int = 0,
|
||||
videos_dir: Path | None = None,
|
||||
return_episode_data: bool = False,
|
||||
@@ -418,7 +417,6 @@ def eval_policy(
|
||||
"eval_ep_s": (time.time() - start) / n_episodes,
|
||||
},
|
||||
}
|
||||
|
||||
if return_episode_data:
|
||||
info["episodes"] = episode_data
|
||||
|
||||
@@ -479,7 +477,6 @@ def eval_main(cfg: EvalPipelineConfig):
|
||||
|
||||
# Check device is available
|
||||
device = get_safe_torch_device(cfg.policy.device, log=True)
|
||||
|
||||
torch.backends.cudnn.benchmark = True
|
||||
torch.backends.cuda.matmul.allow_tf32 = True
|
||||
set_seed(cfg.seed)
|
||||
@@ -490,7 +487,6 @@ def eval_main(cfg: EvalPipelineConfig):
|
||||
envs = make_env(cfg.env, n_envs=cfg.eval.batch_size, use_async_envs=cfg.eval.use_async_envs)
|
||||
|
||||
logging.info("Making policy.")
|
||||
|
||||
policy = make_policy(
|
||||
cfg=cfg.policy,
|
||||
env_cfg=cfg.env,
|
||||
@@ -498,10 +494,7 @@ def eval_main(cfg: EvalPipelineConfig):
|
||||
|
||||
policy.eval()
|
||||
preprocessor, postprocessor = make_pre_post_processors(
|
||||
policy_cfg=cfg.policy,
|
||||
pretrained_path=cfg.policy.pretrained_path,
|
||||
# The inference device is automatically set to match the detected hardware, overriding any previous device settings from training to ensure compatibility.
|
||||
preprocessor_overrides={"device_processor": {"device": str(policy.config.device)}},
|
||||
policy_cfg=cfg.policy, pretrained_path=cfg.policy.pretrained_path
|
||||
)
|
||||
with torch.no_grad(), torch.autocast(device_type=device.type) if cfg.policy.use_amp else nullcontext():
|
||||
info = eval_policy_all(
|
||||
@@ -514,14 +507,17 @@ def eval_main(cfg: EvalPipelineConfig):
|
||||
videos_dir=Path(cfg.output_dir) / "videos",
|
||||
start_seed=cfg.seed,
|
||||
max_parallel_tasks=cfg.env.max_parallel_tasks,
|
||||
verbose=False,
|
||||
)
|
||||
print("Overall Aggregated Metrics:")
|
||||
print(info["overall"])
|
||||
print(info["overall"]["aggregated"])
|
||||
|
||||
# Print per-suite stats
|
||||
for task_group, task_group_info in info.items():
|
||||
if task_group == "overall":
|
||||
continue # Skip the overall stats since we already printed it
|
||||
print(f"\nAggregated Metrics for {task_group}:")
|
||||
print(task_group_info)
|
||||
print(task_group_info["aggregated"])
|
||||
# Close all vec envs
|
||||
close_envs(envs)
|
||||
|
||||
@@ -543,205 +539,216 @@ class TaskMetrics(TypedDict):
|
||||
ACC_KEYS = ("sum_rewards", "max_rewards", "successes", "video_paths")
|
||||
|
||||
|
||||
def eval_one(
|
||||
env: gym.vector.VectorEnv,
|
||||
*,
|
||||
def eval_policy_all(
|
||||
envs: dict[str, dict[int, gym.vector.VectorEnv]],
|
||||
policy: PreTrainedPolicy,
|
||||
preprocessor: PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
|
||||
postprocessor: PolicyProcessorPipeline[PolicyAction, PolicyAction],
|
||||
n_episodes: int,
|
||||
max_episodes_rendered: int,
|
||||
videos_dir: Path | None,
|
||||
return_episode_data: bool,
|
||||
start_seed: int | None,
|
||||
) -> TaskMetrics:
|
||||
"""Evaluates one task_id of one suite using the provided vec env."""
|
||||
|
||||
task_videos_dir = videos_dir
|
||||
|
||||
task_result = eval_policy(
|
||||
env=env,
|
||||
policy=policy,
|
||||
preprocessor=preprocessor,
|
||||
postprocessor=postprocessor,
|
||||
n_episodes=n_episodes,
|
||||
max_episodes_rendered=max_episodes_rendered,
|
||||
videos_dir=task_videos_dir,
|
||||
return_episode_data=return_episode_data,
|
||||
start_seed=start_seed,
|
||||
)
|
||||
|
||||
per_episode = task_result["per_episode"]
|
||||
return TaskMetrics(
|
||||
sum_rewards=[ep["sum_reward"] for ep in per_episode],
|
||||
max_rewards=[ep["max_reward"] for ep in per_episode],
|
||||
successes=[ep["success"] for ep in per_episode],
|
||||
video_paths=task_result.get("video_paths", []),
|
||||
)
|
||||
|
||||
|
||||
def run_one(
|
||||
task_group: str,
|
||||
task_id: int,
|
||||
env,
|
||||
*,
|
||||
policy,
|
||||
preprocessor,
|
||||
postprocessor,
|
||||
n_episodes: int,
|
||||
max_episodes_rendered: int,
|
||||
videos_dir: Path | None,
|
||||
return_episode_data: bool,
|
||||
start_seed: int | None,
|
||||
):
|
||||
"""
|
||||
Run eval_one for a single (task_group, task_id, env).
|
||||
Returns (task_group, task_id, task_metrics_dict).
|
||||
This function is intentionally module-level to make it easy to test.
|
||||
"""
|
||||
task_videos_dir = None
|
||||
if videos_dir is not None:
|
||||
task_videos_dir = videos_dir / f"{task_group}_{task_id}"
|
||||
task_videos_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Call the existing eval_one (assumed to return TaskMetrics-like dict)
|
||||
metrics = eval_one(
|
||||
env,
|
||||
policy=policy,
|
||||
preprocessor=preprocessor,
|
||||
postprocessor=postprocessor,
|
||||
n_episodes=n_episodes,
|
||||
max_episodes_rendered=max_episodes_rendered,
|
||||
videos_dir=task_videos_dir,
|
||||
return_episode_data=return_episode_data,
|
||||
start_seed=start_seed,
|
||||
)
|
||||
# ensure we always provide video_paths key to simplify accumulation
|
||||
if max_episodes_rendered > 0:
|
||||
metrics.setdefault("video_paths", [])
|
||||
return task_group, task_id, metrics
|
||||
|
||||
|
||||
def eval_policy_all(
|
||||
envs: dict[str, dict[int, gym.vector.VectorEnv]],
|
||||
policy,
|
||||
preprocessor: PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
|
||||
postprocessor: PolicyProcessorPipeline[PolicyAction, PolicyAction],
|
||||
n_episodes: int,
|
||||
*,
|
||||
max_episodes_rendered: int = 0,
|
||||
videos_dir: Path | None = None,
|
||||
return_episode_data: bool = False,
|
||||
start_seed: int | None = None,
|
||||
max_parallel_tasks: int = 1,
|
||||
verbose: bool = True,
|
||||
) -> dict:
|
||||
"""
|
||||
Evaluate a nested `envs` dict: {task_group: {task_id: vec_env}}.
|
||||
This implementation flattens tasks, runs them sequentially or via ThreadPoolExecutor,
|
||||
accumulates per-group and overall statistics, and returns the same aggregate metrics
|
||||
schema as the single-env evaluator (avg_sum_reward / avg_max_reward / pc_success / timings)
|
||||
plus per-task infos.
|
||||
Evaluate a policy over a dict-of-dicts of vectorized envs:
|
||||
envs[suite_name][task_id] -> gym.vector.VectorEnv
|
||||
Returns a dict with per-suite aggregates and an 'overall' block.
|
||||
"""
|
||||
start_t = time.time()
|
||||
global_start = time.time()
|
||||
|
||||
# Flatten envs into list of (task_group, task_id, env)
|
||||
tasks = [(tg, tid, vec) for tg, group in envs.items() for tid, vec in group.items()]
|
||||
# inner: evaluate a single (suite, task)
|
||||
def eval_one(
|
||||
task_group: str,
|
||||
task_id: int,
|
||||
env: gym.vector.VectorEnv,
|
||||
*,
|
||||
policy: PreTrainedPolicy,
|
||||
preprocessor: PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
|
||||
postprocessor: PolicyProcessorPipeline[PolicyAction, PolicyAction],
|
||||
n_episodes: int,
|
||||
max_episodes_rendered: int,
|
||||
videos_dir: Path | None,
|
||||
return_episode_data: bool,
|
||||
start_seed: int | None,
|
||||
) -> TaskMetrics:
|
||||
"""Evaluates one task_id of one suite using the provided vec env."""
|
||||
if verbose:
|
||||
print(f"Evaluating: task_group={task_group}, task_id={task_id} ...")
|
||||
|
||||
# accumulators: track metrics at both per-group level and across all groups
|
||||
task_videos_dir = None
|
||||
if videos_dir is not None:
|
||||
task_videos_dir = videos_dir / f"{task_group}_{task_id}"
|
||||
task_videos_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
task_result = eval_policy(
|
||||
env=env,
|
||||
policy=policy,
|
||||
preprocessor=preprocessor,
|
||||
postprocessor=postprocessor,
|
||||
n_episodes=n_episodes,
|
||||
max_episodes_rendered=max_episodes_rendered,
|
||||
videos_dir=task_videos_dir,
|
||||
return_episode_data=return_episode_data,
|
||||
start_seed=start_seed,
|
||||
)
|
||||
|
||||
per_episode = task_result["per_episode"]
|
||||
return TaskMetrics(
|
||||
sum_rewards=[ep["sum_reward"] for ep in per_episode],
|
||||
max_rewards=[ep["max_reward"] for ep in per_episode],
|
||||
successes=[ep["success"] for ep in per_episode],
|
||||
video_paths=task_result.get("video_paths", []),
|
||||
)
|
||||
|
||||
def _eval_monotask(
|
||||
envs, policy, n_episodes, max_episodes_rendered, videos_dir, return_episode_data, start_seed
|
||||
):
|
||||
for task_group, tasks in envs.items():
|
||||
for task_id, vec in tasks.items():
|
||||
yield (
|
||||
task_group,
|
||||
task_id,
|
||||
eval_one(
|
||||
task_group,
|
||||
task_id,
|
||||
vec,
|
||||
policy=policy,
|
||||
preprocessor=preprocessor,
|
||||
postprocessor=postprocessor,
|
||||
n_episodes=n_episodes,
|
||||
max_episodes_rendered=max_episodes_rendered,
|
||||
videos_dir=videos_dir,
|
||||
return_episode_data=return_episode_data,
|
||||
start_seed=start_seed,
|
||||
),
|
||||
)
|
||||
|
||||
def _eval_parallel(
|
||||
envs,
|
||||
policy,
|
||||
preprocessor: PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
|
||||
postprocessor: PolicyProcessorPipeline[PolicyAction, PolicyAction],
|
||||
n_episodes,
|
||||
max_episodes_rendered,
|
||||
videos_dir,
|
||||
return_episode_data,
|
||||
start_seed,
|
||||
max_parallel_tasks,
|
||||
):
|
||||
with cf.ThreadPoolExecutor(max_workers=max_parallel_tasks) as executor:
|
||||
fut2key: dict[cf.Future, tuple[str, int]] = {}
|
||||
for task_group, tasks in envs.items():
|
||||
for task_id, vec in tasks.items():
|
||||
fut = executor.submit(
|
||||
eval_one,
|
||||
task_group,
|
||||
task_id,
|
||||
vec,
|
||||
policy=policy,
|
||||
preprocessor=preprocessor,
|
||||
postprocessor=postprocessor,
|
||||
n_episodes=n_episodes,
|
||||
max_episodes_rendered=max_episodes_rendered,
|
||||
videos_dir=videos_dir,
|
||||
return_episode_data=return_episode_data,
|
||||
start_seed=start_seed,
|
||||
)
|
||||
fut2key[fut] = (task_group, task_id)
|
||||
for fut in cf.as_completed(fut2key):
|
||||
task_group, task_id = fut2key[fut]
|
||||
yield task_group, task_id, fut.result()
|
||||
|
||||
# result producer: sequential or threaded, same consumer
|
||||
def iter_task_results() -> Iterator[tuple[str, int, TaskMetrics]]:
|
||||
"""
|
||||
Yield evaluation results for each (task_group, task_id).
|
||||
|
||||
Depending on `max_parallel_tasks`, runs sequentially or in parallel,
|
||||
but always returns a generator of tuples:
|
||||
(task_group, task_id, TaskMetrics).
|
||||
"""
|
||||
if max_parallel_tasks == 1:
|
||||
yield from _eval_monotask(
|
||||
envs,
|
||||
policy,
|
||||
preprocessor=preprocessor,
|
||||
postprocessor=postprocessor,
|
||||
n_episodes=n_episodes,
|
||||
max_episodes_rendered=max_episodes_rendered,
|
||||
videos_dir=videos_dir,
|
||||
return_episode_data=return_episode_data,
|
||||
start_seed=start_seed,
|
||||
)
|
||||
else:
|
||||
yield from _eval_parallel(
|
||||
envs,
|
||||
policy=policy,
|
||||
preprocessor=preprocessor,
|
||||
postprocessor=postprocessor,
|
||||
n_episodes=n_episodes,
|
||||
max_episodes_rendered=max_episodes_rendered,
|
||||
videos_dir=videos_dir,
|
||||
return_episode_data=return_episode_data,
|
||||
start_seed=start_seed,
|
||||
max_parallel_tasks=max_parallel_tasks,
|
||||
)
|
||||
|
||||
# single accumulator path on the main thread
|
||||
group_acc: dict[str, dict[str, list]] = defaultdict(lambda: {k: [] for k in ACC_KEYS})
|
||||
overall: dict[str, list] = {k: [] for k in ACC_KEYS}
|
||||
per_task_infos: list[dict] = []
|
||||
|
||||
# small inline helper to accumulate one task's metrics into accumulators
|
||||
def _accumulate_to(group: str, metrics: dict):
|
||||
# metrics expected to contain 'sum_rewards', 'max_rewards', 'successes', optionally 'video_paths'
|
||||
# but eval_one may store per-episode lists; we assume metrics uses scalars averaged per task as before.
|
||||
# To be robust, accept scalars or lists.
|
||||
def _append(key, value):
|
||||
if value is None:
|
||||
return
|
||||
if isinstance(value, list):
|
||||
group_acc[group][key].extend(value)
|
||||
overall[key].extend(value)
|
||||
else:
|
||||
group_acc[group][key].append(value)
|
||||
overall[key].append(value)
|
||||
for task_group, _task_id, metrics in iter_task_results():
|
||||
acc = group_acc[task_group]
|
||||
for k in ACC_KEYS:
|
||||
acc[k].extend(metrics[k])
|
||||
overall[k].extend(metrics[k])
|
||||
|
||||
_append("sum_rewards", metrics.get("sum_rewards"))
|
||||
_append("max_rewards", metrics.get("max_rewards"))
|
||||
_append("successes", metrics.get("successes"))
|
||||
# video_paths is list-like
|
||||
paths = metrics.get("video_paths", [])
|
||||
if paths:
|
||||
group_acc[group]["video_paths"].extend(paths)
|
||||
overall["video_paths"].extend(paths)
|
||||
# build outputs
|
||||
results: dict[str, dict] = {}
|
||||
for task_group, data in group_acc.items():
|
||||
suite_rewards = data["sum_rewards"]
|
||||
suite_max = data["max_rewards"]
|
||||
suite_succ = data["successes"]
|
||||
suite_vids = data["video_paths"]
|
||||
|
||||
# Choose runner (sequential vs threaded)
|
||||
task_runner = partial(
|
||||
run_one,
|
||||
policy=policy,
|
||||
preprocessor=preprocessor,
|
||||
postprocessor=postprocessor,
|
||||
n_episodes=n_episodes,
|
||||
max_episodes_rendered=max_episodes_rendered,
|
||||
videos_dir=videos_dir,
|
||||
return_episode_data=return_episode_data,
|
||||
start_seed=start_seed,
|
||||
)
|
||||
suite_eval_s = time.time() - global_start
|
||||
suite_eval_ep_s = suite_eval_s / max(1, len(suite_rewards))
|
||||
|
||||
if max_parallel_tasks <= 1:
|
||||
# sequential path (single accumulator path on the main thread)
|
||||
# NOTE: keeping a single-threaded accumulator avoids concurrent list appends or locks
|
||||
for task_group, task_id, env in tasks:
|
||||
tg, tid, metrics = task_runner(task_group, task_id, env)
|
||||
_accumulate_to(tg, metrics)
|
||||
per_task_infos.append({"task_group": tg, "task_id": tid, "metrics": metrics})
|
||||
else:
|
||||
# threaded path: submit all tasks, consume completions on main thread and accumulate there
|
||||
with cf.ThreadPoolExecutor(max_workers=max_parallel_tasks) as executor:
|
||||
fut2meta = {}
|
||||
for task_group, task_id, env in tasks:
|
||||
fut = executor.submit(task_runner, task_group, task_id, env)
|
||||
fut2meta[fut] = (task_group, task_id)
|
||||
for fut in cf.as_completed(fut2meta):
|
||||
tg, tid, metrics = fut.result()
|
||||
_accumulate_to(tg, metrics)
|
||||
per_task_infos.append({"task_group": tg, "task_id": tid, "metrics": metrics})
|
||||
|
||||
# compute aggregated metrics helper (robust to lists/scalars)
|
||||
def _agg_from_list(xs):
|
||||
if not xs:
|
||||
return float("nan")
|
||||
arr = np.array(xs, dtype=float)
|
||||
return float(np.nanmean(arr))
|
||||
|
||||
# compute per-group aggregates
|
||||
groups_aggregated = {}
|
||||
for group, acc in group_acc.items():
|
||||
groups_aggregated[group] = {
|
||||
"avg_sum_reward": _agg_from_list(acc["sum_rewards"]),
|
||||
"avg_max_reward": _agg_from_list(acc["max_rewards"]),
|
||||
"pc_success": _agg_from_list(acc["successes"]) * 100 if acc["successes"] else float("nan"),
|
||||
"n_episodes": len(acc["sum_rewards"]),
|
||||
"video_paths": list(acc["video_paths"]),
|
||||
results[task_group] = {
|
||||
"aggregated": {
|
||||
"avg_sum_reward": float(np.nanmean(suite_rewards)) if suite_rewards else float("nan"),
|
||||
"avg_max_reward": float(np.nanmean(suite_max)) if suite_max else float("nan"),
|
||||
"pc_success": float(np.nanmean(suite_succ) * 100) if suite_succ else float("nan"),
|
||||
"eval_s": suite_eval_s,
|
||||
"eval_ep_s": suite_eval_ep_s,
|
||||
},
|
||||
"video_paths": suite_vids,
|
||||
"episodes": None,
|
||||
}
|
||||
|
||||
# overall aggregates
|
||||
overall_agg = {
|
||||
"avg_sum_reward": _agg_from_list(overall["sum_rewards"]),
|
||||
"avg_max_reward": _agg_from_list(overall["max_rewards"]),
|
||||
"pc_success": _agg_from_list(overall["successes"]) * 100 if overall["successes"] else float("nan"),
|
||||
"n_episodes": len(overall["sum_rewards"]),
|
||||
"eval_s": time.time() - start_t,
|
||||
"eval_ep_s": (time.time() - start_t) / max(1, len(overall["sum_rewards"])),
|
||||
"video_paths": list(overall["video_paths"]),
|
||||
}
|
||||
|
||||
return {
|
||||
"per_task": per_task_infos,
|
||||
"per_group": groups_aggregated,
|
||||
"overall": overall_agg,
|
||||
global_eval_s = time.time() - global_start
|
||||
global_eval_ep_s = global_eval_s / max(1, len(overall["sum_rewards"]))
|
||||
results["overall"] = {
|
||||
"aggregated": {
|
||||
"avg_sum_reward": float(np.nanmean(overall["sum_rewards"]))
|
||||
if overall["sum_rewards"]
|
||||
else float("nan"),
|
||||
"avg_max_reward": float(np.nanmean(overall["max_rewards"]))
|
||||
if overall["max_rewards"]
|
||||
else float("nan"),
|
||||
"pc_success": float(np.nanmean(overall["successes"]) * 100)
|
||||
if overall["successes"]
|
||||
else float("nan"),
|
||||
"eval_s": global_eval_s,
|
||||
"eval_ep_s": global_eval_ep_s,
|
||||
},
|
||||
"video_paths": overall["video_paths"],
|
||||
"episodes": None,
|
||||
}
|
||||
return results
|
||||
|
||||
|
||||
def main():
|
||||
|
||||
@@ -1,96 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2024 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.
|
||||
|
||||
"""
|
||||
Use this script to get a quick summary of your system config.
|
||||
It should be able to run without any of LeRobot's dependencies or LeRobot itself installed.
|
||||
|
||||
Example:
|
||||
|
||||
```shell
|
||||
lerobot-info
|
||||
```
|
||||
"""
|
||||
|
||||
import importlib
|
||||
import platform
|
||||
|
||||
|
||||
def get_package_version(package_name: str) -> str:
|
||||
"""Get the version of a package if it exists, otherwise return 'N/A'."""
|
||||
try:
|
||||
module = importlib.import_module(package_name)
|
||||
return getattr(module, "__version__", "Installed (version not found)")
|
||||
except ImportError:
|
||||
return "N/A"
|
||||
|
||||
|
||||
def get_sys_info() -> dict:
|
||||
"""Run this to get basic system info to help for tracking issues & bugs."""
|
||||
# General package versions
|
||||
info = {
|
||||
"lerobot version": get_package_version("lerobot"),
|
||||
"Platform": platform.platform(),
|
||||
"Python version": platform.python_version(),
|
||||
"Huggingface Hub version": get_package_version("huggingface_hub"),
|
||||
"Datasets version": get_package_version("datasets"),
|
||||
"Numpy version": get_package_version("numpy"),
|
||||
}
|
||||
|
||||
# PyTorch and GPU specific information
|
||||
torch_version = "N/A"
|
||||
torch_cuda_available = "N/A"
|
||||
cuda_version = "N/A"
|
||||
gpu_model = "N/A"
|
||||
try:
|
||||
import torch
|
||||
|
||||
torch_version = torch.__version__
|
||||
torch_cuda_available = torch.cuda.is_available()
|
||||
if torch_cuda_available:
|
||||
cuda_version = torch.version.cuda
|
||||
# Gets the name of the first available GPU
|
||||
gpu_model = torch.cuda.get_device_name(0)
|
||||
except ImportError:
|
||||
# If torch is not installed, the default "N/A" values will be used.
|
||||
pass
|
||||
|
||||
info.update(
|
||||
{
|
||||
"PyTorch version": torch_version,
|
||||
"Is PyTorch built with CUDA support?": torch_cuda_available,
|
||||
"Cuda version": cuda_version,
|
||||
"GPU model": gpu_model,
|
||||
"Using GPU in script?": "<fill in>",
|
||||
}
|
||||
)
|
||||
|
||||
return info
|
||||
|
||||
|
||||
def format_dict_for_markdown(d: dict) -> str:
|
||||
"""Formats a dictionary into a markdown-friendly bulleted list."""
|
||||
return "\n".join([f"- {prop}: {val}" for prop, val in d.items()])
|
||||
|
||||
|
||||
def main():
|
||||
system_info = get_sys_info()
|
||||
print("\nCopy-and-paste the text below in your GitHub issue and FILL OUT the last point.\n")
|
||||
print(format_dict_for_markdown(system_info))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -24,7 +24,7 @@ Examples of usage:
|
||||
|
||||
- Start an actor server for real robot training with human-in-the-loop intervention:
|
||||
```bash
|
||||
python -m lerobot.rl.actor --config_path src/lerobot/configs/train_config_hilserl_so100.json
|
||||
python -m lerobot.scripts.rl.actor --config_path src/lerobot/configs/train_config_hilserl_so100.json
|
||||
```
|
||||
|
||||
**NOTE**: The actor server requires a running learner server to connect to. Ensure the learner
|
||||
@@ -64,6 +64,12 @@ from lerobot.policies.factory import make_policy
|
||||
from lerobot.policies.sac.modeling_sac import SACPolicy
|
||||
from lerobot.processor import TransitionKey
|
||||
from lerobot.robots import so100_follower # noqa: F401
|
||||
from lerobot.scripts.rl.gym_manipulator import (
|
||||
create_transition,
|
||||
make_processors,
|
||||
make_robot_env,
|
||||
step_env_and_process_transition,
|
||||
)
|
||||
from lerobot.teleoperators import gamepad, so101_leader # noqa: F401
|
||||
from lerobot.teleoperators.utils import TeleopEvents
|
||||
from lerobot.transport import services_pb2, services_pb2_grpc
|
||||
@@ -90,13 +96,6 @@ from lerobot.utils.utils import (
|
||||
init_logging,
|
||||
)
|
||||
|
||||
from .gym_manipulator import (
|
||||
create_transition,
|
||||
make_processors,
|
||||
make_robot_env,
|
||||
step_env_and_process_transition,
|
||||
)
|
||||
|
||||
ACTOR_SHUTDOWN_TIMEOUT = 30
|
||||
|
||||
# Main entry point
|
||||
@@ -25,13 +25,12 @@ from lerobot.robots import ( # noqa: F401
|
||||
make_robot_from_config,
|
||||
so100_follower,
|
||||
)
|
||||
from lerobot.scripts.rl.gym_manipulator import make_robot_env
|
||||
from lerobot.teleoperators import (
|
||||
gamepad, # noqa: F401
|
||||
so101_leader, # noqa: F401
|
||||
)
|
||||
|
||||
from .gym_manipulator import make_robot_env
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
|
||||
|
||||
@@ -44,14 +44,12 @@ from lerobot.processor import (
|
||||
MotorCurrentProcessorStep,
|
||||
Numpy2TorchActionProcessorStep,
|
||||
RewardClassifierProcessorStep,
|
||||
RobotActionToPolicyActionProcessorStep,
|
||||
TimeLimitProcessorStep,
|
||||
Torch2NumpyActionProcessorStep,
|
||||
TransitionKey,
|
||||
VanillaObservationProcessorStep,
|
||||
create_transition,
|
||||
)
|
||||
from lerobot.processor.converters import identity_transition
|
||||
from lerobot.robots import ( # noqa: F401
|
||||
RobotConfig,
|
||||
make_robot_from_config,
|
||||
@@ -59,11 +57,12 @@ from lerobot.robots import ( # noqa: F401
|
||||
)
|
||||
from lerobot.robots.robot import Robot
|
||||
from lerobot.robots.so100_follower.robot_kinematic_processor import (
|
||||
AddRobotObservationAsComplimentaryData,
|
||||
EEBoundsAndSafety,
|
||||
EEReferenceAndDelta,
|
||||
ForwardKinematicsJointsToEEObservation,
|
||||
ForwardKinematicsJointsToEE,
|
||||
GripperVelocityToJoint,
|
||||
InverseKinematicsRLStep,
|
||||
InverseKinematicsEEToJoints,
|
||||
)
|
||||
from lerobot.teleoperators import (
|
||||
gamepad, # noqa: F401
|
||||
@@ -157,20 +156,15 @@ class RobotEnv(gym.Env):
|
||||
|
||||
self.use_gripper = use_gripper
|
||||
|
||||
self._joint_names = list(self.robot.bus.motors.keys())
|
||||
self._raw_joint_positions = None
|
||||
|
||||
self._setup_spaces()
|
||||
|
||||
def _get_observation(self) -> dict[str, Any]:
|
||||
"""Get current robot observation including joint positions and camera images."""
|
||||
obs_dict = self.robot.get_observation()
|
||||
raw_joint_joint_position = {f"{name}.pos": obs_dict[f"{name}.pos"] for name in self._joint_names}
|
||||
joint_positions = np.array([raw_joint_joint_position[f"{name}.pos"] for name in self._joint_names])
|
||||
joint_positions = np.array([obs_dict[name] for name in self._joint_names])
|
||||
|
||||
images = {key: obs_dict[key] for key in self._image_keys}
|
||||
|
||||
return {"agent_pos": joint_positions, "pixels": images, **raw_joint_joint_position}
|
||||
return {"agent_pos": joint_positions, "pixels": images}
|
||||
|
||||
def _setup_spaces(self) -> None:
|
||||
"""Configure observation and action spaces based on robot capabilities."""
|
||||
@@ -245,19 +239,21 @@ class RobotEnv(gym.Env):
|
||||
self.current_step = 0
|
||||
self.episode_data = None
|
||||
obs = self._get_observation()
|
||||
self._raw_joint_positions = {f"{key}.pos": obs[f"{key}.pos"] for key in self._joint_names}
|
||||
return obs, {TeleopEvents.IS_INTERVENTION: False}
|
||||
return obs, {
|
||||
TeleopEvents.IS_INTERVENTION: False,
|
||||
"raw_joint_positions": obs["agent_pos"],
|
||||
}
|
||||
|
||||
def step(self, action) -> tuple[dict[str, np.ndarray], float, bool, bool, dict[str, Any]]:
|
||||
"""Execute one environment step with given action."""
|
||||
joint_targets_dict = {f"{key}.pos": action[i] for i, key in enumerate(self.robot.bus.motors.keys())}
|
||||
joint_targets_dict = {
|
||||
f"{key}.pos": action[f"action.{key}.pos"] for key in self.robot.bus.motors.keys()
|
||||
}
|
||||
|
||||
self.robot.send_action(joint_targets_dict)
|
||||
|
||||
obs = self._get_observation()
|
||||
|
||||
self._raw_joint_positions = {f"{key}.pos": obs[f"{key}.pos"] for key in self._joint_names}
|
||||
|
||||
if self.display_cameras:
|
||||
self.render()
|
||||
|
||||
@@ -272,7 +268,7 @@ class RobotEnv(gym.Env):
|
||||
reward,
|
||||
terminated,
|
||||
truncated,
|
||||
{TeleopEvents.IS_INTERVENTION: False},
|
||||
{TeleopEvents.IS_INTERVENTION: False, "raw_joint_positions": obs["agent_pos"]},
|
||||
)
|
||||
|
||||
def render(self) -> None:
|
||||
@@ -292,10 +288,6 @@ class RobotEnv(gym.Env):
|
||||
if self.robot.is_connected:
|
||||
self.robot.disconnect()
|
||||
|
||||
def get_raw_joint_positions(self) -> dict[str, float]:
|
||||
"""Get raw joint positions."""
|
||||
return self._raw_joint_positions
|
||||
|
||||
|
||||
def make_robot_env(cfg: HILSerlRobotEnvConfig) -> tuple[gym.Env, Any]:
|
||||
"""Create robot environment from configuration.
|
||||
@@ -352,9 +344,7 @@ def make_robot_env(cfg: HILSerlRobotEnvConfig) -> tuple[gym.Env, Any]:
|
||||
|
||||
def make_processors(
|
||||
env: gym.Env, teleop_device: Teleoperator | None, cfg: HILSerlRobotEnvConfig, device: str = "cpu"
|
||||
) -> tuple[
|
||||
DataProcessorPipeline[EnvTransition, EnvTransition], DataProcessorPipeline[EnvTransition, EnvTransition]
|
||||
]:
|
||||
) -> tuple[Any, Any]:
|
||||
"""Create environment and action processors.
|
||||
|
||||
Args:
|
||||
@@ -376,6 +366,7 @@ def make_processors(
|
||||
Torch2NumpyActionProcessorStep(),
|
||||
]
|
||||
|
||||
# Minimal processor pipeline for GymHIL simulation
|
||||
env_pipeline_steps = [
|
||||
Numpy2TorchActionProcessorStep(),
|
||||
VanillaObservationProcessorStep(),
|
||||
@@ -383,10 +374,8 @@ def make_processors(
|
||||
DeviceProcessorStep(device=device),
|
||||
]
|
||||
|
||||
return DataProcessorPipeline(
|
||||
steps=env_pipeline_steps, to_transition=identity_transition, to_output=identity_transition
|
||||
), DataProcessorPipeline(
|
||||
steps=action_pipeline_steps, to_transition=identity_transition, to_output=identity_transition
|
||||
return DataProcessorPipeline(steps=env_pipeline_steps), DataProcessorPipeline(
|
||||
steps=action_pipeline_steps
|
||||
)
|
||||
|
||||
# Full processor pipeline for real robot environment
|
||||
@@ -412,7 +401,7 @@ def make_processors(
|
||||
|
||||
if kinematics_solver is not None:
|
||||
env_pipeline_steps.append(
|
||||
ForwardKinematicsJointsToEEObservation(
|
||||
ForwardKinematicsJointsToEE(
|
||||
kinematics=kinematics_solver,
|
||||
motor_names=motor_names,
|
||||
)
|
||||
@@ -461,6 +450,7 @@ def make_processors(
|
||||
action_pipeline_steps = [
|
||||
AddTeleopActionAsComplimentaryDataStep(teleop_device=teleop_device),
|
||||
AddTeleopEventsAsInfoStep(teleop_device=teleop_device),
|
||||
AddRobotObservationAsComplimentaryData(robot=env.robot),
|
||||
InterventionActionProcessorStep(
|
||||
use_gripper=cfg.processor.gripper.use_gripper if cfg.processor.gripper is not None else False,
|
||||
terminate_on_success=terminate_on_success,
|
||||
@@ -480,37 +470,34 @@ def make_processors(
|
||||
end_effector_step_sizes=cfg.processor.inverse_kinematics.end_effector_step_sizes,
|
||||
motor_names=motor_names,
|
||||
use_latched_reference=False,
|
||||
use_ik_solution=True,
|
||||
),
|
||||
EEBoundsAndSafety(
|
||||
end_effector_bounds=cfg.processor.inverse_kinematics.end_effector_bounds,
|
||||
),
|
||||
InverseKinematicsEEToJoints(
|
||||
kinematics=kinematics_solver,
|
||||
motor_names=motor_names,
|
||||
initial_guess_current_joints=False,
|
||||
),
|
||||
GripperVelocityToJoint(
|
||||
motor_names=motor_names,
|
||||
clip_max=cfg.processor.max_gripper_pos,
|
||||
speed_factor=1.0,
|
||||
discrete_gripper=True,
|
||||
),
|
||||
InverseKinematicsRLStep(
|
||||
kinematics=kinematics_solver, motor_names=motor_names, initial_guess_current_joints=False
|
||||
),
|
||||
]
|
||||
action_pipeline_steps.extend(inverse_kinematics_steps)
|
||||
action_pipeline_steps.append(RobotActionToPolicyActionProcessorStep(motor_names=motor_names))
|
||||
|
||||
return DataProcessorPipeline(
|
||||
steps=env_pipeline_steps, to_transition=identity_transition, to_output=identity_transition
|
||||
), DataProcessorPipeline(
|
||||
steps=action_pipeline_steps, to_transition=identity_transition, to_output=identity_transition
|
||||
)
|
||||
return DataProcessorPipeline(steps=env_pipeline_steps), DataProcessorPipeline(steps=action_pipeline_steps)
|
||||
|
||||
|
||||
def step_env_and_process_transition(
|
||||
env: gym.Env,
|
||||
transition: EnvTransition,
|
||||
action: torch.Tensor,
|
||||
env_processor: DataProcessorPipeline[EnvTransition, EnvTransition],
|
||||
action_processor: DataProcessorPipeline[EnvTransition, EnvTransition],
|
||||
) -> EnvTransition:
|
||||
env_processor: DataProcessorPipeline,
|
||||
action_processor: DataProcessorPipeline,
|
||||
):
|
||||
"""
|
||||
Execute one step with processor pipeline.
|
||||
|
||||
@@ -527,14 +514,14 @@ def step_env_and_process_transition(
|
||||
|
||||
# Create action transition
|
||||
transition[TransitionKey.ACTION] = action
|
||||
transition[TransitionKey.OBSERVATION] = (
|
||||
env.get_raw_joint_positions() if hasattr(env, "get_raw_joint_positions") else {}
|
||||
)
|
||||
processed_action_transition = action_processor(transition)
|
||||
processed_action = processed_action_transition[TransitionKey.ACTION]
|
||||
|
||||
# Step environment with processed action
|
||||
obs, reward, terminated, truncated, info = env.step(processed_action)
|
||||
|
||||
# Combine rewards from environment and action processor
|
||||
|
||||
reward = reward + processed_action_transition[TransitionKey.REWARD]
|
||||
terminated = terminated or processed_action_transition[TransitionKey.DONE]
|
||||
truncated = truncated or processed_action_transition[TransitionKey.TRUNCATED]
|
||||
@@ -558,8 +545,8 @@ def step_env_and_process_transition(
|
||||
|
||||
def control_loop(
|
||||
env: gym.Env,
|
||||
env_processor: DataProcessorPipeline[EnvTransition, EnvTransition],
|
||||
action_processor: DataProcessorPipeline[EnvTransition, EnvTransition],
|
||||
env_processor: DataProcessorPipeline,
|
||||
action_processor: DataProcessorPipeline,
|
||||
teleop_device: Teleoperator,
|
||||
cfg: GymManipulatorConfig,
|
||||
) -> None:
|
||||
@@ -591,7 +578,7 @@ def control_loop(
|
||||
|
||||
# Process initial observation
|
||||
transition = create_transition(observation=obs, info=info, complementary_data=complementary_data)
|
||||
transition = env_processor(data=transition)
|
||||
transition = env_processor(transition)
|
||||
|
||||
# Determine if gripper is used
|
||||
use_gripper = cfg.env.processor.gripper.use_gripper if cfg.env.processor.gripper is not None else True
|
||||
@@ -660,11 +647,7 @@ def control_loop(
|
||||
truncated = transition.get(TransitionKey.TRUNCATED, False)
|
||||
|
||||
if cfg.mode == "record":
|
||||
observations = {
|
||||
k: v.squeeze(0).cpu()
|
||||
for k, v in transition[TransitionKey.OBSERVATION].items()
|
||||
if isinstance(v, torch.Tensor)
|
||||
}
|
||||
observations = {k: v.squeeze(0).cpu() for k, v in transition[TransitionKey.OBSERVATION].items()}
|
||||
# Use teleop_action if available, otherwise use the action from the transition
|
||||
action_to_record = transition[TransitionKey.COMPLEMENTARY_DATA].get(
|
||||
"teleop_action", transition[TransitionKey.ACTION]
|
||||
@@ -678,10 +661,8 @@ def control_loop(
|
||||
if use_gripper:
|
||||
discrete_penalty = transition[TransitionKey.COMPLEMENTARY_DATA].get("discrete_penalty", 0.0)
|
||||
frame["complementary_info.discrete_penalty"] = np.array([discrete_penalty], dtype=np.float32)
|
||||
|
||||
if dataset is not None:
|
||||
frame["task"] = cfg.dataset.task
|
||||
dataset.add_frame(frame)
|
||||
dataset.add_frame(frame, task=cfg.dataset.task)
|
||||
|
||||
episode_step += 1
|
||||
|
||||
@@ -731,19 +712,17 @@ def replay_trajectory(
|
||||
episodes=[cfg.dataset.replay_episode],
|
||||
download_videos=False,
|
||||
)
|
||||
episode_frames = dataset.hf_dataset.filter(lambda x: x["episode_index"] == cfg.dataset.replay_episode)
|
||||
actions = episode_frames.select_columns("action")
|
||||
|
||||
dataset_actions = dataset.hf_dataset.select_columns(["action"])
|
||||
_, info = env.reset()
|
||||
|
||||
for action_data in actions:
|
||||
for action_data in dataset_actions:
|
||||
start_time = time.perf_counter()
|
||||
transition = create_transition(
|
||||
observation=env.get_raw_joint_positions() if hasattr(env, "get_raw_joint_positions") else {},
|
||||
action=action_data["action"],
|
||||
complementary_data={"raw_joint_positions": info["raw_joint_positions"]},
|
||||
)
|
||||
transition = action_processor(transition)
|
||||
env.step(transition[TransitionKey.ACTION])
|
||||
_, _, _, _, info = env.step(transition[TransitionKey.ACTION])
|
||||
busy_wait(1 / cfg.env.fps - (time.perf_counter() - start_time))
|
||||
|
||||
|
||||
@@ -25,7 +25,7 @@ Examples of usage:
|
||||
|
||||
- Start a learner server for training:
|
||||
```bash
|
||||
python -m lerobot.rl.learner --config_path src/lerobot/configs/train_config_hilserl_so100.json
|
||||
python -m lerobot.scripts.rl.learner --config_path src/lerobot/configs/train_config_hilserl_so100.json
|
||||
```
|
||||
|
||||
**NOTE**: Start the learner server before launching the actor server. The learner opens a gRPC server
|
||||
@@ -73,6 +73,7 @@ from lerobot.datasets.lerobot_dataset import LeRobotDataset
|
||||
from lerobot.policies.factory import make_policy
|
||||
from lerobot.policies.sac.modeling_sac import SACPolicy
|
||||
from lerobot.robots import so100_follower # noqa: F401
|
||||
from lerobot.scripts.rl import learner_service
|
||||
from lerobot.teleoperators import gamepad, so101_leader # noqa: F401
|
||||
from lerobot.teleoperators.utils import TeleopEvents
|
||||
from lerobot.transport import services_pb2_grpc
|
||||
@@ -99,8 +100,6 @@ from lerobot.utils.utils import (
|
||||
)
|
||||
from lerobot.utils.wandb_utils import WandBLogger
|
||||
|
||||
from .learner_service import MAX_WORKERS, SHUTDOWN_TIMEOUT, LearnerService
|
||||
|
||||
LOG_PREFIX = "[LEARNER]"
|
||||
|
||||
|
||||
@@ -640,7 +639,7 @@ def start_learner(
|
||||
# TODO: Check if its useful
|
||||
_ = ProcessSignalHandler(False, display_pid=True)
|
||||
|
||||
service = LearnerService(
|
||||
service = learner_service.LearnerService(
|
||||
shutdown_event=shutdown_event,
|
||||
parameters_queue=parameters_queue,
|
||||
seconds_between_pushes=cfg.policy.actor_learner_config.policy_parameters_push_frequency,
|
||||
@@ -650,7 +649,7 @@ def start_learner(
|
||||
)
|
||||
|
||||
server = grpc.server(
|
||||
ThreadPoolExecutor(max_workers=MAX_WORKERS),
|
||||
ThreadPoolExecutor(max_workers=learner_service.MAX_WORKERS),
|
||||
options=[
|
||||
("grpc.max_receive_message_length", MAX_MESSAGE_SIZE),
|
||||
("grpc.max_send_message_length", MAX_MESSAGE_SIZE),
|
||||
@@ -671,7 +670,7 @@ def start_learner(
|
||||
|
||||
shutdown_event.wait()
|
||||
logging.info("[LEARNER] Stopping gRPC server...")
|
||||
server.stop(SHUTDOWN_TIMEOUT)
|
||||
server.stop(learner_service.SHUTDOWN_TIMEOUT)
|
||||
logging.info("[LEARNER] gRPC server stopped")
|
||||
|
||||
|
||||
@@ -67,10 +67,8 @@ def update_policy(
|
||||
) -> tuple[MetricsTracker, dict]:
|
||||
"""
|
||||
Performs a single training step to update the policy's weights.
|
||||
|
||||
This function executes the forward and backward passes, clips gradients, and steps the optimizer and
|
||||
learning rate scheduler. It also handles mixed-precision training via a GradScaler.
|
||||
|
||||
Args:
|
||||
train_metrics: A MetricsTracker instance to record training statistics.
|
||||
policy: The policy model to be trained.
|
||||
@@ -81,7 +79,6 @@ def update_policy(
|
||||
lr_scheduler: An optional learning rate scheduler.
|
||||
use_amp: A boolean indicating whether to use automatic mixed precision.
|
||||
lock: An optional lock for thread-safe optimizer updates.
|
||||
|
||||
Returns:
|
||||
A tuple containing:
|
||||
- The updated MetricsTracker with new statistics for this step.
|
||||
@@ -132,7 +129,6 @@ def update_policy(
|
||||
def train(cfg: TrainPipelineConfig):
|
||||
"""
|
||||
Main function to train a policy.
|
||||
|
||||
This function orchestrates the entire training pipeline, including:
|
||||
- Setting up logging, seeding, and device configuration.
|
||||
- Creating the dataset, evaluation environment (if applicable), policy, and optimizer.
|
||||
@@ -140,7 +136,6 @@ def train(cfg: TrainPipelineConfig):
|
||||
- Running the main training loop, which involves fetching data batches and calling `update_policy`.
|
||||
- Periodically logging metrics, saving model checkpoints, and evaluating the policy.
|
||||
- Pushing the final trained model to the Hugging Face Hub if configured.
|
||||
|
||||
Args:
|
||||
cfg: A `TrainPipelineConfig` object containing all training configurations.
|
||||
"""
|
||||
@@ -163,7 +158,6 @@ def train(cfg: TrainPipelineConfig):
|
||||
|
||||
logging.info("Creating dataset")
|
||||
dataset = make_dataset(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.
|
||||
@@ -173,20 +167,17 @@ def train(cfg: TrainPipelineConfig):
|
||||
eval_env = make_env(cfg.env, n_envs=cfg.eval.batch_size, use_async_envs=cfg.eval.use_async_envs)
|
||||
|
||||
logging.info("Creating policy")
|
||||
|
||||
policy = make_policy(
|
||||
cfg=cfg.policy,
|
||||
ds_meta=dataset.meta,
|
||||
)
|
||||
|
||||
# Create processors - only provide dataset_stats if not resuming from saved processors
|
||||
processor_kwargs = {}
|
||||
if not (cfg.resume and cfg.policy.pretrained_path):
|
||||
# Only provide dataset_stats when not resuming from saved processor state
|
||||
processor_kwargs["dataset_stats"] = dataset.meta.stats
|
||||
|
||||
if cfg.policy.pretrained_path is not None:
|
||||
processor_kwargs["preprocessor_overrides"] = {"device_processor": {"device": device.type}}
|
||||
|
||||
preprocessor, postprocessor = make_pre_post_processors(
|
||||
policy_cfg=cfg.policy, pretrained_path=cfg.policy.pretrained_path, **processor_kwargs
|
||||
)
|
||||
@@ -238,7 +229,6 @@ def train(cfg: TrainPipelineConfig):
|
||||
dl_iter = cycle(dataloader)
|
||||
|
||||
policy.train()
|
||||
|
||||
train_metrics = {
|
||||
"loss": AverageMeter("loss", ":.3f"),
|
||||
"grad_norm": AverageMeter("grdn", ":.3f"),
|
||||
@@ -255,8 +245,8 @@ def train(cfg: TrainPipelineConfig):
|
||||
for _ in range(step, cfg.steps):
|
||||
start_time = time.perf_counter()
|
||||
batch = next(dl_iter)
|
||||
batch = preprocessor(batch)
|
||||
train_tracker.dataloading_s = time.perf_counter() - start_time
|
||||
batch = preprocessor(batch)
|
||||
|
||||
train_tracker, output_dict = update_policy(
|
||||
train_tracker,
|
||||
@@ -304,7 +294,7 @@ def train(cfg: TrainPipelineConfig):
|
||||
torch.autocast(device_type=device.type) if cfg.policy.use_amp else nullcontext(),
|
||||
):
|
||||
eval_info = eval_policy_all(
|
||||
envs=eval_env, # dict[suite][task_id] -> vec_env
|
||||
env=eval_env, # dict[suite][task_id] -> vec_env
|
||||
policy=policy,
|
||||
preprocessor=preprocessor,
|
||||
postprocessor=postprocessor,
|
||||
@@ -313,13 +303,17 @@ def train(cfg: TrainPipelineConfig):
|
||||
max_episodes_rendered=4,
|
||||
start_seed=cfg.seed,
|
||||
max_parallel_tasks=cfg.env.max_parallel_tasks,
|
||||
verbose=False,
|
||||
)
|
||||
|
||||
# overall metrics (suite-agnostic)
|
||||
aggregated = eval_info["overall"]
|
||||
aggregated = eval_info["overall"]["aggregated"]
|
||||
|
||||
# optional: per-suite logging
|
||||
for suite, suite_info in eval_info.items():
|
||||
logging.info("Suite %s aggregated: %s", suite, suite_info)
|
||||
if suite == "overall":
|
||||
continue
|
||||
logging.info("Suite %s aggregated: %s", suite, suite_info["aggregated"])
|
||||
|
||||
# meters/tracker
|
||||
eval_metrics = {
|
||||
@@ -330,13 +324,13 @@ def train(cfg: TrainPipelineConfig):
|
||||
eval_tracker = MetricsTracker(
|
||||
cfg.batch_size, dataset.num_frames, dataset.num_episodes, eval_metrics, initial_step=step
|
||||
)
|
||||
eval_tracker.eval_s = aggregated.pop("eval_s")
|
||||
eval_tracker.avg_sum_reward = aggregated.pop("avg_sum_reward")
|
||||
eval_tracker.pc_success = aggregated.pop("pc_success")
|
||||
eval_tracker.eval_s = aggregated.get("eval_s", 0.0)
|
||||
eval_tracker.avg_sum_reward = aggregated.get("avg_sum_reward", float("nan"))
|
||||
eval_tracker.pc_success = aggregated.get("pc_success", float("nan"))
|
||||
if wandb_logger:
|
||||
wandb_log_dict = {**eval_tracker.to_dict(), **eval_info}
|
||||
wandb_logger.log_dict(wandb_log_dict, step, mode="eval")
|
||||
wandb_logger.log_video(eval_info["overall"]["video_paths"][0], step, mode="eval")
|
||||
wandb_logger.log_video(eval_info["video_paths"][0], step, mode="eval")
|
||||
|
||||
if eval_env:
|
||||
close_envs(eval_env)
|
||||
|
||||
+3
-3
@@ -29,14 +29,14 @@ Examples:
|
||||
|
||||
- Visualize data stored on a local machine:
|
||||
```
|
||||
local$ lerobot-dataset-viz \
|
||||
local$ python -m lerobot.scripts.visualize_dataset \
|
||||
--repo-id lerobot/pusht \
|
||||
--episode-index 0
|
||||
```
|
||||
|
||||
- Visualize data stored on a distant machine with a local viewer:
|
||||
```
|
||||
distant$ lerobot-dataset-viz \
|
||||
distant$ python -m lerobot.scripts.visualize_dataset \
|
||||
--repo-id lerobot/pusht \
|
||||
--episode-index 0 \
|
||||
--save 1 \
|
||||
@@ -50,7 +50,7 @@ local$ rerun lerobot_pusht_episode_0.rrd
|
||||
(You need to forward the websocket port to the distant machine, with
|
||||
`ssh -L 9087:localhost:9087 username@remote-host`)
|
||||
```
|
||||
distant$ lerobot-dataset-viz \
|
||||
distant$ python -m lerobot.scripts.visualize_dataset \
|
||||
--repo-id lerobot/pusht \
|
||||
--episode-index 0 \
|
||||
--mode distant \
|
||||
+5
-9
@@ -20,10 +20,10 @@ Additionally, each individual transform can be visualized separately as well as
|
||||
|
||||
Example:
|
||||
```bash
|
||||
lerobot-imgtransform-viz \
|
||||
--repo_id=lerobot/pusht \
|
||||
--episodes='[0]' \
|
||||
--image_transforms.enable=True
|
||||
python -m lerobot.scripts.visualize_image_transforms \
|
||||
--repo_id=lerobot/pusht \
|
||||
--episodes='[0]' \
|
||||
--image_transforms.enable=True
|
||||
```
|
||||
"""
|
||||
|
||||
@@ -126,9 +126,5 @@ def visualize_image_transforms(cfg: DatasetConfig, output_dir: Path = OUTPUT_DIR
|
||||
save_each_transform(cfg.image_transforms, original_frame, output_dir, n_examples)
|
||||
|
||||
|
||||
def main():
|
||||
visualize_image_transforms()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
visualize_image_transforms()
|
||||
@@ -109,8 +109,8 @@ def teleop_loop(
|
||||
teleop: Teleoperator,
|
||||
robot: Robot,
|
||||
fps: int,
|
||||
teleop_action_processor: RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction],
|
||||
robot_action_processor: RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction],
|
||||
teleop_action_processor: RobotProcessorPipeline[RobotAction, RobotAction],
|
||||
robot_action_processor: RobotProcessorPipeline[RobotAction, RobotAction],
|
||||
robot_observation_processor: RobotProcessorPipeline[RobotObservation, RobotObservation],
|
||||
display_data: bool = False,
|
||||
duration: float | None = None,
|
||||
@@ -137,25 +137,21 @@ def teleop_loop(
|
||||
while True:
|
||||
loop_start = time.perf_counter()
|
||||
|
||||
# Get robot observation
|
||||
# Not really needed for now other than for visualization
|
||||
# teleop_action_processor can take None as an observation
|
||||
# given that it is the identity processor as default
|
||||
obs = robot.get_observation()
|
||||
|
||||
# Get teleop action
|
||||
raw_action = teleop.get_action()
|
||||
|
||||
# Process teleop action through pipeline
|
||||
teleop_action = teleop_action_processor((raw_action, obs))
|
||||
teleop_action = teleop_action_processor(raw_action)
|
||||
|
||||
# Process action for robot through pipeline
|
||||
robot_action_to_send = robot_action_processor((teleop_action, obs))
|
||||
robot_action_to_send = robot_action_processor(teleop_action)
|
||||
|
||||
# Send processed action to robot (robot_action_processor.to_output should return dict[str, Any])
|
||||
_ = robot.send_action(robot_action_to_send)
|
||||
|
||||
if display_data:
|
||||
# Get robot observation
|
||||
obs = robot.get_observation()
|
||||
# Process robot observation through pipeline
|
||||
obs_transition = robot_observation_processor(obs)
|
||||
|
||||
|
||||
@@ -297,8 +297,8 @@ class GamepadController(InputController):
|
||||
try:
|
||||
# Read joystick axes
|
||||
# Left stick X and Y (typically axes 0 and 1)
|
||||
y_input = self.joystick.get_axis(0) # Up/Down (often inverted)
|
||||
x_input = self.joystick.get_axis(1) # Left/Right
|
||||
x_input = self.joystick.get_axis(0) # Left/Right
|
||||
y_input = self.joystick.get_axis(1) # Up/Down (often inverted)
|
||||
|
||||
# Right stick Y (typically axis 3 or 4)
|
||||
z_input = self.joystick.get_axis(3) # Up/Down for Z
|
||||
@@ -310,7 +310,7 @@ class GamepadController(InputController):
|
||||
|
||||
# Calculate deltas (note: may need to invert axes depending on controller)
|
||||
delta_x = -x_input * self.x_step_size # Forward/backward
|
||||
delta_y = -y_input * self.y_step_size # Left/right
|
||||
delta_y = y_input * self.y_step_size # Left/right
|
||||
delta_z = -z_input * self.z_step_size # Up/down
|
||||
|
||||
return delta_x, delta_y, delta_z
|
||||
|
||||
@@ -29,7 +29,7 @@ from lerobot.datasets.transforms import (
|
||||
SharpnessJitter,
|
||||
make_transform_from_config,
|
||||
)
|
||||
from lerobot.scripts.lerobot_imgtransform_viz import (
|
||||
from lerobot.scripts.visualize_image_transforms import (
|
||||
save_all_transforms,
|
||||
save_each_transform,
|
||||
)
|
||||
|
||||
@@ -15,7 +15,7 @@
|
||||
# limitations under the License.
|
||||
import pytest
|
||||
|
||||
from lerobot.scripts.lerobot_dataset_viz import visualize_dataset
|
||||
from lerobot.scripts.visualize_dataset import visualize_dataset
|
||||
|
||||
|
||||
@pytest.mark.skip("TODO: add dummy videos")
|
||||
|
||||
@@ -0,0 +1,147 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2024 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 io
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.fixtures.constants import DUMMY_REPO_ID
|
||||
from tests.utils import require_package
|
||||
|
||||
|
||||
def _find_and_replace(text: str, finds_and_replaces: list[tuple[str, str]]) -> str:
|
||||
for f, r in finds_and_replaces:
|
||||
assert f in text
|
||||
text = text.replace(f, r)
|
||||
return text
|
||||
|
||||
|
||||
# TODO(aliberts): Remove usage of subprocess calls and patch code with fixtures
|
||||
def _run_script(path):
|
||||
subprocess.run([sys.executable, path], check=True)
|
||||
|
||||
|
||||
def _read_file(path):
|
||||
with open(path) as file:
|
||||
return file.read()
|
||||
|
||||
|
||||
@pytest.mark.skip("TODO Fix and remove subprocess / excec calls")
|
||||
def test_example_1(tmp_path, lerobot_dataset_factory):
|
||||
_ = lerobot_dataset_factory(root=tmp_path, repo_id=DUMMY_REPO_ID)
|
||||
path = "examples/1_load_lerobot_dataset.py"
|
||||
file_contents = _read_file(path)
|
||||
file_contents = _find_and_replace(
|
||||
file_contents,
|
||||
[
|
||||
('repo_id = "lerobot/pusht"', f'repo_id = "{DUMMY_REPO_ID}"'),
|
||||
(
|
||||
"LeRobotDataset(repo_id",
|
||||
f"LeRobotDataset(repo_id, root='{str(tmp_path)}'",
|
||||
),
|
||||
],
|
||||
)
|
||||
exec(file_contents, {})
|
||||
assert Path("outputs/examples/1_load_lerobot_dataset/episode_0.mp4").exists()
|
||||
|
||||
|
||||
@pytest.mark.skip("TODO Fix and remove subprocess / excec calls")
|
||||
@require_package("gym_pusht")
|
||||
def test_examples_basic2_basic3_advanced1():
|
||||
"""
|
||||
Train a model with example 3, check the outputs.
|
||||
Evaluate the trained model with example 2, check the outputs.
|
||||
Calculate the validation loss with advanced example 1, check the outputs.
|
||||
"""
|
||||
|
||||
### Test example 3
|
||||
file_contents = _read_file("examples/3_train_policy.py")
|
||||
|
||||
# Do fewer steps, use smaller batch, use CPU, and don't complicate things with dataloader workers.
|
||||
file_contents = _find_and_replace(
|
||||
file_contents,
|
||||
[
|
||||
("training_steps = 5000", "training_steps = 1"),
|
||||
("num_workers=4", "num_workers=0"),
|
||||
('device = torch.device("cuda")', 'device = torch.device("cpu")'),
|
||||
("batch_size=64", "batch_size=1"),
|
||||
],
|
||||
)
|
||||
|
||||
# Pass empty globals to allow dictionary comprehension https://stackoverflow.com/a/32897127/4391249.
|
||||
exec(file_contents, {})
|
||||
|
||||
for file_name in ["model.safetensors", "config.json"]:
|
||||
assert Path(f"outputs/train/example_pusht_diffusion/{file_name}").exists()
|
||||
|
||||
### Test example 2
|
||||
file_contents = _read_file("examples/2_evaluate_pretrained_policy.py")
|
||||
|
||||
# Do fewer evals, use CPU, and use the local model.
|
||||
file_contents = _find_and_replace(
|
||||
file_contents,
|
||||
[
|
||||
(
|
||||
'pretrained_policy_path = Path(snapshot_download("lerobot/diffusion_pusht"))',
|
||||
"",
|
||||
),
|
||||
(
|
||||
'# pretrained_policy_path = Path("outputs/train/example_pusht_diffusion")',
|
||||
'pretrained_policy_path = Path("outputs/train/example_pusht_diffusion")',
|
||||
),
|
||||
('device = torch.device("cuda")', 'device = torch.device("cpu")'),
|
||||
("step += 1", "break"),
|
||||
],
|
||||
)
|
||||
|
||||
exec(file_contents, {})
|
||||
|
||||
assert Path("outputs/eval/example_pusht_diffusion/rollout.mp4").exists()
|
||||
|
||||
## Test example 4
|
||||
file_contents = _read_file("examples/advanced/2_calculate_validation_loss.py")
|
||||
|
||||
# Run on a single example from the last episode, use CPU, and use the local model.
|
||||
file_contents = _find_and_replace(
|
||||
file_contents,
|
||||
[
|
||||
(
|
||||
'pretrained_policy_path = Path(snapshot_download("lerobot/diffusion_pusht"))',
|
||||
"",
|
||||
),
|
||||
(
|
||||
'# pretrained_policy_path = Path("outputs/train/example_pusht_diffusion")',
|
||||
'pretrained_policy_path = Path("outputs/train/example_pusht_diffusion")',
|
||||
),
|
||||
("train_episodes = episodes[:num_train_episodes]", "train_episodes = [0]"),
|
||||
("val_episodes = episodes[num_train_episodes:]", "val_episodes = [1]"),
|
||||
("num_workers=4", "num_workers=0"),
|
||||
('device = torch.device("cuda")', 'device = torch.device("cpu")'),
|
||||
("batch_size=64", "batch_size=1"),
|
||||
],
|
||||
)
|
||||
|
||||
# Capture the output of the script
|
||||
output_buffer = io.StringIO()
|
||||
sys.stdout = output_buffer
|
||||
exec(file_contents, {})
|
||||
printed_output = output_buffer.getvalue()
|
||||
# Restore stdout to its original state
|
||||
sys.stdout = sys.__stdout__
|
||||
assert "Average loss on validation set" in printed_output
|
||||
@@ -194,6 +194,7 @@ def test_policy(ds_repo_id, env_name, env_kwargs, policy_name, policy_kwargs):
|
||||
suite_name = next(iter(envs))
|
||||
task_id = next(iter(envs[suite_name]))
|
||||
env = envs[suite_name][task_id]
|
||||
|
||||
observation, _ = env.reset(seed=train_cfg.seed)
|
||||
|
||||
# apply transform to normalize the observations
|
||||
|
||||
@@ -81,8 +81,8 @@ def test_make_act_processor_basic():
|
||||
|
||||
# Check steps in postprocessor
|
||||
assert len(postprocessor.steps) == 2
|
||||
assert isinstance(postprocessor.steps[0], UnnormalizerProcessorStep)
|
||||
assert isinstance(postprocessor.steps[1], DeviceProcessorStep)
|
||||
assert isinstance(postprocessor.steps[0], DeviceProcessorStep)
|
||||
assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep)
|
||||
|
||||
|
||||
def test_act_processor_normalization():
|
||||
@@ -239,9 +239,7 @@ def test_act_processor_save_and_load():
|
||||
preprocessor.save_pretrained(tmpdir)
|
||||
|
||||
# Load preprocessor
|
||||
loaded_preprocessor = DataProcessorPipeline.from_pretrained(
|
||||
tmpdir, config_filename="policy_preprocessor.json"
|
||||
)
|
||||
loaded_preprocessor = DataProcessorPipeline.from_pretrained(tmpdir)
|
||||
|
||||
# Test that loaded processor works
|
||||
observation = {OBS_STATE: torch.randn(7)}
|
||||
|
||||
@@ -290,10 +290,7 @@ def test_save_and_load_pretrained():
|
||||
|
||||
# Load pipeline
|
||||
loaded_pipeline = DataProcessorPipeline.from_pretrained(
|
||||
tmp_dir,
|
||||
config_filename="batchpipeline.json",
|
||||
to_transition=identity_transition,
|
||||
to_output=identity_transition,
|
||||
tmp_dir, to_transition=identity_transition, to_output=identity_transition
|
||||
)
|
||||
|
||||
assert loaded_pipeline.name == "BatchPipeline"
|
||||
@@ -328,10 +325,7 @@ def test_registry_based_save_load():
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
pipeline.save_pretrained(tmp_dir)
|
||||
loaded_pipeline = DataProcessorPipeline.from_pretrained(
|
||||
tmp_dir,
|
||||
config_filename="dataprocessorpipeline.json",
|
||||
to_transition=identity_transition,
|
||||
to_output=identity_transition,
|
||||
tmp_dir, to_transition=identity_transition, to_output=identity_transition
|
||||
)
|
||||
|
||||
# Verify the loaded processor works
|
||||
|
||||
@@ -250,9 +250,7 @@ def test_classifier_processor_save_and_load():
|
||||
preprocessor.save_pretrained(tmpdir)
|
||||
|
||||
# Load preprocessor
|
||||
loaded_preprocessor = DataProcessorPipeline.from_pretrained(
|
||||
tmpdir, config_filename="classifier_preprocessor.json"
|
||||
)
|
||||
loaded_preprocessor = DataProcessorPipeline.from_pretrained(tmpdir)
|
||||
|
||||
# Test that loaded processor works
|
||||
observation = {
|
||||
|
||||
@@ -324,9 +324,7 @@ def test_save_and_load_pretrained():
|
||||
robot_processor.save_pretrained(tmpdir)
|
||||
|
||||
# Load
|
||||
loaded_processor = DataProcessorPipeline.from_pretrained(
|
||||
tmpdir, config_filename="device_test_processor.json"
|
||||
)
|
||||
loaded_processor = DataProcessorPipeline.from_pretrained(tmpdir)
|
||||
|
||||
assert len(loaded_processor.steps) == 1
|
||||
loaded_device_processor = loaded_processor.steps[0]
|
||||
@@ -1015,147 +1013,3 @@ def test_policy_processor_integration():
|
||||
# Verify action is back on CPU and unnormalized
|
||||
assert output_result[TransitionKey.ACTION].device.type == "cpu"
|
||||
assert output_result[TransitionKey.ACTION].shape == (1, 5)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS not available")
|
||||
def test_mps_float64_compatibility():
|
||||
"""Test MPS device compatibility with float64 tensors (automatic conversion to float32)."""
|
||||
processor = DeviceProcessorStep(device="mps")
|
||||
|
||||
# Create tensors with different dtypes, including float64 which MPS doesn't support
|
||||
observation = {
|
||||
"observation.float64": torch.randn(5, dtype=torch.float64), # Should be converted to float32
|
||||
"observation.float32": torch.randn(5, dtype=torch.float32), # Should remain float32
|
||||
"observation.float16": torch.randn(5, dtype=torch.float16), # Should remain float16
|
||||
"observation.int64": torch.randint(0, 10, (5,), dtype=torch.int64), # Should remain int64
|
||||
"observation.bool": torch.tensor([True, False, True], dtype=torch.bool), # Should remain bool
|
||||
}
|
||||
action = torch.randn(3, dtype=torch.float64) # Should be converted to float32
|
||||
reward = torch.tensor(1.0, dtype=torch.float64) # Should be converted to float32
|
||||
done = torch.tensor(False, dtype=torch.bool) # Should remain bool
|
||||
truncated = torch.tensor(True, dtype=torch.bool) # Should remain bool
|
||||
|
||||
transition = create_transition(
|
||||
observation=observation, action=action, reward=reward, done=done, truncated=truncated
|
||||
)
|
||||
|
||||
result = processor(transition)
|
||||
|
||||
# Check that all tensors are on MPS device
|
||||
assert result[TransitionKey.OBSERVATION]["observation.float64"].device.type == "mps"
|
||||
assert result[TransitionKey.OBSERVATION]["observation.float32"].device.type == "mps"
|
||||
assert result[TransitionKey.OBSERVATION]["observation.float16"].device.type == "mps"
|
||||
assert result[TransitionKey.OBSERVATION]["observation.int64"].device.type == "mps"
|
||||
assert result[TransitionKey.OBSERVATION]["observation.bool"].device.type == "mps"
|
||||
assert result[TransitionKey.ACTION].device.type == "mps"
|
||||
assert result[TransitionKey.REWARD].device.type == "mps"
|
||||
assert result[TransitionKey.DONE].device.type == "mps"
|
||||
assert result[TransitionKey.TRUNCATED].device.type == "mps"
|
||||
|
||||
# Check that float64 tensors were automatically converted to float32
|
||||
assert result[TransitionKey.OBSERVATION]["observation.float64"].dtype == torch.float32
|
||||
assert result[TransitionKey.ACTION].dtype == torch.float32
|
||||
assert result[TransitionKey.REWARD].dtype == torch.float32
|
||||
|
||||
# Check that other dtypes were preserved
|
||||
assert result[TransitionKey.OBSERVATION]["observation.float32"].dtype == torch.float32
|
||||
assert result[TransitionKey.OBSERVATION]["observation.float16"].dtype == torch.float16
|
||||
assert result[TransitionKey.OBSERVATION]["observation.int64"].dtype == torch.int64
|
||||
assert result[TransitionKey.OBSERVATION]["observation.bool"].dtype == torch.bool
|
||||
assert result[TransitionKey.DONE].dtype == torch.bool
|
||||
assert result[TransitionKey.TRUNCATED].dtype == torch.bool
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS not available")
|
||||
def test_mps_float64_with_complementary_data():
|
||||
"""Test MPS float64 conversion with complementary_data tensors."""
|
||||
processor = DeviceProcessorStep(device="mps")
|
||||
|
||||
# Create complementary_data with float64 tensors
|
||||
complementary_data = {
|
||||
"task": ["pick_object"],
|
||||
"index": torch.tensor([42], dtype=torch.int64), # Should remain int64
|
||||
"task_index": torch.tensor([3], dtype=torch.int64), # Should remain int64
|
||||
"float64_tensor": torch.tensor([1.5, 2.5], dtype=torch.float64), # Should convert to float32
|
||||
"float32_tensor": torch.tensor([3.5], dtype=torch.float32), # Should remain float32
|
||||
}
|
||||
|
||||
transition = create_transition(
|
||||
observation={"observation.state": torch.randn(5, dtype=torch.float64)},
|
||||
action=torch.randn(3, dtype=torch.float64),
|
||||
complementary_data=complementary_data,
|
||||
)
|
||||
|
||||
result = processor(transition)
|
||||
|
||||
# Check that all tensors are on MPS device
|
||||
assert result[TransitionKey.OBSERVATION]["observation.state"].device.type == "mps"
|
||||
assert result[TransitionKey.ACTION].device.type == "mps"
|
||||
|
||||
processed_comp_data = result[TransitionKey.COMPLEMENTARY_DATA]
|
||||
assert processed_comp_data["index"].device.type == "mps"
|
||||
assert processed_comp_data["task_index"].device.type == "mps"
|
||||
assert processed_comp_data["float64_tensor"].device.type == "mps"
|
||||
assert processed_comp_data["float32_tensor"].device.type == "mps"
|
||||
|
||||
# Check dtype conversions
|
||||
assert result[TransitionKey.OBSERVATION]["observation.state"].dtype == torch.float32 # Converted
|
||||
assert result[TransitionKey.ACTION].dtype == torch.float32 # Converted
|
||||
assert processed_comp_data["float64_tensor"].dtype == torch.float32 # Converted
|
||||
assert processed_comp_data["float32_tensor"].dtype == torch.float32 # Unchanged
|
||||
assert processed_comp_data["index"].dtype == torch.int64 # Unchanged
|
||||
assert processed_comp_data["task_index"].dtype == torch.int64 # Unchanged
|
||||
|
||||
# Check non-tensor data preserved
|
||||
assert processed_comp_data["task"] == ["pick_object"]
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS not available")
|
||||
def test_mps_with_explicit_float_dtype():
|
||||
"""Test MPS device with explicit float_dtype setting."""
|
||||
# Test that explicit float_dtype still works on MPS
|
||||
processor = DeviceProcessorStep(device="mps", float_dtype="float16")
|
||||
|
||||
observation = {
|
||||
"observation.float64": torch.randn(
|
||||
5, dtype=torch.float64
|
||||
), # First converted to float32, then to float16
|
||||
"observation.float32": torch.randn(5, dtype=torch.float32), # Converted to float16
|
||||
"observation.int32": torch.randint(0, 10, (5,), dtype=torch.int32), # Should remain int32
|
||||
}
|
||||
action = torch.randn(3, dtype=torch.float64)
|
||||
|
||||
transition = create_transition(observation=observation, action=action)
|
||||
result = processor(transition)
|
||||
|
||||
# Check device placement
|
||||
assert result[TransitionKey.OBSERVATION]["observation.float64"].device.type == "mps"
|
||||
assert result[TransitionKey.OBSERVATION]["observation.float32"].device.type == "mps"
|
||||
assert result[TransitionKey.OBSERVATION]["observation.int32"].device.type == "mps"
|
||||
assert result[TransitionKey.ACTION].device.type == "mps"
|
||||
|
||||
# Check that all float tensors end up as float16 (the target dtype)
|
||||
assert result[TransitionKey.OBSERVATION]["observation.float64"].dtype == torch.float16
|
||||
assert result[TransitionKey.OBSERVATION]["observation.float32"].dtype == torch.float16
|
||||
assert result[TransitionKey.ACTION].dtype == torch.float16
|
||||
|
||||
# Check that non-float tensors are preserved
|
||||
assert result[TransitionKey.OBSERVATION]["observation.int32"].dtype == torch.int32
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS not available")
|
||||
def test_mps_serialization():
|
||||
"""Test that MPS device processor can be serialized and loaded correctly."""
|
||||
processor = DeviceProcessorStep(device="mps", float_dtype="float32")
|
||||
|
||||
# Test get_config
|
||||
config = processor.get_config()
|
||||
assert config == {"device": "mps", "float_dtype": "float32"}
|
||||
|
||||
# Test state_dict (should be empty)
|
||||
state = processor.state_dict()
|
||||
assert state == {}
|
||||
|
||||
# Test load_state_dict (should be no-op)
|
||||
processor.load_state_dict({})
|
||||
assert processor.device == "mps"
|
||||
|
||||
@@ -84,8 +84,8 @@ def test_make_diffusion_processor_basic():
|
||||
|
||||
# Check steps in postprocessor
|
||||
assert len(postprocessor.steps) == 2
|
||||
assert isinstance(postprocessor.steps[0], UnnormalizerProcessorStep)
|
||||
assert isinstance(postprocessor.steps[1], DeviceProcessorStep)
|
||||
assert isinstance(postprocessor.steps[0], DeviceProcessorStep)
|
||||
assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep)
|
||||
|
||||
|
||||
def test_diffusion_processor_with_images():
|
||||
@@ -258,9 +258,7 @@ def test_diffusion_processor_save_and_load():
|
||||
preprocessor.save_pretrained(tmpdir)
|
||||
|
||||
# Load preprocessor
|
||||
loaded_preprocessor = DataProcessorPipeline.from_pretrained(
|
||||
tmpdir, config_filename="policy_preprocessor.json"
|
||||
)
|
||||
loaded_preprocessor = DataProcessorPipeline.from_pretrained(tmpdir)
|
||||
|
||||
# Test that loaded processor works
|
||||
observation = {
|
||||
|
||||
@@ -1,341 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""
|
||||
Tests for processor migration detection functionality.
|
||||
"""
|
||||
|
||||
import json
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from lerobot.processor.pipeline import DataProcessorPipeline, ProcessorMigrationError
|
||||
|
||||
|
||||
def test_is_processor_config_valid_configs():
|
||||
"""Test processor config detection with valid configurations."""
|
||||
valid_configs = [
|
||||
{"steps": []}, # Empty steps
|
||||
{"steps": [{"class": "MyClass"}]}, # Class-based step
|
||||
{"steps": [{"registry_name": "my_step"}]}, # Registry-based step
|
||||
{"steps": [{"class": "A"}, {"registry_name": "B"}]}, # Mixed
|
||||
{"name": "Test", "steps": [{"class": "MyClass"}]}, # With name
|
||||
]
|
||||
|
||||
for i, config in enumerate(valid_configs):
|
||||
assert DataProcessorPipeline._is_processor_config(config), (
|
||||
f"Valid config {i} should be detected as processor config: {config}"
|
||||
)
|
||||
|
||||
|
||||
def test_is_processor_config_invalid_configs():
|
||||
"""Test processor config detection with invalid configurations."""
|
||||
invalid_configs = [
|
||||
{}, # No steps field
|
||||
{"steps": "not a list"}, # Steps is not a list
|
||||
{"steps": [{}]}, # Step without class or registry_name
|
||||
{"steps": ["not a dict"]}, # Step is not a dict
|
||||
{"steps": [{"other_field": "value"}]}, # Step with wrong fields
|
||||
{"other_field": "value"}, # Completely different structure
|
||||
]
|
||||
|
||||
for i, config in enumerate(invalid_configs):
|
||||
assert not DataProcessorPipeline._is_processor_config(config), (
|
||||
f"Invalid config {i} should not be detected as processor config: {config}"
|
||||
)
|
||||
|
||||
|
||||
def test_should_suggest_migration_with_processor_config():
|
||||
"""Test that migration is NOT suggested when processor config exists."""
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
# Create a valid processor config
|
||||
processor_config = {
|
||||
"name": "TestProcessor",
|
||||
"steps": [
|
||||
{
|
||||
"class": "lerobot.processor.normalize.NormalizeStep",
|
||||
"config": {"mean": 0.0, "std": 1.0},
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
with open(tmp_path / "processor.json", "w") as f:
|
||||
json.dump(processor_config, f)
|
||||
|
||||
# Should NOT suggest migration (processor config exists)
|
||||
result = DataProcessorPipeline._should_suggest_migration(tmp_path)
|
||||
assert not result
|
||||
|
||||
|
||||
def test_should_suggest_migration_with_empty_processor_config():
|
||||
"""Test that migration is NOT suggested when empty processor config exists."""
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
# Create an empty processor config
|
||||
empty_processor_config = {
|
||||
"name": "EmptyProcessor",
|
||||
"steps": [], # Empty steps is valid
|
||||
}
|
||||
|
||||
with open(tmp_path / "empty_processor.json", "w") as f:
|
||||
json.dump(empty_processor_config, f)
|
||||
|
||||
# Should NOT suggest migration (processor config exists, even if empty)
|
||||
result = DataProcessorPipeline._should_suggest_migration(tmp_path)
|
||||
assert not result
|
||||
|
||||
|
||||
def test_should_suggest_migration_with_model_config_only():
|
||||
"""Test that migration IS suggested when only model config exists."""
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
# Create a model config (like old LeRobot format)
|
||||
model_config = {
|
||||
"type": "act",
|
||||
"input_features": {"observation.state": {"shape": [7]}},
|
||||
"output_features": {"action": {"shape": [7]}},
|
||||
"hidden_dim": 256,
|
||||
"n_obs_steps": 1,
|
||||
"n_action_steps": 1,
|
||||
}
|
||||
|
||||
with open(tmp_path / "config.json", "w") as f:
|
||||
json.dump(model_config, f)
|
||||
|
||||
# SHOULD suggest migration (model config exists but no processor)
|
||||
result = DataProcessorPipeline._should_suggest_migration(tmp_path)
|
||||
assert result
|
||||
|
||||
|
||||
def test_should_suggest_migration_no_json_files():
|
||||
"""Test that migration is NOT suggested when no JSON files exist."""
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
# Create some non-JSON files
|
||||
with open(tmp_path / "model.safetensors", "w") as f:
|
||||
f.write("fake model data")
|
||||
|
||||
with open(tmp_path / "README.md", "w") as f:
|
||||
f.write("# Model README")
|
||||
|
||||
# Should NOT suggest migration (no JSON files)
|
||||
result = DataProcessorPipeline._should_suggest_migration(tmp_path)
|
||||
assert not result
|
||||
|
||||
|
||||
def test_should_suggest_migration_random_json_files():
|
||||
"""Test that migration IS suggested when JSON files exist but none are processor configs."""
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
# Create some random JSON file (not a processor config)
|
||||
random_config = {"some_field": "some_value", "another_field": 123}
|
||||
|
||||
with open(tmp_path / "random.json", "w") as f:
|
||||
json.dump(random_config, f)
|
||||
|
||||
# SHOULD suggest migration (JSON files exist but none are processor configs)
|
||||
result = DataProcessorPipeline._should_suggest_migration(tmp_path)
|
||||
assert result
|
||||
|
||||
|
||||
def test_should_suggest_migration_mixed_configs():
|
||||
"""Test that migration is NOT suggested when processor config exists alongside other configs."""
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
# Create both a processor config and a model config
|
||||
processor_config = {"name": "TestProcessor", "steps": [{"registry_name": "normalize_step"}]}
|
||||
|
||||
model_config = {"type": "diffusion", "hidden_dim": 512}
|
||||
|
||||
with open(tmp_path / "processor.json", "w") as f:
|
||||
json.dump(processor_config, f)
|
||||
|
||||
with open(tmp_path / "config.json", "w") as f:
|
||||
json.dump(model_config, f)
|
||||
|
||||
# Should NOT suggest migration (processor config exists)
|
||||
result = DataProcessorPipeline._should_suggest_migration(tmp_path)
|
||||
assert not result
|
||||
|
||||
|
||||
def test_should_suggest_migration_invalid_json():
|
||||
"""Test that invalid JSON is handled gracefully."""
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
# Create an invalid JSON file
|
||||
with open(tmp_path / "invalid.json", "w") as f:
|
||||
f.write("{ invalid json")
|
||||
|
||||
# Create a valid non-processor config
|
||||
model_config = {"type": "act"}
|
||||
with open(tmp_path / "model.json", "w") as f:
|
||||
json.dump(model_config, f)
|
||||
|
||||
# SHOULD suggest migration (invalid JSON is ignored, but we have non-processor JSON)
|
||||
result = DataProcessorPipeline._should_suggest_migration(tmp_path)
|
||||
assert result
|
||||
|
||||
|
||||
def test_from_pretrained_multiple_json_files_migration_error():
|
||||
"""Test that multiple JSON files trigger ProcessorMigrationError."""
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
# Create multiple non-processor configs
|
||||
model_config = {"type": "act", "hidden_dim": 128}
|
||||
train_config = {"batch_size": 32, "lr": 0.001}
|
||||
|
||||
with open(tmp_path / "config.json", "w") as f:
|
||||
json.dump(model_config, f)
|
||||
|
||||
with open(tmp_path / "train_config.json", "w") as f:
|
||||
json.dump(train_config, f)
|
||||
|
||||
# Should raise ProcessorMigrationError
|
||||
with pytest.raises(ProcessorMigrationError) as exc_info:
|
||||
DataProcessorPipeline.from_pretrained(tmp_path, config_filename="config.json")
|
||||
|
||||
# Check the error details
|
||||
error = exc_info.value
|
||||
assert str(tmp_path) in str(error.model_path)
|
||||
assert "migrate_policy_normalization.py" in error.migration_command
|
||||
assert "not a valid processor configuration" in error.original_error
|
||||
|
||||
|
||||
def test_from_pretrained_no_processor_config_migration_error():
|
||||
"""Test that missing processor config triggers ProcessorMigrationError."""
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
# Create a model config but no processor
|
||||
model_config = {"type": "diffusion", "hidden_dim": 256}
|
||||
|
||||
with open(tmp_path / "config.json", "w") as f:
|
||||
json.dump(model_config, f)
|
||||
|
||||
# Should raise ProcessorMigrationError
|
||||
with pytest.raises(ProcessorMigrationError) as exc_info:
|
||||
DataProcessorPipeline.from_pretrained(tmp_path, config_filename="config.json")
|
||||
|
||||
# Check the error details
|
||||
error = exc_info.value
|
||||
assert str(tmp_path) in str(error.model_path)
|
||||
assert "migrate_policy_normalization.py" in error.migration_command
|
||||
assert "not a valid processor configuration" in error.original_error
|
||||
|
||||
|
||||
def test_from_pretrained_valid_processor_no_migration_error():
|
||||
"""Test that valid processor config does NOT trigger migration error."""
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
# Create a valid processor config
|
||||
processor_config = {
|
||||
"name": "TestProcessor",
|
||||
"steps": [], # Empty is valid
|
||||
}
|
||||
|
||||
with open(tmp_path / "processor.json", "w") as f:
|
||||
json.dump(processor_config, f)
|
||||
|
||||
# Should succeed and create pipeline
|
||||
pipeline = DataProcessorPipeline.from_pretrained(tmp_path, config_filename="processor.json")
|
||||
assert pipeline is not None
|
||||
assert pipeline.name == "TestProcessor"
|
||||
assert len(pipeline) == 0
|
||||
|
||||
|
||||
def test_from_pretrained_no_json_files_no_migration_error():
|
||||
"""Test that directories with no JSON files don't trigger migration errors."""
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
# Create some non-JSON files
|
||||
with open(tmp_path / "model.safetensors", "w") as f:
|
||||
f.write("fake model data")
|
||||
|
||||
# Should raise FileNotFoundError (config file not found)
|
||||
with pytest.raises(FileNotFoundError, match="not found in directory"):
|
||||
DataProcessorPipeline.from_pretrained(tmp_path, config_filename="processor.json")
|
||||
|
||||
|
||||
def test_processor_migration_error_creation():
|
||||
"""Test that ProcessorMigrationError is created correctly."""
|
||||
model_path = "/path/to/model"
|
||||
migration_command = "python migrate.py --path /path/to/model"
|
||||
original_error = "Config not found"
|
||||
|
||||
error = ProcessorMigrationError(model_path, migration_command, original_error)
|
||||
|
||||
assert error.model_path == model_path
|
||||
assert error.migration_command == migration_command
|
||||
assert error.original_error == original_error
|
||||
assert model_path in str(error)
|
||||
assert migration_command in str(error)
|
||||
assert original_error in str(error)
|
||||
|
||||
|
||||
def test_processor_migration_error_attributes():
|
||||
"""Test that ProcessorMigrationError has correct attributes."""
|
||||
model_path = Path("/test/path")
|
||||
migration_command = "python test.py"
|
||||
original_error = "Test error"
|
||||
|
||||
error = ProcessorMigrationError(model_path, migration_command, original_error)
|
||||
|
||||
# Test that attributes are accessible
|
||||
assert hasattr(error, "model_path")
|
||||
assert hasattr(error, "migration_command")
|
||||
assert hasattr(error, "original_error")
|
||||
|
||||
# Test that it's still an Exception
|
||||
assert isinstance(error, Exception)
|
||||
|
||||
|
||||
def test_migration_suggestion_raises_error():
|
||||
"""Test that migration suggestion always raises ProcessorMigrationError."""
|
||||
with pytest.raises(ProcessorMigrationError) as exc_info:
|
||||
DataProcessorPipeline._suggest_processor_migration("/test/path", "Test error")
|
||||
|
||||
error = exc_info.value
|
||||
assert "/test/path" in str(error.model_path)
|
||||
assert "Test error" in error.original_error
|
||||
assert "migrate_policy_normalization.py" in error.migration_command
|
||||
|
||||
|
||||
def test_migration_error_always_raised_for_invalid_configs():
|
||||
"""Test that ProcessorMigrationError is always raised for invalid configs."""
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
# Create a model config
|
||||
model_config = {"type": "test", "param": "value"}
|
||||
with open(tmp_path / "config.json", "w") as f:
|
||||
json.dump(model_config, f)
|
||||
|
||||
# Should always raise ProcessorMigrationError
|
||||
with pytest.raises(ProcessorMigrationError):
|
||||
DataProcessorPipeline.from_pretrained(tmp_path, config_filename="config.json")
|
||||
@@ -1714,9 +1714,7 @@ def test_pipeline_from_pretrained_with_stats_overrides():
|
||||
# Load the pipeline with stat overrides
|
||||
overrides = {"normalizer_processor": {"stats": override_stats}}
|
||||
|
||||
loaded_pipeline = DataProcessorPipeline.from_pretrained(
|
||||
temp_dir, config_filename="test_pipeline.json", overrides=overrides
|
||||
)
|
||||
loaded_pipeline = DataProcessorPipeline.from_pretrained(temp_dir, overrides=overrides)
|
||||
|
||||
# The critical test: the loaded pipeline should use override stats, not original stats
|
||||
loaded_normalizer = loaded_pipeline.steps[0]
|
||||
|
||||
@@ -108,8 +108,8 @@ def test_make_pi0_processor_basic():
|
||||
|
||||
# Check steps in postprocessor
|
||||
assert len(postprocessor.steps) == 2
|
||||
assert isinstance(postprocessor.steps[0], UnnormalizerProcessorStep)
|
||||
assert isinstance(postprocessor.steps[1], DeviceProcessorStep)
|
||||
assert isinstance(postprocessor.steps[0], DeviceProcessorStep)
|
||||
assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep)
|
||||
|
||||
|
||||
def test_pi0_newline_processor_single_task():
|
||||
|
||||
@@ -543,7 +543,7 @@ def test_save_and_load_pretrained():
|
||||
assert config["steps"][1]["config"]["counter"] == 10
|
||||
|
||||
# Load pipeline
|
||||
loaded_pipeline = DataProcessorPipeline.from_pretrained(tmp_dir, config_filename="testpipeline.json")
|
||||
loaded_pipeline = DataProcessorPipeline.from_pretrained(tmp_dir)
|
||||
|
||||
assert loaded_pipeline.name == "TestPipeline"
|
||||
assert len(loaded_pipeline) == 2
|
||||
@@ -571,9 +571,7 @@ def test_step_without_optional_methods():
|
||||
# Save/load should work even without optional methods
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
pipeline.save_pretrained(tmp_dir)
|
||||
loaded_pipeline = DataProcessorPipeline.from_pretrained(
|
||||
tmp_dir, config_filename="dataprocessorpipeline.json"
|
||||
)
|
||||
loaded_pipeline = DataProcessorPipeline.from_pretrained(tmp_dir)
|
||||
assert len(loaded_pipeline) == 1
|
||||
|
||||
|
||||
@@ -602,9 +600,7 @@ def test_mixed_json_and_tensor_state():
|
||||
assert state_path.exists()
|
||||
|
||||
# Load and verify
|
||||
loaded_pipeline = DataProcessorPipeline.from_pretrained(
|
||||
tmp_dir, config_filename="dataprocessorpipeline.json"
|
||||
)
|
||||
loaded_pipeline = DataProcessorPipeline.from_pretrained(tmp_dir)
|
||||
loaded_step = loaded_pipeline.steps[0]
|
||||
|
||||
# Check JSON attributes were restored
|
||||
@@ -892,11 +888,7 @@ def test_from_pretrained_with_overrides():
|
||||
}
|
||||
|
||||
loaded_pipeline = DataProcessorPipeline.from_pretrained(
|
||||
tmp_dir,
|
||||
config_filename="testoverrides.json",
|
||||
overrides=overrides,
|
||||
to_transition=identity_transition,
|
||||
to_output=identity_transition,
|
||||
tmp_dir, overrides=overrides, to_transition=identity_transition, to_output=identity_transition
|
||||
)
|
||||
|
||||
# Verify the pipeline was loaded correctly
|
||||
@@ -934,11 +926,7 @@ def test_from_pretrained_with_partial_overrides():
|
||||
# The current implementation applies overrides to ALL steps with the same class name
|
||||
# Both steps will get the override
|
||||
loaded_pipeline = DataProcessorPipeline.from_pretrained(
|
||||
tmp_dir,
|
||||
config_filename="dataprocessorpipeline.json",
|
||||
overrides=overrides,
|
||||
to_transition=identity_transition,
|
||||
to_output=identity_transition,
|
||||
tmp_dir, overrides=overrides, to_transition=identity_transition, to_output=identity_transition
|
||||
)
|
||||
|
||||
transition = create_transition(reward=1.0)
|
||||
@@ -962,9 +950,7 @@ def test_from_pretrained_invalid_override_key():
|
||||
overrides = {"NonExistentStep": {"param": "value"}}
|
||||
|
||||
with pytest.raises(KeyError, match="Override keys.*do not match any step"):
|
||||
DataProcessorPipeline.from_pretrained(
|
||||
tmp_dir, config_filename="dataprocessorpipeline.json", overrides=overrides
|
||||
)
|
||||
DataProcessorPipeline.from_pretrained(tmp_dir, overrides=overrides)
|
||||
|
||||
|
||||
def test_from_pretrained_multiple_invalid_override_keys():
|
||||
@@ -979,9 +965,7 @@ def test_from_pretrained_multiple_invalid_override_keys():
|
||||
overrides = {"NonExistentStep1": {"param": "value1"}, "NonExistentStep2": {"param": "value2"}}
|
||||
|
||||
with pytest.raises(KeyError) as exc_info:
|
||||
DataProcessorPipeline.from_pretrained(
|
||||
tmp_dir, config_filename="dataprocessorpipeline.json", overrides=overrides
|
||||
)
|
||||
DataProcessorPipeline.from_pretrained(tmp_dir, overrides=overrides)
|
||||
|
||||
error_msg = str(exc_info.value)
|
||||
assert "NonExistentStep1" in error_msg
|
||||
@@ -1001,11 +985,7 @@ def test_from_pretrained_registered_step_override():
|
||||
overrides = {"registered_mock_step": {"value": 999, "device": "cuda"}}
|
||||
|
||||
loaded_pipeline = DataProcessorPipeline.from_pretrained(
|
||||
tmp_dir,
|
||||
config_filename="dataprocessorpipeline.json",
|
||||
overrides=overrides,
|
||||
to_transition=identity_transition,
|
||||
to_output=identity_transition,
|
||||
tmp_dir, overrides=overrides, to_transition=identity_transition, to_output=identity_transition
|
||||
)
|
||||
|
||||
# Test that overrides were applied
|
||||
@@ -1035,11 +1015,7 @@ def test_from_pretrained_mixed_registered_and_unregistered():
|
||||
}
|
||||
|
||||
loaded_pipeline = DataProcessorPipeline.from_pretrained(
|
||||
tmp_dir,
|
||||
config_filename="dataprocessorpipeline.json",
|
||||
overrides=overrides,
|
||||
to_transition=identity_transition,
|
||||
to_output=identity_transition,
|
||||
tmp_dir, overrides=overrides, to_transition=identity_transition, to_output=identity_transition
|
||||
)
|
||||
|
||||
# Test both steps
|
||||
@@ -1062,10 +1038,7 @@ def test_from_pretrained_no_overrides():
|
||||
|
||||
# Load without overrides
|
||||
loaded_pipeline = DataProcessorPipeline.from_pretrained(
|
||||
tmp_dir,
|
||||
config_filename="dataprocessorpipeline.json",
|
||||
to_transition=identity_transition,
|
||||
to_output=identity_transition,
|
||||
tmp_dir, to_transition=identity_transition, to_output=identity_transition
|
||||
)
|
||||
|
||||
assert len(loaded_pipeline) == 1
|
||||
@@ -1087,11 +1060,7 @@ def test_from_pretrained_empty_overrides():
|
||||
|
||||
# Load with empty overrides
|
||||
loaded_pipeline = DataProcessorPipeline.from_pretrained(
|
||||
tmp_dir,
|
||||
config_filename="dataprocessorpipeline.json",
|
||||
overrides={},
|
||||
to_transition=identity_transition,
|
||||
to_output=identity_transition,
|
||||
tmp_dir, overrides={}, to_transition=identity_transition, to_output=identity_transition
|
||||
)
|
||||
|
||||
assert len(loaded_pipeline) == 1
|
||||
@@ -1119,9 +1088,7 @@ def test_from_pretrained_override_instantiation_error():
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="Failed to instantiate processor step"):
|
||||
DataProcessorPipeline.from_pretrained(
|
||||
tmp_dir, config_filename="dataprocessorpipeline.json", overrides=overrides
|
||||
)
|
||||
DataProcessorPipeline.from_pretrained(tmp_dir, overrides=overrides)
|
||||
|
||||
|
||||
def test_from_pretrained_with_state_and_overrides():
|
||||
@@ -1145,9 +1112,7 @@ def test_from_pretrained_with_state_and_overrides():
|
||||
}
|
||||
}
|
||||
|
||||
loaded_pipeline = DataProcessorPipeline.from_pretrained(
|
||||
tmp_dir, config_filename="dataprocessorpipeline.json", overrides=overrides
|
||||
)
|
||||
loaded_pipeline = DataProcessorPipeline.from_pretrained(tmp_dir, overrides=overrides)
|
||||
loaded_step = loaded_pipeline.steps[0]
|
||||
|
||||
# Check that config overrides were applied
|
||||
@@ -1175,9 +1140,7 @@ def test_from_pretrained_override_error_messages():
|
||||
overrides = {"WrongStepName": {"param": "value"}}
|
||||
|
||||
with pytest.raises(KeyError) as exc_info:
|
||||
DataProcessorPipeline.from_pretrained(
|
||||
tmp_dir, config_filename="dataprocessorpipeline.json", overrides=overrides
|
||||
)
|
||||
DataProcessorPipeline.from_pretrained(tmp_dir, overrides=overrides)
|
||||
|
||||
error_msg = str(exc_info.value)
|
||||
assert "WrongStepName" in error_msg
|
||||
@@ -1362,21 +1325,21 @@ def test_multiple_processors_same_directory():
|
||||
assert len(loaded_post) == 1
|
||||
|
||||
|
||||
def test_explicit_config_filename_loading():
|
||||
"""Test explicit config filename loading (no more auto-detection)."""
|
||||
def test_auto_detect_single_config():
|
||||
"""Test automatic config detection when there's only one JSON file."""
|
||||
step = MockStepWithTensorState()
|
||||
pipeline = DataProcessorPipeline([step], name="SingleConfig")
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
pipeline.save_pretrained(tmp_dir)
|
||||
|
||||
# Load with explicit config_filename (now required)
|
||||
loaded = DataProcessorPipeline.from_pretrained(tmp_dir, config_filename="singleconfig.json")
|
||||
# Load without specifying config_filename
|
||||
loaded = DataProcessorPipeline.from_pretrained(tmp_dir)
|
||||
assert loaded.name == "SingleConfig"
|
||||
|
||||
|
||||
def test_explicit_config_selection_with_multiple_configs():
|
||||
"""Test explicit config selection when multiple configs exist."""
|
||||
def test_error_multiple_configs_no_filename():
|
||||
"""Test error when multiple configs exist and no filename specified."""
|
||||
proc1 = DataProcessorPipeline([MockStep()], name="processor1")
|
||||
proc2 = DataProcessorPipeline([MockStep()], name="processor2")
|
||||
|
||||
@@ -1384,12 +1347,9 @@ def test_explicit_config_selection_with_multiple_configs():
|
||||
proc1.save_pretrained(tmp_dir)
|
||||
proc2.save_pretrained(tmp_dir)
|
||||
|
||||
# Can load specific configs explicitly
|
||||
loaded1 = DataProcessorPipeline.from_pretrained(tmp_dir, config_filename="processor1.json")
|
||||
loaded2 = DataProcessorPipeline.from_pretrained(tmp_dir, config_filename="processor2.json")
|
||||
|
||||
assert loaded1.name == "processor1"
|
||||
assert loaded2.name == "processor2"
|
||||
# Should raise error
|
||||
with pytest.raises(ValueError, match="Multiple .json files found"):
|
||||
DataProcessorPipeline.from_pretrained(tmp_dir)
|
||||
|
||||
|
||||
def test_state_file_naming_with_indices():
|
||||
@@ -1523,7 +1483,6 @@ def test_override_with_nested_config():
|
||||
# Load with nested override
|
||||
loaded = DataProcessorPipeline.from_pretrained(
|
||||
tmp_dir,
|
||||
config_filename="dataprocessorpipeline.json",
|
||||
overrides={"complex_config_step": {"nested_config": {"level1": {"level2": "overridden"}}}},
|
||||
to_transition=identity_transition,
|
||||
to_output=identity_transition,
|
||||
@@ -1548,7 +1507,6 @@ def test_override_preserves_defaults():
|
||||
# Override only one parameter
|
||||
loaded = DataProcessorPipeline.from_pretrained(
|
||||
tmp_dir,
|
||||
config_filename="dataprocessorpipeline.json",
|
||||
overrides={
|
||||
"MockStepWithNonSerializableParam": {
|
||||
"multiplier": 5.0 # Only override multiplier
|
||||
@@ -1578,9 +1536,7 @@ def test_override_type_validation():
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="Failed to instantiate"):
|
||||
DataProcessorPipeline.from_pretrained(
|
||||
tmp_dir, config_filename="dataprocessorpipeline.json", overrides=overrides
|
||||
)
|
||||
DataProcessorPipeline.from_pretrained(tmp_dir, overrides=overrides)
|
||||
|
||||
|
||||
def test_override_with_callables():
|
||||
@@ -1631,7 +1587,6 @@ def test_override_with_callables():
|
||||
# Load with callable override
|
||||
loaded = DataProcessorPipeline.from_pretrained(
|
||||
tmp_dir,
|
||||
config_filename="dataprocessorpipeline.json",
|
||||
overrides={"callable_step": {"transform_fn": double_values}},
|
||||
to_transition=identity_transition,
|
||||
to_output=identity_transition,
|
||||
@@ -1656,9 +1611,7 @@ def test_override_multiple_same_class_warning():
|
||||
|
||||
# Override affects all instances of the class
|
||||
loaded = DataProcessorPipeline.from_pretrained(
|
||||
tmp_dir,
|
||||
config_filename="dataprocessorpipeline.json",
|
||||
overrides={"MockStepWithNonSerializableParam": {"multiplier": 10.0}},
|
||||
tmp_dir, overrides={"MockStepWithNonSerializableParam": {"multiplier": 10.0}}
|
||||
)
|
||||
|
||||
# Both steps get the same override
|
||||
@@ -1763,9 +1716,7 @@ def test_override_with_device_strings():
|
||||
# Override device
|
||||
if torch.cuda.is_available():
|
||||
loaded = DataProcessorPipeline.from_pretrained(
|
||||
tmp_dir,
|
||||
config_filename="dataprocessorpipeline.json",
|
||||
overrides={"device_aware_step": {"device": "cuda:0"}},
|
||||
tmp_dir, overrides={"device_aware_step": {"device": "cuda:0"}}
|
||||
)
|
||||
|
||||
loaded_step = loaded.steps[0]
|
||||
@@ -1783,13 +1734,19 @@ def test_from_pretrained_nonexistent_path():
|
||||
|
||||
# Test with an invalid local path - should raise FileNotFoundError
|
||||
with pytest.raises(FileNotFoundError):
|
||||
DataProcessorPipeline.from_pretrained("/path/that/does/not/exist", config_filename="processor.json")
|
||||
DataProcessorPipeline.from_pretrained("/path/that/does/not/exist")
|
||||
|
||||
# Test with a path that doesn't exist as a directory
|
||||
with pytest.raises(FileNotFoundError):
|
||||
DataProcessorPipeline.from_pretrained("user/repo/extra/path", config_filename="processor.json")
|
||||
DataProcessorPipeline.from_pretrained("user/repo/extra/path")
|
||||
|
||||
# Test with a non-existent Hub repo
|
||||
# Test with a Hub repo without specifying config_filename (should raise ValueError)
|
||||
with pytest.raises(
|
||||
ValueError, match="When loading from Hugging Face Hub, 'config_filename' must be specified"
|
||||
):
|
||||
DataProcessorPipeline.from_pretrained("nonexistent-user/nonexistent-repo")
|
||||
|
||||
# Test with a non-existent Hub repo when config_filename is provided
|
||||
with pytest.raises((FileNotFoundError, HfHubHTTPError)):
|
||||
DataProcessorPipeline.from_pretrained(
|
||||
"nonexistent-user/nonexistent-repo", config_filename="processor.json"
|
||||
@@ -1797,9 +1754,9 @@ def test_from_pretrained_nonexistent_path():
|
||||
|
||||
# Test with a local directory that exists but has no config files
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
# Since the directory exists but has no config, it will raise FileNotFoundError
|
||||
# Since the directory exists but has no config, it will try Hub and fail
|
||||
with pytest.raises(FileNotFoundError):
|
||||
DataProcessorPipeline.from_pretrained(tmp_dir, config_filename="processor.json")
|
||||
DataProcessorPipeline.from_pretrained(tmp_dir)
|
||||
|
||||
|
||||
def test_save_load_with_custom_converter_functions():
|
||||
@@ -1836,7 +1793,7 @@ def test_save_load_with_custom_converter_functions():
|
||||
pipeline.save_pretrained(tmp_dir)
|
||||
|
||||
# Load - should use default converters
|
||||
loaded = DataProcessorPipeline.from_pretrained(tmp_dir, config_filename="dataprocessorpipeline.json")
|
||||
loaded = DataProcessorPipeline.from_pretrained(tmp_dir)
|
||||
|
||||
# Verify it uses default converters by checking with standard batch format
|
||||
batch = {
|
||||
|
||||
@@ -1,259 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""
|
||||
Tests for DataProcessorPipeline.from_pretrained helper methods.
|
||||
|
||||
These tests focus on the individual private methods that were extracted from
|
||||
the main from_pretrained method to improve modularity and testability.
|
||||
"""
|
||||
|
||||
import json
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from lerobot.processor.pipeline import DataProcessorPipeline, ProcessorMigrationError
|
||||
|
||||
# Simplified Config Loading Tests
|
||||
|
||||
|
||||
def test_load_config_directory():
|
||||
"""Test loading config from directory."""
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
# Create a config file
|
||||
config_file = tmp_path / "processor.json"
|
||||
test_config = {"name": "TestProcessor", "steps": []}
|
||||
config_file.write_text(json.dumps(test_config))
|
||||
|
||||
# Load from directory
|
||||
loaded_config, base_path = DataProcessorPipeline._load_config(str(tmp_path), "processor.json", {})
|
||||
|
||||
assert loaded_config == test_config
|
||||
assert base_path == tmp_path
|
||||
|
||||
|
||||
def test_load_config_single_file():
|
||||
"""Test loading config from a single file path."""
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
# Create a config file
|
||||
config_file = tmp_path / "processor.json"
|
||||
test_config = {"name": "TestProcessor", "steps": []}
|
||||
config_file.write_text(json.dumps(test_config))
|
||||
|
||||
# Load using file path directly
|
||||
loaded_config, base_path = DataProcessorPipeline._load_config(
|
||||
str(config_file), "any_filename_ignored", {}
|
||||
)
|
||||
|
||||
assert loaded_config == test_config
|
||||
assert base_path == tmp_path
|
||||
|
||||
|
||||
def test_load_config_directory_file_not_found():
|
||||
"""Test directory loading when config file doesn't exist."""
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
# Directory exists but no processor.json
|
||||
with pytest.raises(FileNotFoundError, match="not found in directory"):
|
||||
DataProcessorPipeline._load_config(str(tmp_path), "processor.json", {})
|
||||
|
||||
|
||||
def test_load_config_directory_with_migration_detection():
|
||||
"""Test that missing config triggers migration detection."""
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
# Create old-style config to trigger migration
|
||||
(tmp_path / "config.json").write_text(json.dumps({"type": "act"}))
|
||||
|
||||
# Try to load processor.json (doesn't exist), should trigger migration
|
||||
with pytest.raises(ProcessorMigrationError):
|
||||
DataProcessorPipeline._load_config(str(tmp_path), "processor.json", {})
|
||||
|
||||
|
||||
def test_load_config_nonexistent_path_tries_hub():
|
||||
"""Test that nonexistent paths try Hub (simplified logic)."""
|
||||
# This path doesn't exist locally, should try Hub
|
||||
with pytest.raises(FileNotFoundError, match="on the HuggingFace Hub"):
|
||||
DataProcessorPipeline._load_config("nonexistent/path", "processor.json", {})
|
||||
|
||||
|
||||
# Config Validation Tests
|
||||
|
||||
|
||||
def test_validate_loaded_config_valid_config():
|
||||
"""Test validation with valid processor config."""
|
||||
valid_config = {"name": "TestProcessor", "steps": []}
|
||||
|
||||
# Should not raise any exception
|
||||
DataProcessorPipeline._validate_loaded_config("any-path", valid_config, "processor.json")
|
||||
|
||||
|
||||
def test_validate_loaded_config_invalid_config():
|
||||
"""Test validation with invalid processor config."""
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
# Create non-processor config to trigger migration
|
||||
(tmp_path / "config.json").write_text(json.dumps({"type": "act"}))
|
||||
|
||||
invalid_config = {"type": "act", "hidden_dim": 256}
|
||||
|
||||
with pytest.raises(ProcessorMigrationError):
|
||||
DataProcessorPipeline._validate_loaded_config(str(tmp_path), invalid_config, "config.json")
|
||||
|
||||
|
||||
def test_validate_loaded_config_invalid_config_no_migration():
|
||||
"""Test validation with invalid config when no migration is detected."""
|
||||
# Non-directory path (Hub repo) - no migration detection
|
||||
invalid_config = {"type": "act", "hidden_dim": 256}
|
||||
|
||||
with pytest.raises(ValueError, match="not a valid processor configuration"):
|
||||
DataProcessorPipeline._validate_loaded_config("user/repo", invalid_config, "config.json")
|
||||
|
||||
|
||||
# Step Class Resolution Tests
|
||||
|
||||
|
||||
def test_resolve_step_class_registry_name():
|
||||
"""Test resolution using registry name."""
|
||||
from lerobot.processor.pipeline import ProcessorStep, ProcessorStepRegistry
|
||||
|
||||
# Register a test step
|
||||
@ProcessorStepRegistry.register("test_step")
|
||||
class TestStep(ProcessorStep):
|
||||
def __call__(self, transition):
|
||||
return transition
|
||||
|
||||
def transform_features(self, features):
|
||||
return features
|
||||
|
||||
try:
|
||||
step_entry = {"registry_name": "test_step"}
|
||||
step_class, step_key = DataProcessorPipeline._resolve_step_class(step_entry)
|
||||
|
||||
assert step_class is TestStep
|
||||
assert step_key == "test_step"
|
||||
finally:
|
||||
ProcessorStepRegistry.unregister("test_step")
|
||||
|
||||
|
||||
def test_resolve_step_class_registry_name_not_found():
|
||||
"""Test resolution with non-existent registry name."""
|
||||
step_entry = {"registry_name": "nonexistent_step"}
|
||||
|
||||
with pytest.raises(ImportError, match="Failed to load processor step from registry"):
|
||||
DataProcessorPipeline._resolve_step_class(step_entry)
|
||||
|
||||
|
||||
def test_resolve_step_class_import_path():
|
||||
"""Test resolution using full import path."""
|
||||
# Use a valid existing class (this should work)
|
||||
step_entry = {"class": "lerobot.processor.pipeline.ProcessorStep"}
|
||||
|
||||
# This should succeed - ProcessorStep can be imported, just not instantiated
|
||||
step_class, step_key = DataProcessorPipeline._resolve_step_class(step_entry)
|
||||
|
||||
from lerobot.processor.pipeline import ProcessorStep
|
||||
|
||||
assert step_class is ProcessorStep
|
||||
assert step_key == "ProcessorStep"
|
||||
|
||||
|
||||
def test_resolve_step_class_invalid_import_path():
|
||||
"""Test resolution with invalid import path."""
|
||||
step_entry = {"class": "nonexistent.module.ClassName"}
|
||||
|
||||
with pytest.raises(ImportError, match="Failed to load processor step"):
|
||||
DataProcessorPipeline._resolve_step_class(step_entry)
|
||||
|
||||
|
||||
# Override Validation Tests
|
||||
|
||||
|
||||
def test_validate_overrides_used_all_used():
|
||||
"""Test validation when all overrides are used."""
|
||||
# Empty set means all overrides were used
|
||||
remaining_overrides = set()
|
||||
config = {"steps": [{"class": "SomeStep"}]}
|
||||
|
||||
# Should not raise
|
||||
DataProcessorPipeline._validate_overrides_used(remaining_overrides, config)
|
||||
|
||||
|
||||
def test_validate_overrides_used_some_unused():
|
||||
"""Test validation when some overrides are unused."""
|
||||
remaining_overrides = {"NonExistentStep", "AnotherMissingStep"}
|
||||
config = {
|
||||
"steps": [
|
||||
{"registry_name": "normalize_step"},
|
||||
{"class": "some.module.TransformStep"},
|
||||
]
|
||||
}
|
||||
|
||||
with pytest.raises(KeyError, match="Override keys.*do not match any step"):
|
||||
DataProcessorPipeline._validate_overrides_used(remaining_overrides, config)
|
||||
|
||||
|
||||
def test_validate_overrides_used_helpful_error_message():
|
||||
"""Test that error message includes available step keys."""
|
||||
remaining_overrides = {"WrongStep"}
|
||||
config = {
|
||||
"steps": [
|
||||
{"registry_name": "correct_step"},
|
||||
{"class": "module.path.CorrectClass"},
|
||||
]
|
||||
}
|
||||
|
||||
with pytest.raises(KeyError) as exc_info:
|
||||
DataProcessorPipeline._validate_overrides_used(remaining_overrides, config)
|
||||
|
||||
error_msg = str(exc_info.value)
|
||||
assert "Available step keys" in error_msg
|
||||
assert "correct_step" in error_msg
|
||||
assert "CorrectClass" in error_msg
|
||||
|
||||
|
||||
# Integration Tests for Simplified Logic
|
||||
|
||||
|
||||
def test_simplified_three_way_loading():
|
||||
"""Test that the simplified 3-way loading logic works correctly."""
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
# Test 1: Directory loading
|
||||
config_file = tmp_path / "processor.json"
|
||||
test_config = {"name": "DirectoryTest", "steps": []}
|
||||
config_file.write_text(json.dumps(test_config))
|
||||
|
||||
loaded_config, base_path = DataProcessorPipeline._load_config(str(tmp_path), "processor.json", {})
|
||||
assert loaded_config["name"] == "DirectoryTest"
|
||||
assert base_path == tmp_path
|
||||
|
||||
# Test 2: Single file loading
|
||||
loaded_config, base_path = DataProcessorPipeline._load_config(
|
||||
str(config_file), "ignored_filename", {}
|
||||
)
|
||||
assert loaded_config["name"] == "DirectoryTest"
|
||||
assert base_path == tmp_path
|
||||
@@ -1,525 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from lerobot.configs.types import FeatureType, PipelineFeatureType
|
||||
from lerobot.processor import (
|
||||
DataProcessorPipeline,
|
||||
PolicyActionToRobotActionProcessorStep,
|
||||
ProcessorStepRegistry,
|
||||
RobotActionToPolicyActionProcessorStep,
|
||||
)
|
||||
from lerobot.processor.converters import identity_transition
|
||||
from tests.conftest import assert_contract_is_typed
|
||||
|
||||
|
||||
def test_robot_to_policy_basic_action_conversion():
|
||||
"""Test basic robot action to policy action conversion."""
|
||||
motor_names = ["joint1", "joint2", "joint3"]
|
||||
processor = RobotActionToPolicyActionProcessorStep(motor_names=motor_names)
|
||||
|
||||
robot_action = {
|
||||
"joint1.pos": 1.0,
|
||||
"joint2.pos": 2.0,
|
||||
"joint3.pos": 3.0,
|
||||
}
|
||||
|
||||
policy_action = processor.action(robot_action)
|
||||
|
||||
assert isinstance(policy_action, torch.Tensor)
|
||||
assert policy_action.shape == (3,)
|
||||
torch.testing.assert_close(policy_action, torch.tensor([1.0, 2.0, 3.0]))
|
||||
|
||||
|
||||
def test_robot_to_policy_action_conversion_preserves_order():
|
||||
"""Test that motor names order is preserved in conversion."""
|
||||
motor_names = ["gripper", "arm", "wrist"]
|
||||
processor = RobotActionToPolicyActionProcessorStep(motor_names=motor_names)
|
||||
|
||||
robot_action = {
|
||||
"arm.pos": 10.0,
|
||||
"gripper.pos": 5.0,
|
||||
"wrist.pos": 15.0,
|
||||
}
|
||||
|
||||
policy_action = processor.action(robot_action)
|
||||
|
||||
expected = torch.tensor([5.0, 10.0, 15.0])
|
||||
torch.testing.assert_close(policy_action, expected)
|
||||
|
||||
|
||||
def test_robot_to_policy_action_conversion_with_floats_and_tensors():
|
||||
"""Test conversion with mixed float and tensor values."""
|
||||
motor_names = ["joint1", "joint2"]
|
||||
processor = RobotActionToPolicyActionProcessorStep(motor_names=motor_names)
|
||||
|
||||
robot_action = {
|
||||
"joint1.pos": torch.tensor(1.5),
|
||||
"joint2.pos": 2.5, # Regular float
|
||||
}
|
||||
|
||||
policy_action = processor.action(robot_action)
|
||||
|
||||
assert isinstance(policy_action, torch.Tensor)
|
||||
torch.testing.assert_close(policy_action, torch.tensor([1.5, 2.5]))
|
||||
|
||||
|
||||
def test_robot_to_policy_action_length_mismatch_error():
|
||||
"""Test error when robot action length doesn't match motor names."""
|
||||
motor_names = ["joint1", "joint2", "joint3"]
|
||||
processor = RobotActionToPolicyActionProcessorStep(motor_names=motor_names)
|
||||
|
||||
# Too few actions
|
||||
robot_action = {"joint1.pos": 1.0, "joint2.pos": 2.0}
|
||||
|
||||
with pytest.raises(ValueError, match="Action must have 3 elements, got 2"):
|
||||
processor.action(robot_action)
|
||||
|
||||
robot_action = {
|
||||
"joint1.pos": 1.0,
|
||||
"joint2.pos": 2.0,
|
||||
"joint3.pos": 3.0,
|
||||
"extra.pos": 4.0,
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="Action must have 3 elements, got 4"):
|
||||
processor.action(robot_action)
|
||||
|
||||
|
||||
def test_robot_to_policy_missing_motor_key_error():
|
||||
"""Test error when robot action is missing expected motor keys."""
|
||||
motor_names = ["joint1", "joint2"]
|
||||
processor = RobotActionToPolicyActionProcessorStep(motor_names=motor_names)
|
||||
|
||||
robot_action = {
|
||||
"joint1.pos": 1.0,
|
||||
"wrong_key.pos": 2.0,
|
||||
}
|
||||
|
||||
with pytest.raises(KeyError):
|
||||
processor.action(robot_action)
|
||||
|
||||
|
||||
def test_robot_to_policy_transform_features():
|
||||
"""Test feature transformation for robot to policy action processor."""
|
||||
motor_names = ["joint1", "joint2", "joint3"]
|
||||
processor = RobotActionToPolicyActionProcessorStep(motor_names=motor_names)
|
||||
|
||||
features = {
|
||||
PipelineFeatureType.ACTION: {
|
||||
"joint1.pos": {"type": FeatureType.ACTION, "shape": (1,)},
|
||||
"joint2.pos": {"type": FeatureType.ACTION, "shape": (1,)},
|
||||
"joint3.pos": {"type": FeatureType.ACTION, "shape": (1,)},
|
||||
"other_data": {"type": FeatureType.ENV, "shape": (1,)},
|
||||
}
|
||||
}
|
||||
|
||||
transformed = processor.transform_features(features)
|
||||
|
||||
assert "action" in transformed[PipelineFeatureType.ACTION]
|
||||
action_feature = transformed[PipelineFeatureType.ACTION]["action"]
|
||||
assert action_feature.type == FeatureType.ACTION
|
||||
assert action_feature.shape == (3,)
|
||||
|
||||
assert "joint1.pos" in transformed[PipelineFeatureType.ACTION]
|
||||
assert "joint2.pos" in transformed[PipelineFeatureType.ACTION]
|
||||
assert "joint3.pos" in transformed[PipelineFeatureType.ACTION]
|
||||
|
||||
assert "other_data" in transformed[PipelineFeatureType.ACTION]
|
||||
|
||||
|
||||
def test_robot_to_policy_get_config():
|
||||
"""Test configuration serialization."""
|
||||
motor_names = ["motor1", "motor2"]
|
||||
processor = RobotActionToPolicyActionProcessorStep(motor_names=motor_names)
|
||||
|
||||
config = processor.get_config()
|
||||
assert config == {"motor_names": motor_names}
|
||||
|
||||
|
||||
def test_robot_to_policy_state_dict():
|
||||
"""Test state dict operations."""
|
||||
processor = RobotActionToPolicyActionProcessorStep(motor_names=["joint1"])
|
||||
|
||||
state = processor.state_dict()
|
||||
assert state == {}
|
||||
|
||||
processor.load_state_dict({})
|
||||
|
||||
|
||||
def test_robot_to_policy_single_motor():
|
||||
"""Test with single motor."""
|
||||
processor = RobotActionToPolicyActionProcessorStep(motor_names=["single_joint"])
|
||||
|
||||
robot_action = {"single_joint.pos": 42.0}
|
||||
policy_action = processor.action(robot_action)
|
||||
|
||||
assert policy_action.shape == (1,)
|
||||
torch.testing.assert_close(policy_action, torch.tensor([42.0]))
|
||||
|
||||
|
||||
def test_policy_to_robot_basic_action_conversion():
|
||||
"""Test basic policy action to robot action conversion."""
|
||||
motor_names = ["joint1", "joint2", "joint3"]
|
||||
processor = PolicyActionToRobotActionProcessorStep(motor_names=motor_names)
|
||||
|
||||
policy_action = torch.tensor([1.0, 2.0, 3.0])
|
||||
robot_action = processor.action(policy_action)
|
||||
|
||||
assert isinstance(robot_action, dict)
|
||||
assert len(robot_action) == 3
|
||||
|
||||
expected = {
|
||||
"joint1.pos": 1.0,
|
||||
"joint2.pos": 2.0,
|
||||
"joint3.pos": 3.0,
|
||||
}
|
||||
|
||||
for key, expected_value in expected.items():
|
||||
assert key in robot_action
|
||||
actual_value = robot_action[key]
|
||||
if isinstance(actual_value, torch.Tensor):
|
||||
actual_value = actual_value.item()
|
||||
assert actual_value == pytest.approx(expected_value)
|
||||
|
||||
|
||||
def test_policy_to_robot_action_conversion_preserves_order():
|
||||
"""Test that motor names order corresponds to tensor indices."""
|
||||
motor_names = ["gripper", "arm", "wrist"]
|
||||
processor = PolicyActionToRobotActionProcessorStep(motor_names=motor_names)
|
||||
|
||||
policy_action = torch.tensor([5.0, 10.0, 15.0])
|
||||
robot_action = processor.action(policy_action)
|
||||
|
||||
assert robot_action["gripper.pos"] == pytest.approx(5.0)
|
||||
assert robot_action["arm.pos"] == pytest.approx(10.0)
|
||||
assert robot_action["wrist.pos"] == pytest.approx(15.0)
|
||||
|
||||
|
||||
def test_policy_to_robot_action_conversion_with_numpy_input():
|
||||
"""Test conversion with numpy array input."""
|
||||
import numpy as np
|
||||
|
||||
motor_names = ["joint1", "joint2"]
|
||||
processor = PolicyActionToRobotActionProcessorStep(motor_names=motor_names)
|
||||
|
||||
policy_action = np.array([1.5, 2.5])
|
||||
robot_action = processor.action(policy_action)
|
||||
|
||||
assert robot_action["joint1.pos"] == pytest.approx(1.5)
|
||||
assert robot_action["joint2.pos"] == pytest.approx(2.5)
|
||||
|
||||
|
||||
def test_policy_to_robot_action_length_mismatch_error():
|
||||
"""Test error when policy action length doesn't match motor names."""
|
||||
motor_names = ["joint1", "joint2", "joint3"]
|
||||
processor = PolicyActionToRobotActionProcessorStep(motor_names=motor_names)
|
||||
|
||||
policy_action = torch.tensor([1.0, 2.0])
|
||||
|
||||
with pytest.raises(ValueError, match="Action must have 3 elements, got 2"):
|
||||
processor.action(policy_action)
|
||||
|
||||
policy_action = torch.tensor([1.0, 2.0, 3.0, 4.0])
|
||||
|
||||
with pytest.raises(ValueError, match="Action must have 3 elements, got 4"):
|
||||
processor.action(policy_action)
|
||||
|
||||
|
||||
def test_policy_to_robot_transform_features():
|
||||
"""Test feature transformation for policy to robot action processor."""
|
||||
motor_names = ["joint1", "joint2"]
|
||||
processor = PolicyActionToRobotActionProcessorStep(motor_names=motor_names)
|
||||
|
||||
features = {
|
||||
PipelineFeatureType.ACTION: {
|
||||
"action": {"type": FeatureType.ACTION, "shape": (2,)},
|
||||
"other_data": {"type": FeatureType.ENV, "shape": (1,)},
|
||||
}
|
||||
}
|
||||
|
||||
transformed = processor.transform_features(features)
|
||||
|
||||
assert "joint1.pos" in transformed[PipelineFeatureType.ACTION]
|
||||
assert "joint2.pos" in transformed[PipelineFeatureType.ACTION]
|
||||
|
||||
for motor in motor_names:
|
||||
motor_feature = transformed[PipelineFeatureType.ACTION][f"{motor}.pos"]
|
||||
assert motor_feature.type == FeatureType.ACTION
|
||||
assert motor_feature.shape == (1,)
|
||||
|
||||
assert "action" in transformed[PipelineFeatureType.ACTION]
|
||||
|
||||
assert "other_data" in transformed[PipelineFeatureType.ACTION]
|
||||
|
||||
|
||||
def test_policy_to_robot_get_config():
|
||||
"""Test configuration serialization."""
|
||||
motor_names = ["motor1", "motor2"]
|
||||
processor = PolicyActionToRobotActionProcessorStep(motor_names=motor_names)
|
||||
|
||||
config = processor.get_config()
|
||||
assert config == {"motor_names": motor_names}
|
||||
|
||||
|
||||
def test_policy_to_robot_state_dict():
|
||||
"""Test state dict operations."""
|
||||
processor = PolicyActionToRobotActionProcessorStep(motor_names=["joint1"])
|
||||
|
||||
state = processor.state_dict()
|
||||
assert state == {}
|
||||
|
||||
processor.load_state_dict({})
|
||||
|
||||
|
||||
def test_policy_to_robot_single_motor():
|
||||
"""Test with single motor."""
|
||||
processor = PolicyActionToRobotActionProcessorStep(motor_names=["single_joint"])
|
||||
|
||||
policy_action = torch.tensor([42.0])
|
||||
robot_action = processor.action(policy_action)
|
||||
|
||||
assert len(robot_action) == 1
|
||||
assert robot_action["single_joint.pos"] == pytest.approx(42.0)
|
||||
|
||||
|
||||
def test_robot_to_policy_registry():
|
||||
"""Test RobotActionToPolicyActionProcessorStep registry."""
|
||||
assert "robot_action_to_policy_action_processor" in ProcessorStepRegistry.list()
|
||||
|
||||
retrieved_class = ProcessorStepRegistry.get("robot_action_to_policy_action_processor")
|
||||
assert retrieved_class is RobotActionToPolicyActionProcessorStep
|
||||
|
||||
instance = retrieved_class(motor_names=["test"])
|
||||
assert isinstance(instance, RobotActionToPolicyActionProcessorStep)
|
||||
assert instance.motor_names == ["test"]
|
||||
|
||||
|
||||
def test_policy_to_robot_registry():
|
||||
"""Test PolicyActionToRobotActionProcessorStep registry."""
|
||||
assert "policy_action_to_robot_action_processor" in ProcessorStepRegistry.list()
|
||||
|
||||
retrieved_class = ProcessorStepRegistry.get("policy_action_to_robot_action_processor")
|
||||
assert retrieved_class is PolicyActionToRobotActionProcessorStep
|
||||
|
||||
instance = retrieved_class(motor_names=["test"])
|
||||
assert isinstance(instance, PolicyActionToRobotActionProcessorStep)
|
||||
assert instance.motor_names == ["test"]
|
||||
|
||||
|
||||
def test_save_and_load_robot_to_policy():
|
||||
"""Test saving and loading RobotActionToPolicyActionProcessorStep."""
|
||||
motor_names = ["joint1", "joint2", "joint3"]
|
||||
processor = RobotActionToPolicyActionProcessorStep(motor_names=motor_names)
|
||||
pipeline = DataProcessorPipeline([processor], name="TestRobotToPolicy")
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
# Save pipeline
|
||||
pipeline.save_pretrained(tmp_dir)
|
||||
|
||||
# Check config file exists
|
||||
config_path = Path(tmp_dir) / "testrobottopolicy.json"
|
||||
assert config_path.exists()
|
||||
|
||||
# Load pipeline
|
||||
loaded_pipeline = DataProcessorPipeline.from_pretrained(
|
||||
tmp_dir,
|
||||
"testrobottopolicy.json",
|
||||
to_transition=identity_transition,
|
||||
to_output=identity_transition,
|
||||
)
|
||||
|
||||
assert loaded_pipeline.name == "TestRobotToPolicy"
|
||||
assert len(loaded_pipeline) == 1
|
||||
|
||||
# Check loaded processor
|
||||
loaded_processor = loaded_pipeline.steps[0]
|
||||
assert isinstance(loaded_processor, RobotActionToPolicyActionProcessorStep)
|
||||
assert loaded_processor.motor_names == motor_names
|
||||
|
||||
# Test functionality after loading
|
||||
robot_action = {"joint1.pos": 1.0, "joint2.pos": 2.0, "joint3.pos": 3.0}
|
||||
policy_action = loaded_processor.action(robot_action)
|
||||
torch.testing.assert_close(policy_action, torch.tensor([1.0, 2.0, 3.0]))
|
||||
|
||||
|
||||
def test_save_and_load_policy_to_robot():
|
||||
"""Test saving and loading PolicyActionToRobotActionProcessorStep."""
|
||||
motor_names = ["motor_a", "motor_b"]
|
||||
processor = PolicyActionToRobotActionProcessorStep(motor_names=motor_names)
|
||||
pipeline = DataProcessorPipeline([processor], name="TestPolicyToRobot")
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
# Save pipeline
|
||||
pipeline.save_pretrained(tmp_dir)
|
||||
|
||||
# Load pipeline
|
||||
loaded_pipeline = DataProcessorPipeline.from_pretrained(
|
||||
tmp_dir,
|
||||
"testpolicytorobot.json",
|
||||
to_transition=identity_transition,
|
||||
to_output=identity_transition,
|
||||
)
|
||||
|
||||
loaded_processor = loaded_pipeline.steps[0]
|
||||
assert isinstance(loaded_processor, PolicyActionToRobotActionProcessorStep)
|
||||
assert loaded_processor.motor_names == motor_names
|
||||
|
||||
policy_action = torch.tensor([10.0, 20.0])
|
||||
robot_action = loaded_processor.action(policy_action)
|
||||
assert robot_action["motor_a.pos"] == pytest.approx(10.0)
|
||||
assert robot_action["motor_b.pos"] == pytest.approx(20.0)
|
||||
|
||||
|
||||
# Integration and chaining tests
|
||||
|
||||
|
||||
def test_round_trip_conversion():
|
||||
"""Test that robot->policy->robot conversion preserves values."""
|
||||
motor_names = ["joint1", "joint2", "joint3"]
|
||||
robot_to_policy = RobotActionToPolicyActionProcessorStep(motor_names=motor_names)
|
||||
policy_to_robot = PolicyActionToRobotActionProcessorStep(motor_names=motor_names)
|
||||
|
||||
original_robot_action = {
|
||||
"joint1.pos": 1.5,
|
||||
"joint2.pos": -2.3,
|
||||
"joint3.pos": 0.7,
|
||||
}
|
||||
|
||||
policy_action = robot_to_policy.action(original_robot_action)
|
||||
final_robot_action = policy_to_robot.action(policy_action)
|
||||
|
||||
for key in original_robot_action:
|
||||
original_val = original_robot_action[key]
|
||||
final_val = final_robot_action[key]
|
||||
if isinstance(final_val, torch.Tensor):
|
||||
final_val = final_val.item()
|
||||
assert final_val == pytest.approx(original_val, abs=1e-6)
|
||||
|
||||
|
||||
def test_chained_processors_in_pipeline():
|
||||
"""Test both processors chained in a pipeline."""
|
||||
motor_names = ["joint1", "joint2"]
|
||||
robot_to_policy = RobotActionToPolicyActionProcessorStep(motor_names=motor_names)
|
||||
policy_to_robot = PolicyActionToRobotActionProcessorStep(motor_names=motor_names)
|
||||
|
||||
pipeline = DataProcessorPipeline(
|
||||
[robot_to_policy, policy_to_robot],
|
||||
to_transition=identity_transition,
|
||||
to_output=identity_transition,
|
||||
)
|
||||
|
||||
assert len(pipeline.steps) == 2
|
||||
assert isinstance(pipeline.steps[0], RobotActionToPolicyActionProcessorStep)
|
||||
assert isinstance(pipeline.steps[1], PolicyActionToRobotActionProcessorStep)
|
||||
|
||||
|
||||
def test_robot_to_policy_features_contract(policy_feature_factory):
|
||||
"""Test feature transformation maintains proper typing contract."""
|
||||
processor = RobotActionToPolicyActionProcessorStep(motor_names=["j1", "j2"])
|
||||
features = {
|
||||
PipelineFeatureType.ACTION: {
|
||||
"j1.pos": policy_feature_factory(FeatureType.ACTION, (1,)),
|
||||
"j2.pos": policy_feature_factory(FeatureType.ACTION, (1,)),
|
||||
"other": policy_feature_factory(FeatureType.ENV, (3,)),
|
||||
}
|
||||
}
|
||||
|
||||
out = processor.transform_features(features.copy())
|
||||
|
||||
assert_contract_is_typed(out)
|
||||
|
||||
assert "action" in out[PipelineFeatureType.ACTION]
|
||||
action_feature = out[PipelineFeatureType.ACTION]["action"]
|
||||
assert action_feature.type == FeatureType.ACTION
|
||||
assert action_feature.shape == (2,)
|
||||
|
||||
|
||||
def test_policy_to_robot_features_contract(policy_feature_factory):
|
||||
"""Test feature transformation maintains proper typing contract."""
|
||||
processor = PolicyActionToRobotActionProcessorStep(motor_names=["m1", "m2", "m3"])
|
||||
features = {
|
||||
PipelineFeatureType.ACTION: {
|
||||
"action": policy_feature_factory(FeatureType.ACTION, (3,)),
|
||||
"other": policy_feature_factory(FeatureType.ENV, (1,)),
|
||||
}
|
||||
}
|
||||
|
||||
out = processor.transform_features(features.copy())
|
||||
|
||||
assert_contract_is_typed(out)
|
||||
|
||||
for motor in ["m1", "m2", "m3"]:
|
||||
key = f"{motor}.pos"
|
||||
assert key in out[PipelineFeatureType.ACTION]
|
||||
motor_feature = out[PipelineFeatureType.ACTION][key]
|
||||
assert motor_feature.type == FeatureType.ACTION
|
||||
assert motor_feature.shape == (1,)
|
||||
|
||||
|
||||
def test_empty_motor_names_list():
|
||||
"""Test behavior with empty motor names list."""
|
||||
processor = RobotActionToPolicyActionProcessorStep(motor_names=[])
|
||||
|
||||
robot_action = {}
|
||||
policy_action = processor.action(robot_action)
|
||||
|
||||
assert isinstance(policy_action, torch.Tensor)
|
||||
assert policy_action.shape == (0,)
|
||||
|
||||
|
||||
def test_empty_motor_names_list_policy_to_robot():
|
||||
"""Test PolicyActionToRobotActionProcessorStep with empty motor names."""
|
||||
processor = PolicyActionToRobotActionProcessorStep(motor_names=[])
|
||||
|
||||
policy_action = torch.tensor([])
|
||||
robot_action = processor.action(policy_action)
|
||||
|
||||
assert isinstance(robot_action, dict)
|
||||
assert len(robot_action) == 0
|
||||
|
||||
|
||||
def test_very_long_motor_names():
|
||||
"""Test with many motor names."""
|
||||
motor_names = [f"joint_{i}" for i in range(100)]
|
||||
processor = RobotActionToPolicyActionProcessorStep(motor_names=motor_names)
|
||||
|
||||
robot_action = {f"joint_{i}.pos": float(i) for i in range(100)}
|
||||
policy_action = processor.action(robot_action)
|
||||
|
||||
assert policy_action.shape == (100,)
|
||||
expected = torch.tensor([float(i) for i in range(100)])
|
||||
torch.testing.assert_close(policy_action, expected)
|
||||
|
||||
|
||||
def test_special_characters_in_motor_names():
|
||||
"""Test with special characters in motor names."""
|
||||
motor_names = ["motor-1", "motor_2", "motor.3"]
|
||||
processor = RobotActionToPolicyActionProcessorStep(motor_names=motor_names)
|
||||
|
||||
robot_action = {
|
||||
"motor-1.pos": 1.0,
|
||||
"motor_2.pos": 2.0,
|
||||
"motor.3.pos": 3.0,
|
||||
}
|
||||
|
||||
policy_action = processor.action(robot_action)
|
||||
torch.testing.assert_close(policy_action, torch.tensor([1.0, 2.0, 3.0]))
|
||||
@@ -232,10 +232,7 @@ def test_save_and_load_pretrained():
|
||||
|
||||
# Load pipeline
|
||||
loaded_pipeline = DataProcessorPipeline.from_pretrained(
|
||||
tmp_dir,
|
||||
config_filename="testrenameprocessorstep.json",
|
||||
to_transition=identity_transition,
|
||||
to_output=identity_transition,
|
||||
tmp_dir, to_transition=identity_transition, to_output=identity_transition
|
||||
)
|
||||
|
||||
assert loaded_pipeline.name == "TestRenameProcessorStep"
|
||||
@@ -296,9 +293,7 @@ def test_registry_based_save_load():
|
||||
assert "class" not in config["steps"][0] # Should use registry, not module path
|
||||
|
||||
# Load should work
|
||||
loaded_pipeline = DataProcessorPipeline.from_pretrained(
|
||||
tmp_dir, config_filename="dataprocessorpipeline.json"
|
||||
)
|
||||
loaded_pipeline = DataProcessorPipeline.from_pretrained(tmp_dir)
|
||||
loaded_processor = loaded_pipeline.steps[0]
|
||||
assert isinstance(loaded_processor, RenameObservationsProcessorStep)
|
||||
assert loaded_processor.rename_map == {"key1": "renamed_key1"}
|
||||
|
||||
@@ -84,8 +84,8 @@ def test_make_sac_processor_basic():
|
||||
|
||||
# Check steps in postprocessor
|
||||
assert len(postprocessor.steps) == 2
|
||||
assert isinstance(postprocessor.steps[0], UnnormalizerProcessorStep)
|
||||
assert isinstance(postprocessor.steps[1], DeviceProcessorStep)
|
||||
assert isinstance(postprocessor.steps[0], DeviceProcessorStep)
|
||||
assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep)
|
||||
|
||||
|
||||
def test_sac_processor_normalization_modes():
|
||||
@@ -241,9 +241,7 @@ def test_sac_processor_save_and_load():
|
||||
preprocessor.save_pretrained(tmpdir)
|
||||
|
||||
# Load preprocessor
|
||||
loaded_preprocessor = DataProcessorPipeline.from_pretrained(
|
||||
tmpdir, config_filename="policy_preprocessor.json"
|
||||
)
|
||||
loaded_preprocessor = DataProcessorPipeline.from_pretrained(tmpdir)
|
||||
|
||||
# Test that loaded processor works
|
||||
observation = {OBS_STATE: torch.randn(10)}
|
||||
|
||||
@@ -115,8 +115,8 @@ def test_make_smolvla_processor_basic():
|
||||
|
||||
# Check steps in postprocessor
|
||||
assert len(postprocessor.steps) == 2
|
||||
assert isinstance(postprocessor.steps[0], UnnormalizerProcessorStep)
|
||||
assert isinstance(postprocessor.steps[1], DeviceProcessorStep)
|
||||
assert isinstance(postprocessor.steps[0], DeviceProcessorStep)
|
||||
assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep)
|
||||
|
||||
|
||||
def test_smolvla_newline_processor_single_task():
|
||||
|
||||
@@ -87,8 +87,8 @@ def test_make_tdmpc_processor_basic():
|
||||
|
||||
# Check steps in postprocessor
|
||||
assert len(postprocessor.steps) == 2
|
||||
assert isinstance(postprocessor.steps[0], UnnormalizerProcessorStep)
|
||||
assert isinstance(postprocessor.steps[1], DeviceProcessorStep)
|
||||
assert isinstance(postprocessor.steps[0], DeviceProcessorStep)
|
||||
assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep)
|
||||
|
||||
|
||||
def test_tdmpc_processor_normalization():
|
||||
@@ -269,9 +269,7 @@ def test_tdmpc_processor_save_and_load():
|
||||
preprocessor.save_pretrained(tmpdir)
|
||||
|
||||
# Load preprocessor
|
||||
loaded_preprocessor = DataProcessorPipeline.from_pretrained(
|
||||
tmpdir, config_filename="policy_preprocessor.json"
|
||||
)
|
||||
loaded_preprocessor = DataProcessorPipeline.from_pretrained(tmpdir)
|
||||
|
||||
# Test that loaded processor works
|
||||
observation = {
|
||||
|
||||
@@ -425,10 +425,7 @@ def test_save_and_load_pretrained_with_tokenizer_name(mock_auto_tokenizer):
|
||||
|
||||
# Load processor - tokenizer will be recreated from saved config
|
||||
loaded_processor = DataProcessorPipeline.from_pretrained(
|
||||
temp_dir,
|
||||
config_filename="dataprocessorpipeline.json",
|
||||
to_transition=identity_transition,
|
||||
to_output=identity_transition,
|
||||
temp_dir, to_transition=identity_transition, to_output=identity_transition
|
||||
)
|
||||
|
||||
# Test that loaded processor works
|
||||
@@ -464,7 +461,6 @@ def test_save_and_load_pretrained_with_tokenizer_object():
|
||||
# Load processor with tokenizer override (since tokenizer object wasn't saved)
|
||||
loaded_processor = DataProcessorPipeline.from_pretrained(
|
||||
temp_dir,
|
||||
config_filename="dataprocessorpipeline.json",
|
||||
overrides={"tokenizer_processor": {"tokenizer": mock_tokenizer}},
|
||||
to_transition=identity_transition,
|
||||
to_output=identity_transition,
|
||||
|
||||
@@ -87,8 +87,8 @@ def test_make_vqbet_processor_basic():
|
||||
|
||||
# Check steps in postprocessor
|
||||
assert len(postprocessor.steps) == 2
|
||||
assert isinstance(postprocessor.steps[0], UnnormalizerProcessorStep)
|
||||
assert isinstance(postprocessor.steps[1], DeviceProcessorStep)
|
||||
assert isinstance(postprocessor.steps[0], DeviceProcessorStep)
|
||||
assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep)
|
||||
|
||||
|
||||
def test_vqbet_processor_with_images():
|
||||
@@ -264,9 +264,7 @@ def test_vqbet_processor_save_and_load():
|
||||
preprocessor.save_pretrained(tmpdir)
|
||||
|
||||
# Load preprocessor
|
||||
loaded_preprocessor = DataProcessorPipeline.from_pretrained(
|
||||
tmpdir, config_filename="policy_preprocessor.json"
|
||||
)
|
||||
loaded_preprocessor = DataProcessorPipeline.from_pretrained(tmpdir)
|
||||
|
||||
# Test that loaded processor works
|
||||
observation = {
|
||||
|
||||
@@ -65,7 +65,7 @@ def close_service_stub(channel, server):
|
||||
|
||||
@require_package("grpc")
|
||||
def test_establish_learner_connection_success():
|
||||
from lerobot.rl.actor import establish_learner_connection
|
||||
from lerobot.scripts.rl.actor import establish_learner_connection
|
||||
|
||||
"""Test successful connection establishment."""
|
||||
stub, _servicer, channel, server = create_learner_service_stub()
|
||||
@@ -82,7 +82,7 @@ def test_establish_learner_connection_success():
|
||||
|
||||
@require_package("grpc")
|
||||
def test_establish_learner_connection_failure():
|
||||
from lerobot.rl.actor import establish_learner_connection
|
||||
from lerobot.scripts.rl.actor import establish_learner_connection
|
||||
|
||||
"""Test connection failure."""
|
||||
stub, servicer, channel, server = create_learner_service_stub()
|
||||
@@ -101,7 +101,7 @@ def test_establish_learner_connection_failure():
|
||||
|
||||
@require_package("grpc")
|
||||
def test_push_transitions_to_transport_queue():
|
||||
from lerobot.rl.actor import push_transitions_to_transport_queue
|
||||
from lerobot.scripts.rl.actor import push_transitions_to_transport_queue
|
||||
from lerobot.transport.utils import bytes_to_transitions
|
||||
from tests.transport.test_transport_utils import assert_transitions_equal
|
||||
|
||||
@@ -137,7 +137,7 @@ def test_push_transitions_to_transport_queue():
|
||||
@require_package("grpc")
|
||||
@pytest.mark.timeout(3) # force cross-platform watchdog
|
||||
def test_transitions_stream():
|
||||
from lerobot.rl.actor import transitions_stream
|
||||
from lerobot.scripts.rl.actor import transitions_stream
|
||||
|
||||
"""Test transitions stream functionality."""
|
||||
shutdown_event = Event()
|
||||
@@ -169,7 +169,7 @@ def test_transitions_stream():
|
||||
@require_package("grpc")
|
||||
@pytest.mark.timeout(3) # force cross-platform watchdog
|
||||
def test_interactions_stream():
|
||||
from lerobot.rl.actor import interactions_stream
|
||||
from lerobot.scripts.rl.actor import interactions_stream
|
||||
from lerobot.transport.utils import bytes_to_python_object, python_object_to_bytes
|
||||
|
||||
"""Test interactions stream functionality."""
|
||||
|
||||
@@ -90,13 +90,13 @@ def cfg():
|
||||
@require_package("grpc")
|
||||
@pytest.mark.timeout(10) # force cross-platform watchdog
|
||||
def test_end_to_end_transitions_flow(cfg):
|
||||
from lerobot.rl.actor import (
|
||||
from lerobot.scripts.rl.actor import (
|
||||
establish_learner_connection,
|
||||
learner_service_client,
|
||||
push_transitions_to_transport_queue,
|
||||
send_transitions,
|
||||
)
|
||||
from lerobot.rl.learner import start_learner
|
||||
from lerobot.scripts.rl.learner import start_learner
|
||||
from lerobot.transport.utils import bytes_to_transitions
|
||||
from tests.transport.test_transport_utils import assert_transitions_equal
|
||||
|
||||
@@ -152,12 +152,12 @@ def test_end_to_end_transitions_flow(cfg):
|
||||
@require_package("grpc")
|
||||
@pytest.mark.timeout(10)
|
||||
def test_end_to_end_interactions_flow(cfg):
|
||||
from lerobot.rl.actor import (
|
||||
from lerobot.scripts.rl.actor import (
|
||||
establish_learner_connection,
|
||||
learner_service_client,
|
||||
send_interactions,
|
||||
)
|
||||
from lerobot.rl.learner import start_learner
|
||||
from lerobot.scripts.rl.learner import start_learner
|
||||
from lerobot.transport.utils import bytes_to_python_object, python_object_to_bytes
|
||||
|
||||
"""Test complete interactions flow from actor to learner."""
|
||||
@@ -226,8 +226,8 @@ def test_end_to_end_interactions_flow(cfg):
|
||||
@pytest.mark.parametrize("data_size", ["small", "large"])
|
||||
@pytest.mark.timeout(10)
|
||||
def test_end_to_end_parameters_flow(cfg, data_size):
|
||||
from lerobot.rl.actor import establish_learner_connection, learner_service_client, receive_policy
|
||||
from lerobot.rl.learner import start_learner
|
||||
from lerobot.scripts.rl.actor import establish_learner_connection, learner_service_client, receive_policy
|
||||
from lerobot.scripts.rl.learner import start_learner
|
||||
from lerobot.transport.utils import bytes_to_state_dict, state_to_bytes
|
||||
|
||||
"""Test complete parameter flow from learner to actor, with small and large data."""
|
||||
|
||||
@@ -50,7 +50,7 @@ def create_learner_service_stub(
|
||||
):
|
||||
import grpc
|
||||
|
||||
from lerobot.rl.learner_service import LearnerService
|
||||
from lerobot.scripts.rl.learner_service import LearnerService
|
||||
from lerobot.transport import services_pb2_grpc # generated from .proto
|
||||
|
||||
"""Fixture to start a LearnerService gRPC server and provide a connected stub."""
|
||||
|
||||
Reference in New Issue
Block a user