Compare commits

..

21 Commits

Author SHA1 Message Date
Jade Choghari 0bda187268 xvla log fix 2026-02-26 09:55:25 +00:00
Steven Palma 59b33c0ea3 Merge branch 'chore/bump_transformers_v5' into ci/add_hf_account 2026-02-25 11:56:52 +01:00
Steven Palma 4419901e6b fix(ci): temp fix for pi0 rtc test 2026-02-25 11:42:04 +01:00
Jade Choghari 3f3a159cff fix wall x for transformer v5 (#3008)
* tv5 fix

* various wall x fixes

* Delete tests/policies/pi0_pi05/print_pi05_output_logits.py

Signed-off-by: Jade Choghari <chogharijade@gmail.com>

* sync modeling_florence2.py with chore/bump_transformers_v5

* more

* more fixes

* more

* remove comment

* more

---------

Signed-off-by: Jade Choghari <chogharijade@gmail.com>
2026-02-24 21:16:37 +03:00
Steven Palma 6deebf1e47 Merge branch 'chore/bump_transformers_v5' into ci/add_hf_account 2026-02-24 17:50:48 +01:00
Steven Palma b9cb947bd2 style(test): pre-commit 2026-02-24 15:35:25 +01:00
Steven Palma 481a956100 Merge branch 'chore/bump_transformers_v5' into ci/add_hf_account
Signed-off-by: Steven Palma <imstevenpmwork@ieee.org>
2026-02-24 13:18:07 +01:00
Steven Palma ffde29be49 chore(ci): change hf call + secret name 2026-02-24 12:11:31 +01:00
Steven Palma 2504d00707 feat(ci): log into HF to unblock some CI tests 2026-02-24 11:56:13 +01:00
Steven Palma 11cefed08a style(test): pre-commit check 2026-02-24 10:57:29 +01:00
Steven Palma 7bfedd1388 test(policies): enable wall x CI testing 2026-02-24 10:53:23 +01:00
Jade Choghari 8c95a71c94 chore: fix XVLA in transformers v5 (#3006) 2026-02-24 10:29:48 +01:00
Steven Palma 1d048c7e2b Merge branch 'main' into chore/bump_transformers_v5 2026-02-23 20:54:02 +01:00
Jade Choghari 419305a4c2 Fix: full pi models support for transformer v5 (#2967)
* fix(pi): remove loss truncation

* fix(pi): remove state padding before tokenization

* fix(pi): fix image padding value

* fix from_pretrain

* add transformer v5 changes

* remove reference

* more fixes

* make it work

* add support for rest of pi family

* add pifast work

* more changes

* more changes

* more cleanup

* fix torch params

* dtype fix

* torch compile

* embed mismatch fix

* revert groot

* more nit fixes

* remove unused classes

* more fixes

* revert

* nit

* torch dtype warning fix

* but back dynamic renaming

* add tie embedding

---------

Co-authored-by: Yufei Sun <skieyfly@gmail.com>
2026-02-23 22:44:13 +03:00
Steven Palma 753b996cda test(rl): skip ci tests for resnet10 2026-02-12 21:25:39 +01:00
Steven Palma 099f3ba4d7 fix(policy): xvla forced_bos_token missing 2026-02-12 21:14:53 +01:00
Steven Palma 3f3d08e5a8 chore(style): fix pre-commit 2026-02-12 19:37:57 +01:00
Steven Palma 9e1a67c862 chore(dependencies): bump gr00t to transformers v5 2026-02-12 19:33:56 +01:00
Steven Palma 54c38627bd chore(dependencies): bump wall x to transformers v5 2026-02-12 19:33:34 +01:00
Steven Palma f0ef3717ca chore(dependencies): bump pi0 family to transformers v5 2026-02-12 19:31:23 +01:00
Steven Palma bd8e1ccf70 chore(dependencies): upgrade transformers + hggingface-hub + peft + scipy 2026-02-12 19:17:54 +01:00
133 changed files with 1559 additions and 5293 deletions
+1 -2
View File
@@ -44,7 +44,7 @@ permissions:
# Sets up the environment variables # Sets up the environment variables
env: env:
UV_VERSION: "0.8.0" UV_VERSION: "0.8.0"
PYTHON_VERSION: "3.12" PYTHON_VERSION: "3.10"
# Ensures that only the latest commit for a PR or branch is built, canceling older runs. # Ensures that only the latest commit for a PR or branch is built, canceling older runs.
concurrency: concurrency:
@@ -91,7 +91,6 @@ jobs:
run: uv sync --extra "test" run: uv sync --extra "test"
- name: Login to Hugging Face - name: Login to Hugging Face
if: env.HF_USER_TOKEN != ''
run: | run: |
uv run hf auth login --token "$HF_USER_TOKEN" --add-to-git-credential uv run hf auth login --token "$HF_USER_TOKEN" --add-to-git-credential
uv run hf auth whoami uv run hf auth whoami
+1 -5
View File
@@ -37,7 +37,7 @@ permissions:
# Sets up the environment variables # Sets up the environment variables
env: env:
UV_VERSION: "0.8.0" UV_VERSION: "0.8.0"
PYTHON_VERSION: "3.12" PYTHON_VERSION: "3.10"
DOCKER_IMAGE_NAME: huggingface/lerobot-gpu DOCKER_IMAGE_NAME: huggingface/lerobot-gpu
# Ensures that only the latest action is built, canceling older runs. # Ensures that only the latest action is built, canceling older runs.
@@ -89,7 +89,6 @@ jobs:
run: uv sync --extra all # TODO(Steven): Make flash-attn optional run: uv sync --extra all # TODO(Steven): Make flash-attn optional
- name: Login to Hugging Face - name: Login to Hugging Face
if: env.HF_USER_TOKEN != ''
run: | run: |
uv run hf auth login --token "$HF_USER_TOKEN" --add-to-git-credential uv run hf auth login --token "$HF_USER_TOKEN" --add-to-git-credential
uv run hf auth whoami uv run hf auth whoami
@@ -182,12 +181,9 @@ jobs:
working-directory: /lerobot working-directory: /lerobot
steps: steps:
- name: Login to Hugging Face - name: Login to Hugging Face
if: env.HF_USER_TOKEN != ''
run: | run: |
hf auth login --token "$HF_USER_TOKEN" --add-to-git-credential hf auth login --token "$HF_USER_TOKEN" --add-to-git-credential
hf auth whoami hf auth whoami
- name: Fix ptxas permissions
run: chmod +x /lerobot/.venv/lib/python3.12/site-packages/triton/backends/nvidia/bin/ptxas
- name: Run pytest on GPU - name: Run pytest on GPU
run: pytest tests -vv --maxfail=10 run: pytest tests -vv --maxfail=10
- name: Run end-to-end tests - name: Run end-to-end tests
+4 -20
View File
@@ -28,7 +28,7 @@ on:
# Sets up the environment variables # Sets up the environment variables
env: env:
UV_VERSION: "0.8.0" UV_VERSION: "0.8.0"
PYTHON_VERSION: "3.12" PYTHON_VERSION: "3.10"
DOCKER_IMAGE_NAME_CPU: huggingface/lerobot-cpu:latest DOCKER_IMAGE_NAME_CPU: huggingface/lerobot-cpu:latest
DOCKER_IMAGE_NAME_GPU: huggingface/lerobot-gpu:latest DOCKER_IMAGE_NAME_GPU: huggingface/lerobot-gpu:latest
@@ -119,7 +119,6 @@ jobs:
HF_LEROBOT_HOME: /home/user_lerobot/.cache/huggingface/lerobot HF_LEROBOT_HOME: /home/user_lerobot/.cache/huggingface/lerobot
TORCH_HOME: /home/user_lerobot/.cache/torch TORCH_HOME: /home/user_lerobot/.cache/torch
TRITON_CACHE_DIR: /home/user_lerobot/.cache/triton TRITON_CACHE_DIR: /home/user_lerobot/.cache/triton
HF_USER_TOKEN: ${{ secrets.LEROBOT_HF_USER }}
container: container:
image: ${{ needs.build-docker-cpu-nightly.outputs.image_tag }} # zizmor: ignore[unpinned-images] image: ${{ needs.build-docker-cpu-nightly.outputs.image_tag }} # zizmor: ignore[unpinned-images]
options: --shm-size "16gb" options: --shm-size "16gb"
@@ -131,11 +130,6 @@ jobs:
shell: bash shell: bash
working-directory: /lerobot working-directory: /lerobot
steps: steps:
- name: Login to Hugging Face
if: env.HF_USER_TOKEN != ''
run: |
hf auth login --token "$HF_USER_TOKEN" --add-to-git-credential
hf auth whoami
- name: Run pytest on CPU - name: Run pytest on CPU
run: pytest tests -vv --maxfail=10 run: pytest tests -vv --maxfail=10
- name: Run end-to-end tests - name: Run end-to-end tests
@@ -152,7 +146,6 @@ jobs:
HF_LEROBOT_HOME: /home/user_lerobot/.cache/huggingface/lerobot HF_LEROBOT_HOME: /home/user_lerobot/.cache/huggingface/lerobot
TORCH_HOME: /home/user_lerobot/.cache/torch TORCH_HOME: /home/user_lerobot/.cache/torch
TRITON_CACHE_DIR: /home/user_lerobot/.cache/triton TRITON_CACHE_DIR: /home/user_lerobot/.cache/triton
HF_USER_TOKEN: ${{ secrets.LEROBOT_HF_USER }}
container: container:
image: ${{ needs.build-docker-gpu-nightly.outputs.image_tag }} # zizmor: ignore[unpinned-images] image: ${{ needs.build-docker-gpu-nightly.outputs.image_tag }} # zizmor: ignore[unpinned-images]
options: --gpus all --shm-size "16gb" options: --gpus all --shm-size "16gb"
@@ -164,11 +157,6 @@ jobs:
shell: bash shell: bash
working-directory: /lerobot working-directory: /lerobot
steps: steps:
- name: Login to Hugging Face
if: env.HF_USER_TOKEN != ''
run: |
hf auth login --token "$HF_USER_TOKEN" --add-to-git-credential
hf auth whoami
- name: Run pytest on GPU - name: Run pytest on GPU
run: pytest tests -vv --maxfail=10 run: pytest tests -vv --maxfail=10
- name: Run end-to-end tests - name: Run end-to-end tests
@@ -186,7 +174,6 @@ jobs:
TORCH_HOME: /home/user_lerobot/.cache/torch TORCH_HOME: /home/user_lerobot/.cache/torch
TRITON_CACHE_DIR: /home/user_lerobot/.cache/triton TRITON_CACHE_DIR: /home/user_lerobot/.cache/triton
CUDA_VISIBLE_DEVICES: "0,1,2,3" CUDA_VISIBLE_DEVICES: "0,1,2,3"
HF_USER_TOKEN: ${{ secrets.LEROBOT_HF_USER }}
container: container:
image: ${{ needs.build-docker-gpu-nightly.outputs.image_tag }} # zizmor: ignore[unpinned-images] image: ${{ needs.build-docker-gpu-nightly.outputs.image_tag }} # zizmor: ignore[unpinned-images]
options: --gpus all --shm-size "16gb" options: --gpus all --shm-size "16gb"
@@ -198,15 +185,12 @@ jobs:
shell: bash shell: bash
working-directory: /lerobot working-directory: /lerobot
steps: steps:
- name: Login to Hugging Face
if: env.HF_USER_TOKEN != ''
run: |
hf auth login --token "$HF_USER_TOKEN" --add-to-git-credential
hf auth whoami
- name: Verify GPU availability - name: Verify GPU availability
run: | run: |
nvidia-smi nvidia-smi
python -c "import torch; print(f'PyTorch CUDA available: {torch.cuda.is_available()}'); print(f'Number of GPUs: {torch.cuda.device_count()}')" python -c "import torch; print(f'PyTorch CUDA available: {torch.cuda.is_available()}'); print(f'Number of GPUs: {torch.cuda.device_count()}')"
- name: Run multi-GPU training tests - name: Run multi-GPU training tests
run: pytest -vv tests/training/ # TODO(Steven): Investigate why motors tests are failing in multi-GPU setup
run: pytest tests -vv --maxfail=10 --ignore=tests/motors/
timeout-minutes: 10
+1 -1
View File
@@ -50,7 +50,7 @@ jobs:
- name: Set up Python - name: Set up Python
uses: actions/setup-python@v6 uses: actions/setup-python@v6
with: with:
python-version: '3.12' python-version: '3.10'
- name: Run pre-commit hooks - name: Run pre-commit hooks
uses: pre-commit/action@v3.0.1 # zizmor: ignore[unpinned-uses] uses: pre-commit/action@v3.0.1 # zizmor: ignore[unpinned-uses]
+10 -2
View File
@@ -22,7 +22,7 @@ on:
# Sets up the environment variables # Sets up the environment variables
env: env:
UV_VERSION: "0.8.0" UV_VERSION: "0.8.0"
PYTHON_VERSION: "3.12" PYTHON_VERSION: "3.10"
jobs: jobs:
# This job builds the Python package and publishes it to PyPI # This job builds the Python package and publishes it to PyPI
@@ -45,7 +45,7 @@ jobs:
- name: Set up Python - name: Set up Python
uses: actions/setup-python@v6 uses: actions/setup-python@v6
with: with:
python-version: '3.12' python-version: '3.10'
- name: Extract Version - name: Extract Version
id: extract_info id: extract_info
@@ -83,6 +83,14 @@ jobs:
exit 1 exit 1
fi fi
- name: Remove Tags with Git dependencies
# TODO(Steven): Temporary patch to remove pi from PyPi 0.4.0 release due to its reliance on git dependencies.
run: |
echo "::info:: Checking for Git dependencies to remove from pyproject.toml..."
grep -E '@ git\+https|lerobot\[pi\]' pyproject.toml | sed 's/^/::warning:: Removing line: /' || true
sed -E -i '/@ git\+https|lerobot\[pi\]/d' pyproject.toml
echo "::info:: Git dependencies removed. Proceeding with build."
- name: Install build dependencies - name: Install build dependencies
run: python -m pip install build run: python -m pip install build
+2 -13
View File
@@ -29,7 +29,7 @@ permissions:
# Sets up the environment variables # Sets up the environment variables
env: env:
UV_VERSION: "0.8.0" UV_VERSION: "0.8.0"
PYTHON_VERSION: "3.12" PYTHON_VERSION: "3.10"
DOCKER_IMAGE_NAME: huggingface/lerobot-gpu:unbound DOCKER_IMAGE_NAME: huggingface/lerobot-gpu:unbound
# Ensures that only the latest action is built, canceling older runs. # Ensures that only the latest action is built, canceling older runs.
@@ -48,7 +48,6 @@ jobs:
MUJOCO_GL: egl MUJOCO_GL: egl
HF_HOME: /mnt/cache/.cache/huggingface HF_HOME: /mnt/cache/.cache/huggingface
HF_LEROBOT_HOME: /mnt/cache/.cache/huggingface/lerobot HF_LEROBOT_HOME: /mnt/cache/.cache/huggingface/lerobot
HF_USER_TOKEN: ${{ secrets.LEROBOT_HF_USER }}
steps: steps:
- uses: actions/checkout@v6 - uses: actions/checkout@v6
with: with:
@@ -80,11 +79,7 @@ jobs:
- name: Install lerobot with all extras - name: Install lerobot with all extras
run: uv sync --extra all # TODO(Steven): Make flash-attn optional run: uv sync --extra all # TODO(Steven): Make flash-attn optional
- name: Login to Hugging Face
if: env.HF_USER_TOKEN != ''
run: |
uv run hf auth login --token "$HF_USER_TOKEN" --add-to-git-credential
uv run hf auth whoami
- name: Run pytest (all extras) - name: Run pytest (all extras)
run: uv run pytest tests -vv run: uv run pytest tests -vv
@@ -142,7 +137,6 @@ jobs:
HF_LEROBOT_HOME: /home/user_lerobot/.cache/huggingface/lerobot HF_LEROBOT_HOME: /home/user_lerobot/.cache/huggingface/lerobot
TORCH_HOME: /home/user_lerobot/.cache/torch TORCH_HOME: /home/user_lerobot/.cache/torch
TRITON_CACHE_DIR: /home/user_lerobot/.cache/triton TRITON_CACHE_DIR: /home/user_lerobot/.cache/triton
HF_USER_TOKEN: ${{ secrets.LEROBOT_HF_USER }}
container: container:
image: ${{ needs.build-and-push-docker.outputs.image_tag }} # zizmor: ignore[unpinned-images] image: ${{ needs.build-and-push-docker.outputs.image_tag }} # zizmor: ignore[unpinned-images]
options: --gpus all --shm-size "16gb" options: --gpus all --shm-size "16gb"
@@ -154,11 +148,6 @@ jobs:
shell: bash shell: bash
working-directory: /lerobot working-directory: /lerobot
steps: steps:
- name: Login to Hugging Face
if: env.HF_USER_TOKEN != ''
run: |
hf auth login --token "$HF_USER_TOKEN" --add-to-git-credential
hf auth whoami
- name: Run pytest on GPU - name: Run pytest on GPU
run: pytest tests -vv run: pytest tests -vv
- name: Run end-to-end tests - name: Run end-to-end tests
+2 -2
View File
@@ -13,7 +13,7 @@
# limitations under the License. # limitations under the License.
default_language_version: default_language_version:
python: python3.12 python: python3.10
exclude: "tests/artifacts/.*\\.safetensors$" exclude: "tests/artifacts/.*\\.safetensors$"
@@ -55,7 +55,7 @@ repos:
rev: v3.21.0 rev: v3.21.0
hooks: hooks:
- id: pyupgrade - id: pyupgrade
args: [--py312-plus] args: [--py310-plus]
##### Markdown Quality ##### ##### Markdown Quality #####
- repo: https://github.com/rbubley/mirrors-prettier - repo: https://github.com/rbubley/mirrors-prettier
-25
View File
@@ -1,25 +0,0 @@
# AI Usage Policy
The LeRobot project welcomes contributions from everyone, and we have a few guidelines regarding AI usage to ensure high code quality, clear communication, and a healthy open-source ecosystem:
- **Please disclose significant AI assistance.** If you used AI tools (e.g., Copilot, Claude, Cursor, ChatGPT) to generate a substantial portion of your code or text, let us know in your PR description. Transparency helps us review your changes more effectively.
- **Own your code (The Human-in-the-Loop).** You must fully understand all the changes you are proposing. If you cannot explain what your AI-assisted code does or how it interacts with LeRobot's broader architecture, please take the time to learn and test it before submitting.
- **Keep issues and discussions focused.** You are welcome to use AI to help draft issues or PR descriptions, but please review and edit them carefully before posting. AI can often be overly verbose; trimming the noise and getting straight to the point helps our maintainers address your needs faster.
Our core maintainers also use AI tools to aid their workflows, but they do so while bringing deep contextual knowledge of the LeRobot codebase to validate the output. We ask all contributors to apply that same level of rigor.
## Remember the Human Maintainers
Please remember that LeRobot is maintained by a dedicated team of humans.
Every discussion, issue, and pull request is read and reviewed by real people. While AI tools can generate thousands of lines of code in seconds, reviewing that code still takes human time and energy. Submitting unverified or low-effort AI output puts an unfair burden on our maintainers.
Today, the quality of the AI output still heavily depends on the developer driving the tool. We ask that you respect our maintainers' time by thoroughly vetting, testing, and refining your submissions.
## AI is Welcome Here
LeRobot operates at the cutting edge of AI and robotics, and many of our maintainers actively embrace AI coding assistants as valuable productivity tools. We are a pro-AI project!
Our reason for having an AI policy is not an anti-AI stance. Rather, it exists to ensure that AI is used to enhance human contributions, not replace them with unverified noise. It's about how the tools are used, not the tools themselves.
We value the unique human insight you bring to the LeRobot community. Let AI empower your workflow, but always let your own judgment take the wheel.
+1 -1
View File
@@ -2,7 +2,7 @@
Everyone is welcome to contribute, and we value everybody's contribution. Code is not the only way to help the community. Answering questions, helping others, reaching out, and improving the documentation are immensely valuable. Everyone is welcome to contribute, and we value everybody's contribution. Code is not the only way to help the community. Answering questions, helping others, reaching out, and improving the documentation are immensely valuable.
Whichever way you choose to contribute, please be mindful to respect our [code of conduct](./CODE_OF_CONDUCT.md) and our [AI policy](./AI_POLICY.md). Whichever way you choose to contribute, please be mindful to respect our [code of conduct](./CODE_OF_CONDUCT.md).
## Ways to Contribute ## Ways to Contribute
-1
View File
@@ -1,3 +1,2 @@
include src/lerobot/templates/lerobot_modelcard_template.md include src/lerobot/templates/lerobot_modelcard_template.md
include src/lerobot/datasets/card_template.md include src/lerobot/datasets/card_template.md
include src/lerobot/envs/metaworld_config.json
+1 -18
View File
@@ -135,7 +135,7 @@ Learn how to implement your own simulation environment or benchmark and distribu
## Citation ## Citation
If you use LeRobot in your project, please cite the GitHub repository to acknowledge the ongoing development and contributors: If you use LeRobot in your research, please cite:
```bibtex ```bibtex
@misc{cadene2024lerobot, @misc{cadene2024lerobot,
@@ -146,23 +146,6 @@ If you use LeRobot in your project, please cite the GitHub repository to acknowl
} }
``` ```
If you are referencing our research or the academic paper, please also cite our ICLR publication:
<details>
<summary><b>ICLR 2026 Paper</b></summary>
```bibtex
@inproceedings{cadenelerobot,
title={LeRobot: An Open-Source Library for End-to-End Robot Learning},
author={Cadene, Remi and Alibert, Simon and Capuano, Francesco and Aractingi, Michel and Zouitine, Adil and Kooijmans, Pepijn and Choghari, Jade and Russi, Martino and Pascal, Caroline and Palma, Steven and Shukor, Mustafa and Moss, Jess and Soare, Alexander and Aubakirova, Dana and Lhoest, Quentin and Gallou\'edec, Quentin and Wolf, Thomas},
booktitle={The Fourteenth International Conference on Learning Representations},
year={2026},
url={https://arxiv.org/abs/2602.22818}
}
```
</details>
## Contribute ## Contribute
We welcome contributions from everyone in the community! To get started, please read our [CONTRIBUTING.md](./CONTRIBUTING.md) guide. Whether you're adding a new feature, improving documentation, or fixing a bug, your help and feedback are invaluable. We're incredibly excited about the future of open-source robotics and can't wait to work with you on what's next—thank you for your support! We welcome contributions from everyone in the community! To get started, please read our [CONTRIBUTING.md](./CONTRIBUTING.md) guide. Whether you're adding a new feature, improving documentation, or fixing a bug, your help and feedback are invaluable. We're incredibly excited about the future of open-source robotics and can't wait to work with you on what's next—thank you for your support!
+1 -3
View File
@@ -24,7 +24,7 @@ ARG OS_VERSION=22.04
FROM nvidia/cuda:${CUDA_VERSION}-base-ubuntu${OS_VERSION} FROM nvidia/cuda:${CUDA_VERSION}-base-ubuntu${OS_VERSION}
# Define Python version argument # Define Python version argument
ARG PYTHON_VERSION=3.12 ARG PYTHON_VERSION=3.10
# Configure environment variables # Configure environment variables
ENV DEBIAN_FRONTEND=noninteractive \ ENV DEBIAN_FRONTEND=noninteractive \
@@ -85,8 +85,6 @@ RUN if [ "$UNBOUND_DEPS" = "true" ]; then \
RUN uv pip install --no-cache ".[all]" RUN uv pip install --no-cache ".[all]"
RUN chmod +x /lerobot/.venv/lib/python${PYTHON_VERSION}/site-packages/triton/backends/nvidia/bin/ptxas
# Copy the rest of the application source code # Copy the rest of the application source code
# Make sure to have the git-LFS files for testing # Make sure to have the git-LFS files for testing
COPY --chown=user_lerobot:user_lerobot . . COPY --chown=user_lerobot:user_lerobot . .
+1 -1
View File
@@ -19,7 +19,7 @@
# docker run -it --rm lerobot-user # docker run -it --rm lerobot-user
# Configure the base image # Configure the base image
ARG PYTHON_VERSION=3.12 ARG PYTHON_VERSION=3.10
FROM python:${PYTHON_VERSION}-slim FROM python:${PYTHON_VERSION}-slim
# Configure environment variables # Configure environment variables
-2
View File
@@ -31,8 +31,6 @@
title: Using Subtasks in the Dataset title: Using Subtasks in the Dataset
- local: streaming_video_encoding - local: streaming_video_encoding
title: Streaming Video Encoding title: Streaming Video Encoding
- local: multi_dataset_training
title: Multi-Dataset Training
title: "Datasets" title: "Datasets"
- sections: - sections:
- local: act - local: act
+1 -1
View File
@@ -48,7 +48,7 @@ python -m lerobot.async_inference.robot_client \
--task="dummy" \ # POLICY: The task to run the policy on (`Fold my t-shirt`). Not necessarily defined for all policies, such as `act` --task="dummy" \ # POLICY: The task to run the policy on (`Fold my t-shirt`). Not necessarily defined for all policies, such as `act`
--policy_type=your_policy_type \ # POLICY: the type of policy to run (smolvla, act, etc) --policy_type=your_policy_type \ # POLICY: the type of policy to run (smolvla, act, etc)
--pretrained_name_or_path=user/model \ # POLICY: the model name/path on server to the checkpoint to run (e.g., lerobot/smolvla_base) --pretrained_name_or_path=user/model \ # POLICY: the model name/path on server to the checkpoint to run (e.g., lerobot/smolvla_base)
--policy_device=mps \ # POLICY: the device to run the policy on, on the server (cuda, mps, xpu, cpu) --policy_device=mps \ # POLICY: the device to run the policy on, on the server
--actions_per_chunk=50 \ # POLICY: the number of actions to output at once --actions_per_chunk=50 \ # POLICY: the number of actions to output at once
--chunk_size_threshold=0.5 \ # CLIENT: the threshold for the chunk size before sending a new observation to the server --chunk_size_threshold=0.5 \ # CLIENT: the threshold for the chunk size before sending a new observation to the server
--aggregate_fn_name=weighted_average \ # CLIENT: the function to aggregate actions on overlapping portions --aggregate_fn_name=weighted_average \ # CLIENT: the function to aggregate actions on overlapping portions
+4 -4
View File
@@ -32,7 +32,7 @@ version = "0.1.0"
dependencies = [ dependencies = [
# your policy-specific dependencies # your policy-specific dependencies
] ]
requires-python = ">= 3.12" requires-python = ">= 3.11"
[build-system] [build-system]
build-backend = # your-build-backend build-backend = # your-build-backend
@@ -82,7 +82,7 @@ Create your policy implementation by inheriting from LeRobot's base `PreTrainedP
# modeling_my_custom_policy.py # modeling_my_custom_policy.py
import torch import torch
import torch.nn as nn import torch.nn as nn
from typing import Any from typing import Dict, Any
from lerobot.policies.pretrained import PreTrainedPolicy from lerobot.policies.pretrained import PreTrainedPolicy
from .configuration_my_custom_policy import MyCustomPolicyConfig from .configuration_my_custom_policy import MyCustomPolicyConfig
@@ -91,7 +91,7 @@ class MyCustomPolicy(PreTrainedPolicy):
config_class = MyCustomPolicyConfig config_class = MyCustomPolicyConfig
name = "my_custom_policy" name = "my_custom_policy"
def __init__(self, config: MyCustomPolicyConfig, dataset_stats: dict[str, Any] = None): def __init__(self, config: MyCustomPolicyConfig, dataset_stats: Dict[str, Any] = None):
super().__init__(config, dataset_stats) super().__init__(config, dataset_stats)
... ...
``` ```
@@ -102,7 +102,7 @@ Create processor functions:
```python ```python
# processor_my_custom_policy.py # processor_my_custom_policy.py
from typing import Any from typing import Dict, Any
import torch import torch
+3 -3
View File
@@ -13,7 +13,7 @@ The EarthRover Mini Plus is a fully open source mobile robot that connects throu
### Hardware ### Hardware
- EarthRover Mini robot - EarthRover Mini robot
- Computer with Python 3.12 or newer - Computer with Python 3.10 or newer
- Internet connection - Internet connection
### Setting Up the Frodobots SDK ### Setting Up the Frodobots SDK
@@ -170,13 +170,13 @@ Once you can drive the robot well, you can start recording data to train AI mode
We use Hugging Face to store your data online. First, log in with your token from [Hugging Face settings](https://huggingface.co/settings/tokens): We use Hugging Face to store your data online. First, log in with your token from [Hugging Face settings](https://huggingface.co/settings/tokens):
```bash ```bash
hf auth login --token ${HUGGINGFACE_TOKEN} --add-to-git-credential huggingface-cli login --token ${HUGGINGFACE_TOKEN} --add-to-git-credential
``` ```
Store your Hugging Face username: Store your Hugging Face username:
```bash ```bash
HF_USER=$(hf auth whoami | awk -F': *' 'NR==1 {print $2}') HF_USER=$(huggingface-cli whoami | head -n 1)
echo $HF_USER echo $HF_USER
``` ```
+2 -2
View File
@@ -155,10 +155,10 @@ Upload your repository to Hugging Face:
pip install huggingface_hub pip install huggingface_hub
# Login to Hugging Face # Login to Hugging Face
hf auth login huggingface-cli login
# Create a new repository # Create a new repository
hf repo create my-org/my-custom-env huggingface-cli repo create my-custom-env --type space --org my-org
# Initialize git and push # Initialize git and push
git init git init
+4 -4
View File
@@ -159,7 +159,7 @@ We use the Hugging Face hub features for uploading your dataset. If you haven't
Add your token to the CLI by running this command: Add your token to the CLI by running this command:
```bash ```bash
hf auth login --token ${HUGGINGFACE_TOKEN} --add-to-git-credential huggingface-cli login --token ${HUGGINGFACE_TOKEN} --add-to-git-credential
``` ```
Then store your Hugging Face repository name in a variable: Then store your Hugging Face repository name in a variable:
@@ -327,7 +327,7 @@ You can look for other LeRobot datasets on the hub by searching for `LeRobot` [t
You can also push your local dataset to the Hub manually, running: You can also push your local dataset to the Hub manually, running:
```bash ```bash
hf upload ${HF_USER}/record-test ~/.cache/huggingface/lerobot/{repo-id} --repo-type dataset huggingface-cli upload ${HF_USER}/record-test ~/.cache/huggingface/lerobot/{repo-id} --repo-type dataset
``` ```
#### Record function #### Record function
@@ -491,7 +491,7 @@ If your local computer doesn't have a powerful GPU you could utilize Google Cola
Once training is done, upload the latest checkpoint with: Once training is done, upload the latest checkpoint with:
```bash ```bash
hf upload ${HF_USER}/act_so101_test \ huggingface-cli upload ${HF_USER}/act_so101_test \
outputs/train/act_so101_test/checkpoints/last/pretrained_model outputs/train/act_so101_test/checkpoints/last/pretrained_model
``` ```
@@ -499,7 +499,7 @@ You can also upload intermediate checkpoints with:
```bash ```bash
CKPT=010000 CKPT=010000
hf upload ${HF_USER}/act_so101_test${CKPT} \ huggingface-cli upload ${HF_USER}/act_so101_test${CKPT} \
outputs/train/act_so101_test/checkpoints/${CKPT}/pretrained_model outputs/train/act_so101_test/checkpoints/${CKPT}/pretrained_model
``` ```
+13 -60
View File
@@ -1,8 +1,8 @@
# Installation # Installation
This guide uses `conda` (via miniforge) to manage environments (recommended). If you prefer another environment manager (e.g. `uv`, `venv`), ensure you have Python >=3.12 and `ffmpeg` installed with the `libsvtav1` encoder, then skip ahead to [Environment Setup](#step-2-environment-setup). This guide uses conda (via miniforge) to manage environments. If you prefer another environment manager (e.g. `uv`, `venv`), ensure you have Python >=3.10 and ffmpeg installed with the `libsvtav1` encoder, then skip ahead to [Install LeRobot](#step-3-install-lerobot-).
## Step 1 (`conda` only): Install [`miniforge`](https://conda-forge.org/download/) ## Step 1: Install [`miniforge`](https://conda-forge.org/download/)
```bash ```bash
wget "https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-$(uname)-$(uname -m).sh" wget "https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-$(uname)-$(uname -m).sh"
@@ -11,47 +11,22 @@ bash Miniforge3-$(uname)-$(uname -m).sh
## Step 2: Environment Setup ## Step 2: Environment Setup
Create a virtual environment with Python 3.12: Create a virtual environment with Python 3.10, using conda:
<!-- prettier-ignore-start -->
<hfoptions id="create_venv">
<hfoption id="conda">
```bash ```bash
conda create -y -n lerobot python=3.12 conda create -y -n lerobot python=3.10
``` ```
</hfoption>
<hfoption id="uv"> Then activate your conda environment, you have to do this each time you open a shell to use lerobot:
```bash ```bash
uv python install 3.12
uv venv --python 3.12
```
</hfoption>
</hfoptions>
<!-- prettier-ignore-end -->
Then activate your virtual environment, you have to do this each time you open a shell to use lerobot:
<!-- prettier-ignore-start -->
<hfoptions id="activate_venv">
<hfoption id="conda">```bash
conda activate lerobot conda activate lerobot
```</hfoption>
<hfoption id="uv">
```bash
# Linux/macOSsource
source .venv/bin/activate
# Windows PowerShell
source .venv\Scripts\Activate.ps1
``` ```
</hfoption>
</hfoptions>
<!-- prettier-ignore-end -->
When using `conda`, install `ffmpeg` in your environment: When using `conda`, install `ffmpeg` in your environment:
```bash ```bash
conda install ffmpeg -c conda-forge conda install ffmpeg -c conda-forge
ffmpeg -version # ffmpeg 8.X is not yet supported !
``` ```
> [!TIP] > [!TIP]
@@ -72,9 +47,6 @@ ffmpeg -version # ffmpeg 8.X is not yet supported !
> conda install evdev -c conda-forge > conda install evdev -c conda-forge
> ``` > ```
> [!IMPORTANT]
> If you are using `uv` you will have to install `ffmpeg` system-wide (outside of the virtual environment). You rely on `uv` and `torchcodec` ability to dynamically link to the system `ffmpeg`.
## Step 3: Install LeRobot 🤗 ## Step 3: Install LeRobot 🤗
### From Source ### From Source
@@ -88,45 +60,23 @@ cd lerobot
Then, install the library in editable mode. This is useful if you plan to contribute to the code. Then, install the library in editable mode. This is useful if you plan to contribute to the code.
<!-- prettier-ignore-start -->
<hfoptions id="install_lerobot_src">
<hfoption id="conda">
```bash ```bash
pip install -e . pip install -e .
``` ```
</hfoption>
<hfoption id="uv">
```bash
uv pip install -e .
```
</hfoption>
</hfoptions>
<!-- prettier-ignore-end -->
### Installation from PyPI ### Installation from PyPI
**Core Library:** **Core Library:**
Install the base package with: Install the base package with:
<!-- prettier-ignore-start -->
<hfoptions id="install_lerobot_pypi">
<hfoption id="conda">
```bash ```bash
pip install lerobot pip install lerobot
``` ```
</hfoption>
<hfoption id="uv">
```bash
uv pip install lerobot
```
</hfoption>
</hfoptions>
<!-- prettier-ignore-end -->
_This installs only the default dependencies._ _This installs only the default dependencies._
**Extra Features:** **Extra Features:**
To install additional functionality, use one of the following (If you are using `uv`, replace `pip install` with `uv pip install` in the commands below.): To install additional functionality, use one of the following:
```bash ```bash
pip install 'lerobot[all]' # All available features pip install 'lerobot[all]' # All available features
@@ -140,10 +90,13 @@ _Replace `[...]` with your desired features._
For a full list of optional dependencies, see: For a full list of optional dependencies, see:
https://pypi.org/project/lerobot/ https://pypi.org/project/lerobot/
> [!NOTE]
> For lerobot 0.4.0, if you want to install pi, you will have to do: `pip install "lerobot[pi]@git+https://github.com/huggingface/lerobot.git"`
### Troubleshooting ### Troubleshooting
If you encounter build errors, you may need to install additional dependencies: `cmake`, `build-essential`, and `ffmpeg libs`. If you encounter build errors, you may need to install additional dependencies: `cmake`, `build-essential`, and `ffmpeg libs`.
To install these for Linux run: To install these for linux run:
```bash ```bash
sudo apt-get install cmake build-essential python3-dev pkg-config libavformat-dev libavcodec-dev libavdevice-dev libavutil-dev libswscale-dev libswresample-dev libavfilter-dev sudo apt-get install cmake build-essential python3-dev pkg-config libavformat-dev libavcodec-dev libavdevice-dev libavutil-dev libswscale-dev libswresample-dev libavfilter-dev
@@ -153,7 +106,7 @@ For other systems, see: [Compiling PyAV](https://pyav.org/docs/develop/overview/
## Optional dependencies ## Optional dependencies
LeRobot provides optional extras for specific functionalities. Multiple extras can be combined (e.g., `.[aloha,feetech]`). For all available extras, refer to `pyproject.toml`. If you are using `uv`, replace `pip install` with `uv pip install` in the commands below. LeRobot provides optional extras for specific functionalities. Multiple extras can be combined (e.g., `.[aloha,feetech]`). For all available extras, refer to `pyproject.toml`.
### Simulations ### Simulations
+2 -2
View File
@@ -279,13 +279,13 @@ We use the Hugging Face hub features for uploading your dataset. If you haven't
Add your token to the CLI by running this command: Add your token to the CLI by running this command:
```bash ```bash
hf auth login --token ${HUGGINGFACE_TOKEN} --add-to-git-credential huggingface-cli login --token ${HUGGINGFACE_TOKEN} --add-to-git-credential
``` ```
Then store your Hugging Face repository name in a variable: Then store your Hugging Face repository name in a variable:
```bash ```bash
HF_USER=$(hf auth whoami | awk -F': *' 'NR==1 {print $2}') HF_USER=$(huggingface-cli whoami | head -n 1)
echo $HF_USER echo $HF_USER
``` ```
-232
View File
@@ -1,232 +0,0 @@
# Multi-Dataset Training
This guide covers how to train a single policy on multiple heterogeneous datasets using `MultiLeRobotDataset`.
## Overview
Real-world robot learning datasets come from different environments, robots, and camera setups. A RoboCasa dataset might have three cameras named `robot0_agentview_left`, `robot0_agentview_right`, and `robot0_eye_in_hand`, while a LIBERO dataset uses `observation.images.front` and `observation.images.wrist`, and a RoboMME dataset uses bare `image` and `wrist_image` keys. State and action dimensions also differ.
`MultiLeRobotDataset` lets you train on all of them jointly by:
- **Mapping** each dataset's feature keys into a shared namespace
- **Padding** features that a dataset doesn't have with zeros
- **Weighting** how often each dataset is sampled
- **Transforming** samples per-dataset (e.g. padding actions to a common dimension)
- **Aggregating** statistics across all sub-datasets for normalization
## Configuration
Multi-dataset training is configured via `MultiDatasetConfig` in a YAML config file. Instead of a single `dataset.repo_id`, you provide a `datasets` list where each entry is a `SubDatasetConfig`.
### SubDatasetConfig fields
| Field | Type | Default | Description |
|-------|------|---------|-------------|
| `repo_id` | `str` | required | HuggingFace repo ID or local dataset name |
| `root` | `str \| None` | `None` | Local root directory for the dataset |
| `episodes` | `list[int] \| None` | `None` | Subset of episode indices to use |
| `revision` | `str \| None` | `None` | Dataset version / revision |
| `video_backend` | `str` | auto | Video decoding backend (`pyav`, `torchcodec`, etc.) |
| `weight` | `float` | `1.0` | Relative sampling weight for this dataset |
| `feature_map` | `dict[str, str]` | `{}` | Maps dataset keys to unified policy keys |
| `transforms` | `list` | `None` | Per-dataset transform steps (applied per sample) |
### Example: Three-dataset config
```yaml
dataset:
type: multi
use_imagenet_stats: true
datasets:
# RoboCasa: 3 cameras, state(16), action(12)
- repo_id: pepijn223/robocasa_PrepareCoffee
root: /data/robocasa_PrepareCoffee
weight: 1.0
feature_map:
observation.images.robot0_agentview_left: observation.images.front_left
observation.images.robot0_agentview_right: observation.images.front_right
observation.images.robot0_eye_in_hand: observation.images.wrist
# LIBERO-plus: 2 cameras, state(8), action(7)
- repo_id: pepijn223/libero_plus_lerobot
root: /data/libero_plus_lerobot
weight: 0.5
feature_map:
observation.images.front: observation.images.front_left
observation.images.wrist: observation.images.wrist
transforms:
- type: pad_action
kwargs: {target_dim: 12}
- type: pad_state
kwargs: {target_dim: 16}
# RoboMME: 2 cameras (non-standard keys), state(8), action(8)
- repo_id: pepijn223/robomme_data_lerobot
root: /data/robomme_data_lerobot
weight: 0.3
feature_map:
image: observation.images.front_left
wrist_image: observation.images.wrist
state: observation.state
actions: action
transforms:
- type: pad_action
kwargs: {target_dim: 12}
- type: pad_state
kwargs: {target_dim: 16}
```
## Feature Mapping
The `feature_map` dictionary renames dataset-local keys into a shared namespace. Keys not listed pass through unchanged. In the example above, all three datasets end up with the same camera key names (`observation.images.front_left`, `observation.images.wrist`) even though they use different conventions internally.
After mapping, the **union** of all features across datasets defines the unified schema. If a feature exists in some datasets but not others, it is automatically zero-padded for datasets that lack it, and a boolean `{key}_is_pad` flag is added to the sample so the policy can optionally mask padded features.
## Automatic Padding
When a sub-dataset doesn't have a feature that exists in the unified schema:
- **Images/videos**: padded with a black frame (zeros) matching the expected resolution
- **Float tensors** (state, action): padded with zeros
- **Integer/bool tensors**: padded with zeros / False
A companion `{key}_is_pad = True` tensor is added so the model can distinguish real data from padding.
## Per-Dataset Transforms
Each sub-dataset can have its own `transforms` pipeline that runs after feature renaming but before cross-dataset padding. This is useful for making shapes compatible before PyTorch's collate function stacks the batch.
### Built-in transforms
| Name | Description | Parameters |
|------|-------------|------------|
| `pad_action` | Zero-pad `action` to a target dimension | `target_dim: int` |
| `pad_state` | Zero-pad `observation.state` to a target dimension | `target_dim: int` |
| `resize_images` | Resize all `observation.images.*` tensors | `height: int`, `width: int` |
### Custom transforms
You can register your own transforms in `lerobot/datasets/transforms.py`:
```python
from lerobot.datasets.transforms import DatasetTransformStep, register_dataset_transform
@register_dataset_transform("my_transform")
class MyTransform(DatasetTransformStep):
def __init__(self, some_param: int):
self.some_param = some_param
def __call__(self, sample: dict) -> dict:
# Modify sample in-place or return a new dict
sample["action"] = sample["action"] * self.some_param
return sample
```
Then reference it in the config:
```yaml
transforms:
- type: my_transform
kwargs: {some_param: 2}
```
## Weighted Sampling
The `weight` field on each sub-dataset controls how often it is sampled during training. Weights are relative and automatically normalized to probabilities. For example, with weights `[1.0, 0.5, 0.3]`, the first dataset is sampled roughly 56% of the time, the second 28%, and the third 16%.
This uses `WeightedEpisodeAwareSampler`, which respects episode boundaries (so `drop_n_last_frames` and similar policy settings work correctly) while sampling across datasets proportionally.
## Stats Aggregation
Normalization statistics (mean, std, min, max, quantiles) are automatically aggregated across all sub-datasets using the mapped feature keys. The aggregation uses a weighted parallel variance algorithm so that datasets with more frames contribute proportionally to the global statistics.
The aggregated stats are used by the standard LeRobot preprocessor for normalization during training.
## Training
Launch training the same way as single-dataset training. The factory and training script automatically detect `MultiDatasetConfig` and set up the weighted sampler:
```bash
python -m lerobot.scripts.lerobot_train \
--config_path path/to/multi_dataset_config.yaml
```
## Architecture
The data flow during training with `MultiLeRobotDataset`:
```
┌─────────────────────────────────────────────────────────┐
│ MultiLeRobotDataset.__getitem__(global_idx) │
│ │
│ 1. Map global_idx → (dataset_idx, local_idx) │
│ 2. Fetch sample from sub-dataset │
│ 3. Rename keys via feature_map │
│ 4. Apply per-dataset transforms (pad_action, etc.) │
│ 5. Zero-pad missing features + add _is_pad flags │
│ 6. Add dataset_index tag │
└─────────────────────┬───────────────────────────────────┘
┌────────────▼────────────┐
│ PyTorch DataLoader │
│ (collates into batch) │
└────────────┬────────────┘
┌────────────▼────────────┐
│ LeRobot Preprocessor │
│ (normalize, tokenize) │
└────────────┬────────────┘
┌────────────▼────────────┐
│ Policy forward + loss │
└─────────────────────────┘
```
## API Reference
### `NewMultiLeRobotDataset`
```python
from lerobot.datasets.multi_dataset import NewMultiLeRobotDataset
dataset = NewMultiLeRobotDataset(
configs=[...], # list[SubDatasetConfig]
image_transforms=None, # optional image augmentation
delta_timestamps=None, # optional temporal neighbors
tolerance_s=1e-4, # timestamp tolerance
)
dataset.num_frames # total frames across all sub-datasets
dataset.num_episodes # total episodes
dataset.meta # MultiDatasetMeta (stats, features, episodes)
dataset.dataset_weights # list of per-dataset weights
dataset.features # unified feature dict (union of all mapped features)
dataset.camera_keys # unified camera key list
```
### `WeightedEpisodeAwareSampler`
```python
from lerobot.datasets.sampler import WeightedEpisodeAwareSampler
sampler = WeightedEpisodeAwareSampler(
dataset_from_indices=dataset.meta.episodes["dataset_from_index"],
dataset_to_indices=dataset.meta.episodes["dataset_to_index"],
dataset_membership=dataset.meta.episodes["dataset_source"],
dataset_weights=dataset.dataset_weights,
shuffle=True,
)
```
### `DatasetTransformPipeline`
```python
from lerobot.datasets.transforms import DatasetTransformPipeline, DatasetTransformStepConfig
pipeline = DatasetTransformPipeline([
DatasetTransformStepConfig(type="pad_action", kwargs={"target_dim": 12}),
DatasetTransformStepConfig(type="pad_state", kwargs={"target_dim": 16}),
])
sample = pipeline(sample) # modifies the sample dict
```
+5
View File
@@ -34,6 +34,11 @@ As described by Physical Intelligence, while AI has achieved remarkable success
pip install -e ".[pi]" pip install -e ".[pi]"
``` ```
> [!NOTE]
> For lerobot 0.4.0, if you want to install pi tag, you will have to do: `pip install "lerobot[pi]@git+https://github.com/huggingface/lerobot.git"`.
>
> This will be solved in the next patch release
## Training Data and Capabilities ## Training Data and Capabilities
π₀ is trained on the largest robot interaction dataset to date, combining three key data sources: π₀ is trained on the largest robot interaction dataset to date, combining three key data sources:
+5
View File
@@ -36,6 +36,11 @@ This diverse training mixture creates a "curriculum" that enables generalization
pip install -e ".[pi]" pip install -e ".[pi]"
``` ```
> [!NOTE]
> For lerobot 0.4.0, if you want to install pi tag, you will have to do: `pip install "lerobot[pi]@git+https://github.com/huggingface/lerobot.git"`.
>
> This will be solved in the next patch release
## Usage ## Usage
To use π₀.₅ in your LeRobot configuration, specify the policy type as: To use π₀.₅ in your LeRobot configuration, specify the policy type as:
+5
View File
@@ -43,6 +43,11 @@ This approach can transform **any existing VLM** into a VLA by training it to pr
pip install -e ".[pi]" pip install -e ".[pi]"
``` ```
> [!NOTE]
> For lerobot 0.4.0, if you want to install the pi tag, you will have to do: `pip install "lerobot[pi]@git+https://github.com/huggingface/lerobot.git"`.
>
> This will be solved in the next patch release
## Training a Custom FAST Tokenizer ## Training a Custom FAST Tokenizer
You have two options for the FAST tokenizer: You have two options for the FAST tokenizer:
+186 -140
View File
@@ -1,49 +1,23 @@
# Unitree G1 # Unitree G1
<img This guide covers the complete setup process for the Unitree G1 humanoid, from initial connection to running gr00t_wbc locomotion.
src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/lerobot/unitree_thumbnail.jpg"
alt="Unitree G1 locomanipulation demo"
style={{ width: "100%" }}
/>
The Unitree G1 humanoid is now supported in LeRobot! You can teleoperate, train locomanipulation policies, test in sim, and more. Both 29 and 23 DoF variants are supported. ## About
We support both 29 and 23 DOF G1 EDU version. We introduce:
- **`unitree g1` robot class, handling low level read/write from/to the humanoid**
- **ZMQ socket bridge** for remote communication and camera streaming, allowing for remote policy deployment over wlan, eth or directly on the robot
- **Locomotion policies** from NVIDIA gr00t and Amazon FAR Holosoma
- **Simulation mode** for testing policies without the physical robot in mujoco
--- ---
## Part 1: Getting Started ## Connection guide
### Install LeRobot on Your Machine ### Step 1: Configure Ethernet Interface
```bash Set a static IP on the same subnet as the robot:
conda create -y -n lerobot python=3.12
conda activate lerobot
git clone https://github.com/unitreerobotics/unitree_sdk2_python.git
cd unitree_sdk2_python && pip install -e .
git clone https://github.com/huggingface/lerobot.git
cd lerobot
pip install -e '.[unitree_g1]'
```
### Test the Installation (Simulation)
```bash
lerobot-teleoperate \
--robot.type=unitree_g1 \
--robot.is_simulation=true \
--teleop.type=unitree_g1 \
--teleop.id=wbc_unitree \
--robot.cameras='{"global_view": {"type": "zmq", "server_address": "localhost", "port": 5555, "camera_name": "head_camera", "width": 640, "height": 480, "fps": 30}}' \
--display_data=true
```
This will launch a [MuJoCo sim instance](https://huggingface.co/lerobot/unitree-g1-mujoco/tree/main) for the G1.
- Press `9` to release the robot
- Press `7` / `8` to increase / decrease waist height
### Connect to the Robot
The G1's Ethernet IP is fixed at `192.168.123.164`. Your machine must have a static IP on the same subnet: `192.168.123.x` where `x ≠ 164`.
```bash ```bash
# Replace 'enp131s0' with your ethernet interface name (check with `ip a`) # Replace 'enp131s0' with your ethernet interface name (check with `ip a`)
@@ -52,200 +26,272 @@ sudo ip addr add 192.168.123.200/24 dev enp131s0
sudo ip link set enp131s0 up sudo ip link set enp131s0 up
``` ```
### SSH into the Robot **Note**: The G1's Ethernet IP is fixed at `192.168.123.164`. Your computer must use `192.168.123.x` with x ≠ 164.
### Step 2: SSH into the Robot
```bash ```bash
ssh unitree@192.168.123.164 ssh unitree@192.168.123.164
# Password: 123 # Password: 123
``` ```
### Install LeRobot on the G1 You should now be connected to the G1's Orin.
From the robot:
```bash
conda create -y -n lerobot python=3.12
conda activate lerobot
git clone https://github.com/unitreerobotics/unitree_sdk2_python.git
cd unitree_sdk2_python && pip install -e .
git clone https://github.com/huggingface/lerobot.git
cd lerobot
pip install -e '.[unitree_g1]'
```
> **Note:** The Unitree SDK requires CycloneDDS v0.10.2. See the [Unitree SDK docs](https://github.com/unitreerobotics/unitree_sdk2_python) for details.
--- ---
## Part 2: Enable WiFi on the Robot ## Part 2: Enable WiFi on the Robot
Wi-Fi connectivity is blocked by default on the G1. To activate: Wlan0 is disabled by default on the G1. To enable it:
### Step 1: Enable WiFi Hardware
```bash ```bash
sudo rfkill unblock wifi
sudo rfkill unblock all sudo rfkill unblock all
# Bring up wlan0
sudo ip link set wlan0 up sudo ip link set wlan0 up
# Enable NetworkManager control of wlan0
sudo nmcli radio wifi on sudo nmcli radio wifi on
sudo nmcli device set wlan0 managed yes sudo nmcli device set wlan0 managed yes
sudo systemctl restart NetworkManager sudo systemctl restart NetworkManager
``` ```
**On your laptop** (share internet via Ethernet): ### Step 2: Enable Internet Forwarding
**On your laptop:**
```bash ```bash
# Enable IP forwarding
sudo sysctl -w net.ipv4.ip_forward=1 sudo sysctl -w net.ipv4.ip_forward=1
# Replace wlp132s0f0 with your WiFi interface name # Set up NAT (replace wlp132s0f0 with your WiFi interface)
sudo iptables -t nat -A POSTROUTING -o wlp132s0f0 -s 192.168.123.0/24 -j MASQUERADE sudo iptables -t nat -A POSTROUTING -o wlp132s0f0 -s 192.168.123.0/24 -j MASQUERADE
sudo iptables -A FORWARD -i wlp132s0f0 -o enp131s0 -m state --state RELATED,ESTABLISHED -j ACCEPT sudo iptables -A FORWARD -i wlp132s0f0 -o enp131s0 -m state --state RELATED,ESTABLISHED -j ACCEPT
sudo iptables -A FORWARD -i enp131s0 -o wlp132s0f0 -j ACCEPT sudo iptables -A FORWARD -i enp131s0 -o wlp132s0f0 -j ACCEPT
``` ```
**On the G1** (set default route through your laptop): **On the G1:**
```bash ```bash
# Add laptop as default gateway
sudo ip route del default 2>/dev/null || true sudo ip route del default 2>/dev/null || true
sudo ip route add default via 192.168.123.200 dev eth0 sudo ip route add default via 192.168.123.200 dev eth0
echo "nameserver 8.8.8.8" | sudo tee /etc/resolv.conf echo "nameserver 8.8.8.8" | sudo tee /etc/resolv.conf
# Verify # Test connection
ping -c 3 8.8.8.8 ping -c 3 8.8.8.8
``` ```
**Connect to a WiFi network:** ### Step 3: Connect to WiFi Network
```bash ```bash
# List available networks
nmcli device wifi list nmcli device wifi list
# Connect to your WiFi (example)
sudo nmcli connection add type wifi ifname wlan0 con-name "YourNetwork" ssid "YourNetwork" sudo nmcli connection add type wifi ifname wlan0 con-name "YourNetwork" ssid "YourNetwork"
sudo nmcli connection modify "YourNetwork" wifi-sec.key-mgmt wpa-psk sudo nmcli connection modify "YourNetwork" wifi-sec.key-mgmt wpa-psk
sudo nmcli connection modify "YourNetwork" wifi-sec.psk "YourPassword" sudo nmcli connection modify "YourNetwork" wifi-sec.psk "YourPassword"
sudo nmcli connection modify "YourNetwork" connection.autoconnect yes sudo nmcli connection modify "YourNetwork" connection.autoconnect yes
sudo nmcli connection up "YourNetwork" sudo nmcli connection up "YourNetwork"
# Check WiFi IP address
ip a show wlan0 ip a show wlan0
``` ```
You can now SSH over WiFi: ### Step 4: SSH Over WiFi
Once connected to WiFi, note the robot's IP address and disconnect the Ethernet cable. You can now SSH over WiFi:
```bash ```bash
ssh unitree@<ROBOT_WIFI_IP> ssh unitree@<YOUR_ROBOT_IP>
# Password: 123 # Password: 123
``` ```
Replace `<YOUR_ROBOT_IP>` with your robot's actual WiFi IP address.
--- ---
## Part 3: Teleoperation & Locomotion ## Part 3: Robot Server Setup
### Run the Robot Server ### Step 1: Install LeRobot on the Orin
SSH into the robot and install LeRobot:
```bash
ssh unitree@<YOUR_ROBOT_IP>
conda create -y -n lerobot python=3.10
conda activate lerobot
git clone https://github.com/huggingface/lerobot.git
cd lerobot
pip install -e '.[unitree_g1]'
git clone https://github.com/unitreerobotics/unitree_sdk2_python.git
cd unitree_sdk2_python && pip install -e .
```
**Note**: The Unitree SDK requires CycloneDDS v0.10.2 to be installed. See the [Unitree SDK documentation](https://github.com/unitreerobotics/unitree_sdk2_python) for details.
### Step 2: Run the Robot Server
On the robot: On the robot:
```bash ```bash
python src/lerobot/robots/unitree_g1/run_g1_server.py --camera python src/lerobot/robots/unitree_g1/run_g1_server.py
``` ```
### Run the Locomotion Policy **Important**: Keep this terminal running. The server must be active for remote control.
```bash
lerobot-teleoperate \
--robot.type=unitree_g1 \
--robot.is_simulation=false \
--robot.robot_ip=<ROBOT_IP> \
--teleop.type=unitree_g1 \
--teleop.id=wbc_unitree \
--robot.cameras='{"global_view": {"type": "zmq", "server_address": "<ROBOT_IP>", "port": 5555, "camera_name": "head_camera", "width": 640, "height": 480, "fps": 30}}' \
--display_data=true \
--robot.controller=HolosomaLocomotionController
```
We support both [HolosomaLocomotionController](https://github.com/amazon-far/holosoma) and [GrootLocomotionController](https://github.com/NVlabs/GR00T-WholeBodyControl).
--- ---
## Part 4: Loco-Manipulation with the Homunculus Exoskeleton ## Part 4: Controlling the robot
We provide a loco-manipulation solution via the Homunculus Exoskeleton — an open-source 7 DoF exoskeleton for whole-body control. Assembly instructions [here](https://github.com/nepyope/hmc_exo). With the robot server running, you can now control the robot remotely. Let's launch a locomotion policy
### Calibrate ### Step 1: Install LeRobot on your machine
```bash
conda create -y -n lerobot python=3.10
conda activate lerobot
git clone https://github.com/huggingface/lerobot.git
cd lerobot
pip install -e '.[unitree_g1]'
git clone https://github.com/unitreerobotics/unitree_sdk2_python.git
cd unitree_sdk2_python && pip install -e .
```
### Step 2: Update Robot IP in Config
Edit the config file to match your robot's WiFi IP:
```python
# In src/lerobot/robots/unitree_g1/config_unitree_g1.py
robot_ip: str = "<YOUR_ROBOT_IP>" # Replace with your robot's WiFi IP.
```
### Step 3: Run the Locomotion Policy
```bash
# Run GR00T locomotion controller
python examples/unitree_g1/gr00t_locomotion.py --repo-id "nepyope/GR00T-WholeBodyControl_g1"
# Run Holosoma locomotion controller
python examples/unitree_g1/holosoma_locomotion.py
```
Press `Ctrl+C` to stop the policy.
---
## Running in Simulation Mode (MuJoCo)
You can test policies before deploying on the physical robot using MuJoCo simulation. Set `is_simulation=True` in config or pass `--robot.is_simulation=true` via CLI.
### Calibrate Exoskeleton Teleoperator
```bash ```bash
lerobot-calibrate \ lerobot-calibrate \
--teleop.type=unitree_g1 \ --teleop.type=unitree_g1 \
--teleop.left_arm_config.port=/dev/ttyACM1 \ --teleop.left_arm_config.port=/dev/ttyACM1 \
--teleop.right_arm_config.port=/dev/ttyACM0 \ --teleop.right_arm_config.port=/dev/ttyACM0 \
--teleop.id=exo --teleop.id=exo
``` ```
During calibration move each joint through its entire range. After fitting, move the joint in a neutral position and press `n` to advance. ### Teleoperate in Simulation
### Record a Dataset ```bash
lerobot-teleoperate \
--robot.type=unitree_g1 \
--robot.is_simulation=true \
--teleop.type=unitree_g1 \
--teleop.left_arm_config.port=/dev/ttyACM1 \
--teleop.right_arm_config.port=/dev/ttyACM0 \
--teleop.id=exo \
--fps=100
```
### Record Dataset in Simulation
```bash ```bash
lerobot-record \ lerobot-record \
--robot.type=unitree_g1 \ --robot.type=unitree_g1 \
--robot.is_simulation=true \ --robot.is_simulation=true \
--robot.cameras='{"global_view": {"type": "zmq", "server_address": "localhost", "port": 5555, "camera_name": "head_camera", "width": 640, "height": 480, "fps": 30}}' \ --robot.cameras='{"global_view": {"type": "zmq", "server_address": "localhost", "port": 5555, "camera_name": "head_camera", "width": 640, "height": 480, "fps": 30}}' \
--teleop.type=unitree_g1 \ --teleop.type=unitree_g1 \
--teleop.left_arm_config.port=/dev/ttyACM1 \ --teleop.left_arm_config.port=/dev/ttyACM1 \
--teleop.right_arm_config.port=/dev/ttyACM0 \ --teleop.right_arm_config.port=/dev/ttyACM0 \
--teleop.id=exo \ --teleop.id=exo \
--dataset.repo_id=your-username/dataset-name \ --dataset.repo_id=your-username/dataset-name \
--dataset.single_task="Test" \ --dataset.single_task="Test" \
--dataset.num_episodes=2 \ --dataset.num_episodes=2 \
--dataset.episode_time_s=5 \ --dataset.episode_time_s=5 \
--dataset.reset_time_s=5 \ --dataset.reset_time_s=5 \
--dataset.push_to_hub=true \ --dataset.push_to_hub=true \
--dataset.streaming_encoding=true \ --dataset.streaming_encoding=true \
--dataset.encoder_threads=2 # --dataset.vcodec=auto \
--dataset.encoder_threads=2
``` ```
> **Note:** Omit `--teleop.left_arm_config.port` and `--teleop.right_arm_config.port` if you're only using the joystick. Example simulation dataset: [nepyope/teleop_test_sim](https://huggingface.co/datasets/nepyope/teleop_test_sim)
Example dataset: [nepyope/unitree_box_move_blue_full](https://huggingface.co/datasets/nepyope/unitree_box_move_blue_full)
--- ---
## Part 5: Training & Inference ## Running on Real Robot
### Train Once the robot server is running on the G1 (see Part 3), you can teleoperate and record on the real robot.
### Start the Camera Server
On the robot, start the ZMQ image server:
```bash ```bash
python src/lerobot/scripts/lerobot_train.py \ python src/lerobot/cameras/zmq/image_server.py
--dataset.repo_id=your-username/dataset-name \
--policy.type=pi05 \
--output_dir=./outputs/pi05_training \
--job_name=pi05_training \
--policy.repo_id=your-username/your-repo-id \
--policy.pretrained_path=lerobot/pi05_base \
--policy.compile_model=true \
--policy.gradient_checkpointing=true \
--wandb.enable=true \
--policy.dtype=bfloat16 \
--policy.freeze_vision_encoder=false \
--policy.train_expert_only=false \
--steps=3000 \
--policy.device=cuda \
--batch_size=32
``` ```
### Inference with RTC Keep this running in a separate terminal for camera streaming during recording.
Once trained, we recommend deploying policies using inference-time RTC: ### Teleoperate Real Robot
```bash ```bash
python examples/rtc/eval_with_real_robot.py \ lerobot-teleoperate \
--policy.path=your-username/your-repo-id \ --robot.type=unitree_g1 \
--policy.device=cuda \ --robot.is_simulation=false \
--robot.type=unitree_g1 \ --teleop.type=unitree_g1 \
--robot.is_simulation=false \ --teleop.left_arm_config.port=/dev/ttyACM1 \
--robot.controller=HolosomaLocomotionController \ --teleop.right_arm_config.port=/dev/ttyACM0 \
--robot.cameras='{"global_view": {"type": "zmq", "server_address": "<ROBOT_IP>", "port": 5555, "camera_name": "head_camera", "width": 640, "height": 480, "fps": 30}}' \ --teleop.id=exo \
--task="task_description" \ --fps=100
--duration=1000 \
--fps=30 \
--rtc.enabled=true
``` ```
### Record Dataset on Real Robot
```bash
lerobot-record \
--robot.type=unitree_g1 \
--robot.is_simulation=false \
--robot.cameras='{"global_view": {"type": "zmq", "server_address": "172.18.129.215", "port": 5555, "camera_name": "head_camera", "width": 640, "height": 480, "fps": 30}}' \
--teleop.type=unitree_g1 \
--teleop.left_arm_config.port=/dev/ttyACM1 \
--teleop.right_arm_config.port=/dev/ttyACM0 \
--teleop.id=exo \
--dataset.repo_id=your-username/dataset-name \
--dataset.single_task="Test" \
--dataset.num_episodes=2 \
--dataset.episode_time_s=5 \
--dataset.reset_time_s=5 \
--dataset.push_to_hub=true \
--dataset.streaming_encoding=true \
# --dataset.vcodec=auto \
--dataset.encoder_threads=2
```
**Note**: Update `server_address` to match your robot's camera server IP.
Example real robot dataset: [nepyope/teleop_test_real](https://huggingface.co/datasets/nepyope/teleop_test_real)
--- ---
## Additional Resources ## Additional Resources
@@ -254,8 +300,8 @@ python examples/rtc/eval_with_real_robot.py \
- [GR00T-WholeBodyControl](https://github.com/NVlabs/GR00T-WholeBodyControl) - [GR00T-WholeBodyControl](https://github.com/NVlabs/GR00T-WholeBodyControl)
- [Holosoma](https://github.com/amazon-far/holosoma) - [Holosoma](https://github.com/amazon-far/holosoma)
- [LeRobot Documentation](https://github.com/huggingface/lerobot) - [LeRobot Documentation](https://github.com/huggingface/lerobot)
- [Unitree IL LeRobot](https://github.com/unitreerobotics/unitree_IL_lerobot) - [Unitree_IL_Lerobot](https://github.com/unitreerobotics/unitree_IL_lerobot)
--- ---
_Last updated: March 2026_ _Last updated: December 2025_
+1 -1
View File
@@ -57,7 +57,7 @@ class DatasetReplayConfig:
repo_id: str repo_id: str
# Episode to replay. # Episode to replay.
episode: int episode: int
# Root directory where the dataset will be stored (e.g. 'dataset/path'). If None, defaults to $HF_LEROBOT_HOME/repo_id. # Root directory where the dataset will be stored (e.g. 'dataset/path').
root: str | Path | None = None root: str | Path | None = None
# Limit the frames per second. By default, uses the policy fps. # Limit the frames per second. By default, uses the policy fps.
fps: int = 30 fps: int = 30
-490
View File
@@ -1,490 +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.
"""
SLURM-distributed SARM RA-BC annotation pipeline.
Computes SARM progress values for all frames in a dataset, distributed across
SLURM workers, then merges the shards into a single sarm_progress.parquet.
Two subcommands, each a separate SLURM submission:
compute N workers, each computes progress for a subset of episodes
aggregate 1 worker, merges N shards into sarm_progress.parquet, pushes to hub
Usage:
python slurm_compute_rabc.py compute \\
--repo-id user/dataset --reward-model-path user/sarm_model \\
--stride 10 --device cpu --workers 50 --partition cpu
python slurm_compute_rabc.py aggregate \\
--repo-id user/dataset --reward-model-path user/sarm_model \\
--partition cpu --push-to-hub
"""
import argparse
from pathlib import Path
from datatrove.executor import LocalPipelineExecutor
from datatrove.executor.slurm import SlurmPipelineExecutor
from datatrove.pipeline.base import PipelineStep
class ComputeProgressShards(PipelineStep):
"""Each worker computes SARM progress for its assigned episodes."""
def __init__(
self, repo_id, reward_model_path, stride=1, head_mode="sparse", device="cpu", shard_dir="rabc_shards"
):
super().__init__()
if stride < 1:
raise ValueError(f"stride must be >= 1, got {stride}")
self.repo_id = repo_id
self.reward_model_path = reward_model_path
self.stride = stride
self.head_mode = head_mode
self.device = device
self.shard_dir = shard_dir
def run(self, data=None, rank: int = 0, world_size: int = 1):
import logging
from pathlib import Path
import numpy as np
import pyarrow as pa
import pyarrow.parquet as pq
import torch
from tqdm import tqdm
from lerobot.policies.sarm.compute_rabc_weights import (
generate_all_frame_indices,
interpolate_progress,
load_sarm_resources,
)
from lerobot.utils.utils import init_logging
init_logging()
dataset, reward_model, preprocess = load_sarm_resources(
self.repo_id,
self.reward_model_path,
self.device,
)
if hasattr(preprocess, "eval"):
preprocess.eval()
for step in preprocess.steps:
if hasattr(step, "eval"):
step.eval()
image_key = reward_model.config.image_key
state_key = reward_model.config.state_key
frame_gap = reward_model.config.frame_gap
center_idx = reward_model.config.n_obs_steps // 2
dual_mode = reward_model.config.uses_dual_heads
compute_sparse = self.head_mode in ("sparse", "both") or not dual_mode
compute_dense = self.head_mode in ("dense", "both") and dual_mode
my_episodes = list(range(dataset.num_episodes))[rank::world_size]
if not my_episodes:
logging.info(f"Rank {rank}: no episodes assigned")
return
logging.info(f"Rank {rank}: {len(my_episodes)} / {dataset.num_episodes} episodes")
all_rows = []
for ep_idx in tqdm(my_episodes, desc=f"Rank {rank}"):
ep = dataset.meta.episodes[ep_idx]
ep_start, ep_end = ep["dataset_from_index"], ep["dataset_to_index"]
task = dataset[ep_start].get("task", "perform the task")
all_ep_indices = generate_all_frame_indices(ep_start, ep_end, frame_gap)
if self.stride > 1:
compute_indices = [i for i in all_ep_indices if (i - ep_start) % self.stride == 0]
if (ep_end - 1) not in compute_indices:
compute_indices.append(ep_end - 1)
compute_indices = sorted(set(compute_indices))
else:
compute_indices = all_ep_indices
frame_results = {}
for qi in tqdm(compute_indices, desc=f" Ep {ep_idx}", leave=False):
try:
sample = dataset[qi]
batch = {
image_key: sample[image_key],
"task": task,
"index": qi,
"episode_index": ep_idx,
}
if state_key in sample:
batch[state_key] = sample[state_key]
with torch.no_grad():
processed = preprocess(batch)
vf = processed["video_features"].to(self.device)
tf = processed["text_features"].to(self.device)
sf = processed.get("state_features")
if sf is not None:
sf = sf.to(self.device)
lengths = processed.get("lengths")
sparse_val = dense_val = np.nan
if compute_sparse:
r = reward_model.calculate_rewards(
text_embeddings=tf,
video_embeddings=vf,
state_features=sf,
lengths=lengths,
return_all_frames=True,
head_mode="sparse",
)
sparse_val = float(r[0, center_idx] if r.ndim == 2 else r[center_idx])
if compute_dense:
r = reward_model.calculate_rewards(
text_embeddings=tf,
video_embeddings=vf,
state_features=sf,
lengths=lengths,
return_all_frames=True,
head_mode="dense",
)
dense_val = float(r[0, center_idx] if r.ndim == 2 else r[center_idx])
frame_results[qi] = (sparse_val, dense_val)
except Exception as e:
logging.warning(f"Failed frame {qi}: {e}")
if not frame_results:
logging.warning(f"Episode {ep_idx}: all frames failed, skipping")
continue
# Interpolate to all frames in this episode
computed_idx = np.array(sorted(frame_results.keys()))
all_frame_arr = np.arange(ep_start, ep_end)
sparse_vals = np.array([frame_results[i][0] for i in computed_idx]) if compute_sparse else None
dense_vals = np.array([frame_results[i][1] for i in computed_idx]) if compute_dense else None
if self.stride > 1 and len(computed_idx) > 1:
if compute_sparse:
sparse_vals = interpolate_progress(computed_idx, sparse_vals, all_frame_arr)
if compute_dense:
dense_vals = interpolate_progress(computed_idx, dense_vals, all_frame_arr)
output_frames = all_frame_arr
else:
# Use only successfully computed frames to avoid indexing mismatch on failures
output_frames = computed_idx
for i, fi in enumerate(output_frames):
row = {"index": int(fi), "episode_index": ep_idx, "frame_index": int(fi - ep_start)}
if compute_sparse:
row["progress_sparse"] = float(sparse_vals[i])
if compute_dense:
row["progress_dense"] = float(dense_vals[i])
all_rows.append(row)
if all_rows:
import pandas as pd
df = pd.DataFrame(all_rows).sort_values("index").reset_index(drop=True)
table = pa.Table.from_pandas(df, preserve_index=False)
table = table.replace_schema_metadata({b"reward_model_path": self.reward_model_path.encode()})
shard_dir = Path(self.shard_dir)
shard_dir.mkdir(parents=True, exist_ok=True)
out = shard_dir / f"shard_{rank:05d}.parquet"
pq.write_table(table, out)
logging.info(f"Rank {rank}: saved {len(df)} rows to {out}")
class AggregateProgress(PipelineStep):
"""Merge all shard parquets into final sarm_progress.parquet."""
def __init__(self, repo_id, reward_model_path, shard_dir="rabc_shards", push_to_hub=False):
super().__init__()
self.repo_id = repo_id
self.reward_model_path = reward_model_path
self.shard_dir = shard_dir
self.push_to_hub = push_to_hub
def run(self, data=None, rank: int = 0, world_size: int = 1):
import datetime
import logging
import os
from pathlib import Path
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq
from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.utils.utils import init_logging
init_logging()
if rank != 0:
return
shard_dir = Path(self.shard_dir)
shards = sorted(shard_dir.glob("shard_*.parquet"))
if not shards:
raise FileNotFoundError(f"No shards found in {shard_dir}")
# Log shard modification time range to help detect stale files
mtimes = [os.path.getmtime(s) for s in shards]
oldest = datetime.datetime.fromtimestamp(min(mtimes)).isoformat(timespec="seconds")
newest = datetime.datetime.fromtimestamp(max(mtimes)).isoformat(timespec="seconds")
logging.info(f"Aggregating {len(shards)} shards (oldest: {oldest}, newest: {newest})")
df = pd.concat([pd.read_parquet(s) for s in shards], ignore_index=True)
df = df.sort_values("index").reset_index(drop=True)
table = pa.Table.from_pandas(df, preserve_index=False)
table = table.replace_schema_metadata({b"reward_model_path": self.reward_model_path.encode()})
temp_ds = LeRobotDataset(self.repo_id, download_videos=False)
out_path = Path(temp_ds.root) / "sarm_progress.parquet"
out_path.parent.mkdir(parents=True, exist_ok=True)
pq.write_table(table, out_path)
logging.info(f"Saved {len(df)} rows to {out_path}")
for col in ["progress_sparse", "progress_dense"]:
if col in df.columns:
v = df[col].dropna()
logging.info(
f"{col}: mean={v.mean():.4f} std={v.std():.4f} min={v.min():.4f} max={v.max():.4f}"
)
if self.push_to_hub:
from huggingface_hub import HfApi
api = HfApi()
hub_path = "sarm_progress.parquet"
logging.info(f"Uploading to {self.repo_id}/{hub_path}")
api.upload_file(
path_or_fileobj=str(out_path),
path_in_repo=hub_path,
repo_id=self.repo_id,
repo_type="dataset",
)
logging.info(f"Uploaded: https://huggingface.co/datasets/{self.repo_id}/blob/main/{hub_path}")
def make_compute_executor(
repo_id,
reward_model_path,
stride,
head_mode,
device,
shard_dir,
logs_dir,
job_name,
slurm,
workers,
partition,
cpus_per_task,
mem_per_cpu,
):
kwargs = {
"pipeline": [
ComputeProgressShards(repo_id, reward_model_path, stride, head_mode, device, str(shard_dir)),
],
"logging_dir": str(logs_dir / job_name),
}
if slurm:
kwargs.update(
{
"job_name": job_name,
"tasks": workers,
"workers": workers,
"time": "24:00:00",
"partition": partition,
"cpus_per_task": cpus_per_task,
"sbatch_args": {"mem-per-cpu": mem_per_cpu},
}
)
return SlurmPipelineExecutor(**kwargs)
kwargs.update({"tasks": workers, "workers": 1})
return LocalPipelineExecutor(**kwargs)
def make_aggregate_executor(
repo_id,
reward_model_path,
shard_dir,
logs_dir,
job_name,
slurm,
partition,
cpus_per_task,
mem_per_cpu,
push_to_hub,
):
kwargs = {
"pipeline": [
AggregateProgress(repo_id, reward_model_path, str(shard_dir), push_to_hub),
],
"logging_dir": str(logs_dir / job_name),
}
if slurm:
kwargs.update(
{
"job_name": job_name,
"tasks": 1,
"workers": 1,
"time": "02:00:00",
"partition": partition,
"cpus_per_task": cpus_per_task,
"sbatch_args": {"mem-per-cpu": mem_per_cpu},
}
)
return SlurmPipelineExecutor(**kwargs)
kwargs.update({"tasks": 1, "workers": 1})
return LocalPipelineExecutor(**kwargs)
def _add_shared_args(p):
p.add_argument(
"--repo-id",
type=str,
required=True,
help="Hugging Face repository identifier, e.g. 'user/dataset'.",
)
p.add_argument(
"--shard-dir",
type=Path,
default=Path("rabc_shards"),
help="Directory to read/write per-rank parquet shards.",
)
p.add_argument(
"--logs-dir",
type=Path,
default=Path("logs"),
help="Directory for datatrove logs.",
)
p.add_argument(
"--job-name",
type=str,
default=None,
help="SLURM job name (defaults to rabc_<subcommand>).",
)
p.add_argument(
"--slurm",
type=int,
default=1,
help="1 = submit via SLURM; 0 = run locally (useful for debugging).",
)
p.add_argument(
"--partition",
type=str,
default=None,
help="SLURM partition to submit to.",
)
p.add_argument(
"--cpus-per-task",
type=int,
default=4,
help="Number of CPUs per SLURM task.",
)
p.add_argument(
"--mem-per-cpu",
type=str,
default="4G",
help="Memory per CPU, e.g. '4G' or '1950M'.",
)
def main():
parser = argparse.ArgumentParser(
description="SLURM-distributed SARM RA-BC annotation pipeline",
formatter_class=argparse.RawDescriptionHelpFormatter,
)
sub = parser.add_subparsers(dest="command", required=True)
# compute subcommand
cp = sub.add_parser(
"compute",
help="Distribute progress computation across SLURM workers.",
)
_add_shared_args(cp)
cp.add_argument(
"--reward-model-path",
type=str,
required=True,
help="Path or HF repo id of the SARM reward model.",
)
cp.add_argument(
"--stride",
type=int,
default=1,
help="Compute every Nth frame; intermediate frames are interpolated (must be >= 1).",
)
cp.add_argument(
"--head-mode",
type=str,
default="sparse",
choices=["sparse", "dense", "both"],
help="Which reward head(s) to compute.",
)
cp.add_argument(
"--device",
type=str,
default="cpu",
help="Device for reward model inference, e.g. 'cpu' or 'cuda'.",
)
cp.add_argument(
"--workers",
type=int,
default=50,
help="Number of parallel SLURM tasks (one shard per worker).",
)
# aggregate subcommand
ap = sub.add_parser(
"aggregate",
help="Merge per-rank shards into a single sarm_progress.parquet.",
)
_add_shared_args(ap)
ap.add_argument(
"--reward-model-path",
type=str,
required=True,
help="Path or HF repo id of the SARM reward model (stored in parquet metadata).",
)
ap.add_argument(
"--push-to-hub",
action="store_true",
help="Upload sarm_progress.parquet to the Hugging Face Hub after aggregation.",
)
args = parser.parse_args()
job_name = args.job_name or f"rabc_{args.command}"
kwargs = vars(args)
kwargs["slurm"] = kwargs.pop("slurm") == 1
kwargs["job_name"] = job_name
command = kwargs.pop("command")
executor = make_compute_executor(**kwargs) if command == "compute" else make_aggregate_executor(**kwargs)
executor.run()
if __name__ == "__main__":
main()
-2
View File
@@ -78,7 +78,6 @@ from torch import Tensor
from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig # noqa: F401 from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig # noqa: F401
from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig # noqa: F401 from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig # noqa: F401
from lerobot.cameras.zmq.configuration_zmq import ZMQCameraConfig # noqa: F401
from lerobot.configs import parser from lerobot.configs import parser
from lerobot.configs.policies import PreTrainedConfig from lerobot.configs.policies import PreTrainedConfig
from lerobot.configs.types import RTCAttentionSchedule from lerobot.configs.types import RTCAttentionSchedule
@@ -98,7 +97,6 @@ from lerobot.robots import ( # noqa: F401
bi_so_follower, bi_so_follower,
koch_follower, koch_follower,
so_follower, so_follower,
unitree_g1,
) )
from lerobot.robots.utils import make_robot_from_config from lerobot.robots.utils import make_robot_from_config
from lerobot.utils.constants import OBS_IMAGES from lerobot.utils.constants import OBS_IMAGES
@@ -14,20 +14,20 @@
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
import argparse
import logging import logging
import time
from collections import deque from collections import deque
import numpy as np import numpy as np
import onnxruntime as ort import onnxruntime as ort
from huggingface_hub import hf_hub_download from huggingface_hub import hf_hub_download
from lerobot.robots.unitree_g1.g1_utils import ( from lerobot.robots.unitree_g1.config_unitree_g1 import UnitreeG1Config
REMOTE_AXES, from lerobot.robots.unitree_g1.g1_utils import G1_29_JointIndex
REMOTE_BUTTONS, from lerobot.robots.unitree_g1.unitree_g1 import UnitreeG1
G1_29_JointIndex,
get_gravity_orientation,
)
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -36,13 +36,18 @@ GROOT_DEFAULT_ANGLES[[0, 6]] = -0.1 # Hip pitch
GROOT_DEFAULT_ANGLES[[3, 9]] = 0.3 # Knee GROOT_DEFAULT_ANGLES[[3, 9]] = 0.3 # Knee
GROOT_DEFAULT_ANGLES[[4, 10]] = -0.2 # Ankle pitch GROOT_DEFAULT_ANGLES[[4, 10]] = -0.2 # Ankle pitch
MISSING_JOINTS = []
G1_MODEL = "g1_23" # Or "g1_29"
if G1_MODEL == "g1_23":
MISSING_JOINTS = [12, 14, 20, 21, 27, 28] # Waist yaw/pitch, wrist pitch/yaw
# Control parameters # Control parameters
ACTION_SCALE = 0.25 ACTION_SCALE = 0.25
CONTROL_DT = 0.02 # 50Hz CONTROL_DT = 0.02 # 50Hz
ANG_VEL_SCALE: float = 0.25 ANG_VEL_SCALE: float = 0.25
DOF_POS_SCALE: float = 1.0 DOF_POS_SCALE: float = 1.0
DOF_VEL_SCALE: float = 0.05 DOF_VEL_SCALE: float = 0.05
CMD_SCALE: list[float] = [2.0, 2.0, 0.25] CMD_SCALE: list = [2.0, 2.0, 0.25]
DEFAULT_GROOT_REPO_ID = "nepyope/GR00T-WholeBodyControl_g1" DEFAULT_GROOT_REPO_ID = "nepyope/GR00T-WholeBodyControl_g1"
@@ -80,11 +85,11 @@ def load_groot_policies(
class GrootLocomotionController: class GrootLocomotionController:
"""GR00T lower-body locomotion controller for the Unitree G1.""" """GR00T lower-body locomotion controller for the Unitree G1."""
control_dt = CONTROL_DT # Expose for unitree_g1.py def __init__(self, policy_balance, policy_walk, robot, config):
self.policy_balance = policy_balance
def __init__(self): self.policy_walk = policy_walk
# Load policies self.robot = robot
self.policy_balance, self.policy_walk = load_groot_policies() self.config = config
self.cmd = np.array([0.0, 0.0, 0.0], dtype=np.float32) # vx, vy, theta_dot self.cmd = np.array([0.0, 0.0, 0.0], dtype=np.float32) # vx, vy, theta_dot
@@ -104,60 +109,45 @@ class GrootLocomotionController:
logger.info("GrootLocomotionController initialized") logger.info("GrootLocomotionController initialized")
def reset(self) -> None: def run_step(self):
"""Reset internal state for a new episode.""" # Get current observation
self.cmd[:] = 0.0 obs = self.robot.get_observation()
self.groot_qj_all[:] = 0.0
self.groot_dqj_all[:] = 0.0
self.groot_action[:] = 0.0
self.groot_obs_single[:] = 0.0
self.groot_obs_stacked[:] = 0.0
self.groot_height_cmd = 0.74
self.groot_orientation_cmd[:] = 0.0
self.groot_obs_history.clear()
for _ in range(6):
self.groot_obs_history.append(np.zeros(86, dtype=np.float32))
def run_step(self, action: dict, lowstate) -> dict: if not obs:
"""Run one step of the locomotion controller. return
Args: # Get command from remote controller
action: Action dict containing remote.lx/ly/rx/ry and buttons if obs["remote.buttons"][0]: # R1 - raise waist
lowstate: Robot lowstate containing motor positions/velocities and IMU
Returns:
Action dict for lower body joints (0-14)
"""
if lowstate is None:
return {}
buttons = [int(action.get(k, 0)) for k in REMOTE_BUTTONS]
if buttons[0]: # R1 - raise waist
self.groot_height_cmd += 0.001 self.groot_height_cmd += 0.001
self.groot_height_cmd = np.clip(self.groot_height_cmd, 0.50, 1.00) self.groot_height_cmd = np.clip(self.groot_height_cmd, 0.50, 1.00)
if buttons[4]: # R2 - lower waist if obs["remote.buttons"][4]: # R2 - lower waist
self.groot_height_cmd -= 0.001 self.groot_height_cmd -= 0.001
self.groot_height_cmd = np.clip(self.groot_height_cmd, 0.50, 1.00) self.groot_height_cmd = np.clip(self.groot_height_cmd, 0.50, 1.00)
lx, ly, rx, _ry = (action.get(k, 0.0) for k in REMOTE_AXES) self.cmd[0] = obs["remote.ly"] # Forward/backward
self.cmd[0] = ly # Forward/backward self.cmd[1] = obs["remote.lx"] * -1 # Left/right
self.cmd[1] = -lx # Left/right (negated) self.cmd[2] = obs["remote.rx"] * -1 # Rotation rate
self.cmd[2] = -rx # Rotation rate (negated)
# Get joint positions and velocities from lowstate # Get joint positions and velocities from flat dict
for motor in G1_29_JointIndex: for motor in G1_29_JointIndex:
name = motor.name
idx = motor.value idx = motor.value
self.groot_qj_all[idx] = lowstate.motor_state[idx].q self.groot_qj_all[idx] = obs[f"{name}.q"]
self.groot_dqj_all[idx] = lowstate.motor_state[idx].dq self.groot_dqj_all[idx] = obs[f"{name}.dq"]
# Adapt observation for g1_23dof
for idx in MISSING_JOINTS:
self.groot_qj_all[idx] = 0.0
self.groot_dqj_all[idx] = 0.0
# Scale joint positions and velocities # Scale joint positions and velocities
qj_obs = self.groot_qj_all.copy() qj_obs = self.groot_qj_all.copy()
dqj_obs = self.groot_dqj_all.copy() dqj_obs = self.groot_dqj_all.copy()
# Express IMU data in gravity frame of reference # Express IMU data in gravity frame of reference
quat = lowstate.imu_state.quaternion quat = [obs["imu.quat.w"], obs["imu.quat.x"], obs["imu.quat.y"], obs["imu.quat.z"]]
ang_vel = np.array(lowstate.imu_state.gyroscope, dtype=np.float32) ang_vel = np.array([obs["imu.gyro.x"], obs["imu.gyro.y"], obs["imu.gyro.z"]], dtype=np.float32)
gravity_orientation = get_gravity_orientation(quat) gravity_orientation = self.robot.get_gravity_orientation(quat)
# Scale joint positions and velocities before policy inference # Scale joint positions and velocities before policy inference
qj_obs = (qj_obs - GROOT_DEFAULT_ANGLES) * DOF_POS_SCALE qj_obs = (qj_obs - GROOT_DEFAULT_ANGLES) * DOF_POS_SCALE
@@ -196,10 +186,73 @@ class GrootLocomotionController:
# Transform action back to target joint positions # Transform action back to target joint positions
target_dof_pos_15 = GROOT_DEFAULT_ANGLES[:15] + self.groot_action * ACTION_SCALE target_dof_pos_15 = GROOT_DEFAULT_ANGLES[:15] + self.groot_action * ACTION_SCALE
# Build action dict # Build action dict (only first 15 joints for GR00T)
action_dict = {} action_dict = {}
for i in range(15): for i in range(15):
motor_name = G1_29_JointIndex(i).name motor_name = G1_29_JointIndex(i).name
action_dict[f"{motor_name}.q"] = float(target_dof_pos_15[i]) action_dict[f"{motor_name}.q"] = float(target_dof_pos_15[i])
return action_dict # Zero out missing joints for g1_23dof
for joint_idx in MISSING_JOINTS:
motor_name = G1_29_JointIndex(joint_idx).name
action_dict[f"{motor_name}.q"] = 0.0
# Send action to robot
self.robot.send_action(action_dict)
def run(repo_id: str = DEFAULT_GROOT_REPO_ID) -> None:
"""Main function to run the GR00T locomotion controller.
Args:
repo_id: Hugging Face Hub repository ID for GR00T policies.
"""
# Load policies
policy_balance, policy_walk = load_groot_policies(repo_id=repo_id)
# Initialize robot
config = UnitreeG1Config()
robot = UnitreeG1(config)
robot.connect()
# Initialize gr00T locomotion controller
groot_controller = GrootLocomotionController(
policy_balance=policy_balance,
policy_walk=policy_walk,
robot=robot,
config=config,
)
try:
robot.reset(CONTROL_DT, GROOT_DEFAULT_ANGLES)
logger.info("Use joystick: LY=fwd/back, LX=left/right, RX=rotate, R1=raise waist, R2=lower waist")
logger.info("Press Ctrl+C to stop")
# Run step
while not robot._shutdown_event.is_set():
start_time = time.time()
groot_controller.run_step()
elapsed = time.time() - start_time
sleep_time = max(0, CONTROL_DT - elapsed)
time.sleep(sleep_time)
except KeyboardInterrupt:
logger.info("Stopping locomotion...")
finally:
if robot.is_connected:
robot.disconnect()
logger.info("Done!")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="GR00T Locomotion Controller for Unitree G1")
parser.add_argument(
"--repo-id",
type=str,
default=DEFAULT_GROOT_REPO_ID,
help=f"Hugging Face Hub repo ID for GR00T policies (default: {DEFAULT_GROOT_REPO_ID})",
)
args = parser.parse_args()
run(repo_id=args.repo_id)
@@ -14,21 +14,21 @@
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
import argparse
import json import json
import logging import logging
import time
import numpy as np import numpy as np
import onnx import onnx
import onnxruntime as ort import onnxruntime as ort
from huggingface_hub import hf_hub_download from huggingface_hub import hf_hub_download
from lerobot.robots.unitree_g1.g1_utils import ( from lerobot.robots.unitree_g1.config_unitree_g1 import UnitreeG1Config
REMOTE_AXES, from lerobot.robots.unitree_g1.g1_utils import G1_29_JointIndex
G1_29_JointArmIndex, from lerobot.robots.unitree_g1.unitree_g1 import UnitreeG1
G1_29_JointIndex,
get_gravity_orientation,
)
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
DEFAULT_ANGLES = np.zeros(29, dtype=np.float32) DEFAULT_ANGLES = np.zeros(29, dtype=np.float32)
@@ -40,13 +40,18 @@ DEFAULT_ANGLES[16] = 0.2 # Left shoulder roll
DEFAULT_ANGLES[23] = -0.2 # Right shoulder roll DEFAULT_ANGLES[23] = -0.2 # Right shoulder roll
DEFAULT_ANGLES[[18, 25]] = 0.6 # Elbow DEFAULT_ANGLES[[18, 25]] = 0.6 # Elbow
MISSING_JOINTS = []
G1_MODEL = "g1_23" # Or "g1_29"
if G1_MODEL == "g1_23":
MISSING_JOINTS = [12, 14, 20, 21, 27, 28] # Waist yaw/pitch, wrist pitch/yaw
# Control parameters # Control parameters
ACTION_SCALE = 0.25 ACTION_SCALE = 0.25
CONTROL_DT = 0.005 # 200Hz CONTROL_DT = 0.02 # 50Hz
ANG_VEL_SCALE = 0.25 ANG_VEL_SCALE = 0.25
DOF_POS_SCALE = 1.0 DOF_POS_SCALE = 1.0
DOF_VEL_SCALE = 0.05 DOF_VEL_SCALE = 0.05
GAIT_PERIOD = 0.5 GAIT_PERIOD = 1.0
DEFAULT_HOLOSOMA_REPO_ID = "nepyope/holosoma_locomotion" DEFAULT_HOLOSOMA_REPO_ID = "nepyope/holosoma_locomotion"
@@ -82,7 +87,7 @@ def load_policy(
logger.info(f"Policy loaded: {policy.get_inputs()[0].shape}{policy.get_outputs()[0].shape}") logger.info(f"Policy loaded: {policy.get_inputs()[0].shape}{policy.get_outputs()[0].shape}")
# Extract KP/KD from ONNX metadata # Extract KP/KD from ONNX metadata
model = onnx.load(policy_path, load_external_data=False) model = onnx.load(policy_path)
metadata = {prop.key: prop.value for prop in model.metadata_props} metadata = {prop.key: prop.value for prop in model.metadata_props}
if "kp" not in metadata or "kd" not in metadata: if "kp" not in metadata or "kd" not in metadata:
@@ -96,13 +101,15 @@ def load_policy(
class HolosomaLocomotionController: class HolosomaLocomotionController:
"""Holosoma lower-body locomotion controller for Unitree G1.""" """Holosoma whole-body locomotion controller for Unitree G1."""
control_dt = CONTROL_DT # Expose for unitree_g1.py def __init__(self, policy, robot, kp: np.ndarray, kd: np.ndarray):
self.policy = policy
self.robot = robot
def __init__(self): # Override robot's PD gains with policy gains
# Load policy and gains self.robot.kp = kp
self.policy, self.kp, self.kd = load_policy() self.robot.kd = kd
self.cmd = np.zeros(3, dtype=np.float32) self.cmd = np.zeros(3, dtype=np.float32)
@@ -117,55 +124,35 @@ class HolosomaLocomotionController:
self.phase_dt = 2 * np.pi / ((1.0 / CONTROL_DT) * GAIT_PERIOD) self.phase_dt = 2 * np.pi / ((1.0 / CONTROL_DT) * GAIT_PERIOD)
self.is_standing = True self.is_standing = True
logger.info("HolosomaLocomotionController initialized") def run_step(self):
# Get current observation
obs = self.robot.get_observation()
def reset(self) -> None: if not obs:
"""Reset internal state for a new episode.""" return
self.cmd[:] = 0.0
self.qj[:] = 0.0
self.dqj[:] = 0.0
self.obs[:] = 0.0
self.last_action[:] = 0.0
self.phase = np.array([[0.0, np.pi]], dtype=np.float32)
self.is_standing = True
def run_step(self, action: dict, lowstate) -> dict: # Get command from remote controller
"""Run one step of the locomotion controller. ly = obs["remote.ly"] if abs(obs["remote.ly"]) > 0.1 else 0.0
lx = obs["remote.lx"] if abs(obs["remote.lx"]) > 0.1 else 0.0
Args: rx = obs["remote.rx"] if abs(obs["remote.rx"]) > 0.1 else 0.0
action: Action dict containing remote.lx/ly/rx/ry
lowstate: Robot lowstate containing motor positions/velocities and IMU
Returns:
Action dict for lower body joints (0-14)
"""
if lowstate is None:
return {}
lx, ly, rx, _ry = (action.get(k, 0.0) for k in REMOTE_AXES)
ly = ly if abs(ly) > 0.1 else 0.0
lx = lx if abs(lx) > 0.1 else 0.0
rx = rx if abs(rx) > 0.1 else 0.0
ly = np.clip(ly, -0.3, 0.3)
lx = np.clip(lx, -0.3, 0.3)
self.cmd[:] = [ly, -lx, -rx] self.cmd[:] = [ly, -lx, -rx]
# Get joint positions and velocities from lowstate # Get joint positions and velocities
for motor in G1_29_JointIndex: for motor in G1_29_JointIndex:
name = motor.name
idx = motor.value idx = motor.value
self.qj[idx] = lowstate.motor_state[idx].q self.qj[idx] = obs[f"{name}.q"]
self.dqj[idx] = lowstate.motor_state[idx].dq self.dqj[idx] = obs[f"{name}.dq"]
# Hide arm positions from policy (show DEFAULT_ANGLES instead) # Adapt observation for g1_23dof
# This prevents policy from reacting to teleop arm movements for idx in MISSING_JOINTS:
for arm_joint in G1_29_JointArmIndex: self.qj[idx] = 0.0
self.qj[arm_joint.value] = DEFAULT_ANGLES[arm_joint.value] self.dqj[idx] = 0.0
self.dqj[arm_joint.value] = 0.0
# Express IMU data in gravity frame of reference # Express IMU data in gravity frame of reference
quat = lowstate.imu_state.quaternion quat = [obs["imu.quat.w"], obs["imu.quat.x"], obs["imu.quat.y"], obs["imu.quat.z"]]
ang_vel = np.array(lowstate.imu_state.gyroscope, dtype=np.float32) ang_vel = np.array([obs["imu.gyro.x"], obs["imu.gyro.y"], obs["imu.gyro.z"]], dtype=np.float32)
gravity = get_gravity_orientation(quat) gravity = self.robot.get_gravity_orientation(quat)
# Scale joint positions and velocities before policy inference # Scale joint positions and velocities before policy inference
qj_obs = (self.qj - DEFAULT_ANGLES) * DOF_POS_SCALE qj_obs = (self.qj - DEFAULT_ANGLES) * DOF_POS_SCALE
@@ -199,16 +186,79 @@ class HolosomaLocomotionController:
# Run policy inference # Run policy inference
ort_in = {self.policy.get_inputs()[0].name: self.obs.reshape(1, -1).astype(np.float32)} ort_in = {self.policy.get_inputs()[0].name: self.obs.reshape(1, -1).astype(np.float32)}
raw_action = self.policy.run(None, ort_in)[0].squeeze() raw_action = self.policy.run(None, ort_in)[0].squeeze()
policy_action = np.clip(raw_action, -100.0, 100.0) action = np.clip(raw_action, -100.0, 100.0)
self.last_action = policy_action.copy() self.last_action = action.copy()
# Transform action back to target joint positions # Transform action back to target joint positions
target = DEFAULT_ANGLES + policy_action * ACTION_SCALE target = DEFAULT_ANGLES + action * ACTION_SCALE
# Build action dict (first 15 joints only) # Build action dict
action_dict = {} action_dict = {}
for i in range(15): for motor in G1_29_JointIndex:
motor_name = G1_29_JointIndex(i).name action_dict[f"{motor.name}.q"] = float(target[motor.value])
action_dict[f"{motor_name}.q"] = float(target[i])
return action_dict # Zero out missing joints for g1_23dof
for joint_idx in MISSING_JOINTS:
motor_name = G1_29_JointIndex(joint_idx).name
action_dict[f"{motor_name}.q"] = 0.0
# Send action to robot
self.robot.send_action(action_dict)
def run(repo_id: str = DEFAULT_HOLOSOMA_REPO_ID, policy_type: str = "fastsac") -> None:
"""Main function to run the Holosoma locomotion controller.
Args:
repo_id: Hugging Face Hub repository ID for Holosoma policies.
policy_type: Policy type to use ('fastsac' or 'ppo').
"""
# Load policy and gains
policy, kp, kd = load_policy(repo_id=repo_id, policy_type=policy_type)
# Initialize robot
config = UnitreeG1Config()
robot = UnitreeG1(config)
robot.connect()
holosoma_controller = HolosomaLocomotionController(policy, robot, kp, kd)
try:
robot.reset(CONTROL_DT, DEFAULT_ANGLES)
logger.info("Use joystick: LY=fwd/back, LX=left/right, RX=rotate")
logger.info("Press Ctrl+C to stop")
# Run step
while not robot._shutdown_event.is_set():
start_time = time.time()
holosoma_controller.run_step()
elapsed = time.time() - start_time
sleep_time = max(0, CONTROL_DT - elapsed)
time.sleep(sleep_time)
except KeyboardInterrupt:
logger.info("Stopping locomotion...")
finally:
if robot.is_connected:
robot.disconnect()
logger.info("Done!")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Holosoma Locomotion Controller for Unitree G1")
parser.add_argument(
"--repo-id",
type=str,
default=DEFAULT_HOLOSOMA_REPO_ID,
help=f"Hugging Face Hub repo ID for Holosoma policies (default: {DEFAULT_HOLOSOMA_REPO_ID})",
)
parser.add_argument(
"--policy",
type=str,
choices=["fastsac", "ppo"],
default="fastsac",
help="Policy type to use: 'fastsac' (default) or 'ppo'",
)
args = parser.parse_args()
run(repo_id=args.repo_id, policy_type=args.policy)
+32 -59
View File
@@ -25,11 +25,11 @@ discord = "https://discord.gg/s3KuuzsPFb"
[project] [project]
name = "lerobot" name = "lerobot"
version = "0.5.1" version = "0.4.4"
description = "🤗 LeRobot: State-of-the-art Machine Learning for Real-World Robotics in Pytorch" description = "🤗 LeRobot: State-of-the-art Machine Learning for Real-World Robotics in Pytorch"
dynamic = ["readme"] dynamic = ["readme"]
license = { text = "Apache-2.0" } license = { text = "Apache-2.0" }
requires-python = ">=3.12" requires-python = ">=3.10"
authors = [ authors = [
{ name = "Rémi Cadène", email = "re.cadene@gmail.com" }, { name = "Rémi Cadène", email = "re.cadene@gmail.com" },
{ name = "Simon Alibert", email = "alibert.sim@gmail.com" }, { name = "Simon Alibert", email = "alibert.sim@gmail.com" },
@@ -50,8 +50,7 @@ classifiers = [
"Intended Audience :: Education", "Intended Audience :: Education",
"Intended Audience :: Science/Research", "Intended Audience :: Science/Research",
"License :: OSI Approved :: Apache Software License", "License :: OSI Approved :: Apache Software License",
"Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.10",
"Programming Language :: Python :: 3.13",
"Topic :: Software Development :: Build Tools", "Topic :: Software Development :: Build Tools",
"Topic :: Scientific/Engineering :: Artificial Intelligence", "Topic :: Scientific/Engineering :: Artificial Intelligence",
] ]
@@ -62,28 +61,26 @@ dependencies = [
# Hugging Face dependencies # Hugging Face dependencies
"datasets>=4.0.0,<5.0.0", "datasets>=4.0.0,<5.0.0",
"diffusers>=0.27.2,<0.36.0", "diffusers>=0.27.2,<0.36.0",
"huggingface-hub>=1.0.0,<2.0.0", "huggingface-hub[cli]>=1.0.0,<2.0.0",
"accelerate>=1.10.0,<2.0.0", "accelerate>=1.10.0,<2.0.0",
# Core dependencies # Core dependencies
"numpy>=2.0.0,<2.3.0", # NOTE: Explicitly listing numpy helps the resolver converge faster. Upper bound imposed by opencv-python-headless.
"setuptools>=71.0.0,<81.0.0", "setuptools>=71.0.0,<81.0.0",
"cmake>=3.29.0.1,<4.2.0", "cmake>=3.29.0.1,<4.2.0",
"packaging>=24.2,<26.0",
"torch>=2.2.1,<2.11.0",
"torchcodec>=0.2.1,<0.11.0; sys_platform != 'win32' and (sys_platform != 'linux' or (platform_machine != 'aarch64' and platform_machine != 'arm64' and platform_machine != 'armv7l')) and (sys_platform != 'darwin' or platform_machine != 'x86_64')",
"torchvision>=0.21.0,<0.26.0",
"einops>=0.8.0,<0.9.0", "einops>=0.8.0,<0.9.0",
"opencv-python-headless>=4.9.0,<4.13.0", "opencv-python-headless>=4.9.0,<4.13.0",
"av>=15.0.0,<16.0.0", "av>=15.0.0,<16.0.0",
"jsonlines>=4.0.0,<5.0.0", "jsonlines>=4.0.0,<5.0.0",
"pynput>=1.7.8,<1.9.0", "packaging>=24.2,<26.0",
"pynput>=1.7.7,<1.9.0",
"pyserial>=3.5,<4.0", "pyserial>=3.5,<4.0",
"wandb>=0.24.0,<0.25.0", "wandb>=0.24.0,<0.25.0",
"draccus==0.10.0", # TODO: Relax version constraint
"torch>=2.2.1,<2.11.0", # TODO: Bump dependency
"torchcodec>=0.2.1,<0.11.0; sys_platform != 'win32' and (sys_platform != 'linux' or (platform_machine != 'aarch64' and platform_machine != 'arm64' and platform_machine != 'armv7l')) and (sys_platform != 'darwin' or platform_machine != 'x86_64')", # TODO: Bump dependency
"torchvision>=0.21.0,<0.26.0", # TODO: Bump dependency
"draccus==0.10.0", # TODO: Remove ==
"gymnasium>=1.1.1,<2.0.0", "gymnasium>=1.1.1,<2.0.0",
"rerun-sdk>=0.24.0,<0.27.0", "rerun-sdk>=0.24.0,<0.27.0",
@@ -98,14 +95,10 @@ dependencies = [
# Common # Common
pygame-dep = ["pygame>=2.5.1,<2.7.0"] pygame-dep = ["pygame>=2.5.1,<2.7.0"]
placo-dep = ["placo>=0.9.6,<0.9.17"] placo-dep = ["placo>=0.9.6,<0.10.0"]
transformers-dep = ["transformers>=5.3.0,<6.0.0"] transformers-dep = ["transformers>=5.1.0,<6.0.0"]
grpcio-dep = ["grpcio==1.73.1", "protobuf>=6.31.1,<6.32.0"] grpcio-dep = ["grpcio==1.73.1", "protobuf>=6.31.1,<6.32.0"]
can-dep = ["python-can>=4.2.0,<5.0.0"] can-dep = ["python-can>=4.2.0,<5.0.0"]
peft-dep = ["peft>=0.18.0,<1.0.0"]
scipy-dep = ["scipy>=1.14.0,<2.0.0"]
qwen-vl-utils-dep = ["qwen-vl-utils>=0.0.11,<0.1.0"]
matplotlib-dep = ["matplotlib>=3.10.3,<4.0.0", "contourpy>=1.3.0,<2.0.0"] # NOTE: Explicitly listing contourpy helps the resolver converge faster.
# Motors # Motors
feetech = ["feetech-servo-sdk>=1.0.0,<2.0.0"] feetech = ["feetech-servo-sdk>=1.0.0,<2.0.0"]
@@ -119,36 +112,34 @@ gamepad = ["lerobot[pygame-dep]", "hidapi>=0.14.0,<0.15.0"]
hopejr = ["lerobot[feetech]", "lerobot[pygame-dep]"] hopejr = ["lerobot[feetech]", "lerobot[pygame-dep]"]
lekiwi = ["lerobot[feetech]", "pyzmq>=26.2.1,<28.0.0"] lekiwi = ["lerobot[feetech]", "pyzmq>=26.2.1,<28.0.0"]
unitree_g1 = [ unitree_g1 = [
"unitree-sdk2==1.0.1",
"pyzmq>=26.2.1,<28.0.0", "pyzmq>=26.2.1,<28.0.0",
"onnxruntime>=1.16.0,<2.0.0", "onnxruntime>=1.16.0,<2.0.0",
"pin>=3.0.0,<4.0.0", "pin>=3.0.0,<4.0.0",
"meshcat>=0.3.0,<0.4.0", "meshcat>=0.3.0,<0.4.0",
"lerobot[matplotlib-dep]", "matplotlib>=3.9.0,<4.0.0",
"lerobot[pygame-dep]",
"casadi>=3.6.0,<4.0.0", "casadi>=3.6.0,<4.0.0",
] ]
reachy2 = ["reachy2_sdk>=1.0.15,<1.1.0"] reachy2 = ["reachy2_sdk>=1.0.15,<1.1.0"]
kinematics = ["lerobot[placo-dep]"] kinematics = ["lerobot[placo-dep]"]
intelrealsense = [ intelrealsense = [
"pyrealsense2>=2.55.1.6486,<2.57.0 ; sys_platform != 'darwin'", "pyrealsense2>=2.55.1.6486,<2.57.0 ; sys_platform != 'darwin'",
"pyrealsense2-macosx>=2.54,<2.57.0 ; sys_platform == 'darwin'", "pyrealsense2-macosx>=2.54,<2.55.0 ; sys_platform == 'darwin'",
] ]
phone = ["hebi-py>=2.8.0,<2.12.0", "teleop>=0.1.0,<0.2.0", "fastapi<1.0", "lerobot[scipy-dep]"] phone = ["hebi-py>=2.8.0,<2.12.0", "teleop>=0.1.0,<0.2.0", "fastapi<1.0"]
# Policies # Policies
wallx = [ wallx = [
"lerobot[transformers-dep]", "lerobot[transformers-dep]",
"lerobot[peft]", "peft>=0.18.0,<1.0.0",
"lerobot[scipy-dep]", "scipy==1.15.3", # TODO: Relax version
"torchdiffeq>=0.2.4,<0.3.0", "torchdiffeq==0.2.5", # TODO: Relax version
"lerobot[qwen-vl-utils-dep]", "qwen-vl-utils==0.0.11" # TODO: Relax version
] ]
pi = ["lerobot[transformers-dep]", "lerobot[scipy-dep]"] pi = ["lerobot[transformers-dep]", "scipy==1.15.3"] # TODO: Relax scipy version
smolvla = ["lerobot[transformers-dep]", "num2words>=0.5.14,<0.6.0", "accelerate>=1.7.0,<2.0.0", "safetensors>=0.4.3,<1.0.0"] smolvla = ["lerobot[transformers-dep]", "num2words>=0.5.14,<0.6.0", "accelerate>=1.7.0,<2.0.0", "safetensors>=0.4.3,<1.0.0"]
groot = [ groot = [
"lerobot[transformers-dep]", "lerobot[transformers-dep]",
"lerobot[peft]", "peft>=0.13.0,<1.0.0",
"dm-tree>=0.1.8,<1.0.0", "dm-tree>=0.1.8,<1.0.0",
"timm>=1.0.0,<1.1.0", "timm>=1.0.0,<1.1.0",
"safetensors>=0.4.3,<1.0.0", "safetensors>=0.4.3,<1.0.0",
@@ -157,13 +148,13 @@ groot = [
"ninja>=1.11.1,<2.0.0", "ninja>=1.11.1,<2.0.0",
"flash-attn>=2.5.9,<3.0.0 ; sys_platform != 'darwin'" "flash-attn>=2.5.9,<3.0.0 ; sys_platform != 'darwin'"
] ]
sarm = ["lerobot[transformers-dep]", "faker>=33.0.0,<35.0.0", "lerobot[matplotlib-dep]", "lerobot[qwen-vl-utils-dep]"] sarm = ["lerobot[transformers-dep]", "faker>=33.0.0,<35.0.0", "matplotlib>=3.10.3,<4.0.0", "qwen-vl-utils>=0.0.11,<0.1.0"]
xvla = ["lerobot[transformers-dep]"] xvla = ["lerobot[transformers-dep]"]
hilserl = ["lerobot[transformers-dep]", "gym-hil>=0.1.13,<0.2.0", "lerobot[grpcio-dep]", "lerobot[placo-dep]"] hilserl = ["lerobot[transformers-dep]", "gym-hil>=0.1.13,<0.2.0", "lerobot[grpcio-dep]", "lerobot[placo-dep]"]
# Features # Features
async = ["lerobot[grpcio-dep]", "lerobot[matplotlib-dep]"] async = ["lerobot[grpcio-dep]", "matplotlib>=3.10.3,<4.0.0"]
peft = ["lerobot[transformers-dep]", "lerobot[peft-dep]"] peft = ["lerobot[transformers-dep]", "peft>=0.18.0,<1.0.0"]
# Development # Development
dev = ["pre-commit>=3.7.0,<5.0.0", "debugpy>=1.8.1,<1.9.0", "lerobot[grpcio-dep]", "grpcio-tools==1.73.1", "mypy>=1.19.1"] dev = ["pre-commit>=3.7.0,<5.0.0", "debugpy>=1.8.1,<1.9.0", "lerobot[grpcio-dep]", "grpcio-tools==1.73.1", "mypy>=1.19.1"]
@@ -171,27 +162,13 @@ test = ["pytest>=8.1.0,<9.0.0", "pytest-timeout>=2.4.0,<3.0.0", "pytest-cov>=5.0
video_benchmark = ["scikit-image>=0.23.2,<0.26.0", "pandas>=2.2.2,<2.4.0"] video_benchmark = ["scikit-image>=0.23.2,<0.26.0", "pandas>=2.2.2,<2.4.0"]
# Simulation # Simulation
# NOTE: Explicitly listing scipy helps flatten the dependecy tree. aloha = ["gym-aloha>=0.1.2,<0.2.0"]
aloha = ["gym-aloha>=0.1.2,<0.2.0", "lerobot[scipy-dep]"]
pusht = ["gym-pusht>=0.1.5,<0.2.0", "pymunk>=6.6.0,<7.0.0"] # TODO: Fix pymunk version in gym-pusht instead pusht = ["gym-pusht>=0.1.5,<0.2.0", "pymunk>=6.6.0,<7.0.0"] # TODO: Fix pymunk version in gym-pusht instead
libero = ["lerobot[transformers-dep]", "hf-libero>=0.1.3,<0.2.0; sys_platform == 'linux'", "lerobot[scipy-dep]"] libero = ["lerobot[transformers-dep]", "hf-libero>=0.1.3,<0.2.0"]
libero_plus = [ metaworld = ["metaworld==3.0.0"]
"lerobot[transformers-dep]",
"libero @ git+https://github.com/sylvestf/LIBERO-plus.git@main ; sys_platform == 'linux'",
"lerobot[scipy-dep]",
]
robomme = [
"robomme @ git+https://github.com/RoboMME/robomme_benchmark.git@main ; sys_platform == 'linux'",
]
metaworld = ["metaworld==3.0.0", "lerobot[scipy-dep]"]
# All # All
all = [ all = [
# NOTE(resolver hint): scipy is pulled in transitively via lerobot[scipy-dep] through
# multiple extras (aloha, metaworld, pi, wallx, phone). Listing it explicitly
# helps pip's resolver converge by constraining scipy early, before it encounters
# the loose scipy requirements from transitive deps like dm-control and metaworld.
"scipy>=1.14.0,<2.0.0",
"lerobot[dynamixel]", "lerobot[dynamixel]",
"lerobot[gamepad]", "lerobot[gamepad]",
"lerobot[hopejr]", "lerobot[hopejr]",
@@ -212,11 +189,10 @@ all = [
"lerobot[aloha]", "lerobot[aloha]",
"lerobot[pusht]", "lerobot[pusht]",
"lerobot[phone]", "lerobot[phone]",
"lerobot[libero]; sys_platform == 'linux'", "lerobot[libero]",
"lerobot[metaworld]", "lerobot[metaworld]",
"lerobot[sarm]", "lerobot[sarm]",
"lerobot[peft]", "lerobot[peft]",
# "lerobot[unitree_g1]", TODO: Unitree requires specific installation instructions for unitree_sdk2
] ]
[project.scripts] [project.scripts]
@@ -238,14 +214,11 @@ lerobot-edit-dataset="lerobot.scripts.lerobot_edit_dataset:main"
lerobot-setup-can="lerobot.scripts.lerobot_setup_can:main" lerobot-setup-can="lerobot.scripts.lerobot_setup_can:main"
# ---------------- Tool Configurations ---------------- # ---------------- Tool Configurations ----------------
[tool.setuptools.package-data]
lerobot = ["envs/*.json"]
[tool.setuptools.packages.find] [tool.setuptools.packages.find]
where = ["src"] where = ["src"]
[tool.ruff] [tool.ruff]
target-version = "py312" target-version = "py310"
line-length = 110 line-length = 110
exclude = ["tests/artifacts/**/*.safetensors", "*_pb2.py", "*_pb2_grpc.py"] exclude = ["tests/artifacts/**/*.safetensors", "*_pb2.py", "*_pb2_grpc.py"]
@@ -337,7 +310,7 @@ default.extend-ignore-identifiers-re = [
# Uncomment [tool.mypy] first, then uncomment individual module overrides as they get proper type annotations # Uncomment [tool.mypy] first, then uncomment individual module overrides as they get proper type annotations
[tool.mypy] [tool.mypy]
python_version = "3.12" python_version = "3.10"
ignore_missing_imports = true ignore_missing_imports = true
follow_imports = "skip" follow_imports = "skip"
# warn_return_any = true # warn_return_any = true
+276 -175
View File
@@ -1,73 +1,76 @@
# #
# This file is autogenerated by pip-compile with Python 3.12 # This file is autogenerated by pip-compile with Python 3.10
# by the following command: # by the following command:
# #
# pip-compile --output-file=requirements-macos.txt requirements.in # pip-compile --output-file=requirements-macos.txt requirements.in
# #
-e .[all] -e .[all]
# via -[all] # via -[all]
absl-py==2.4.0 absl-py==2.3.1
# via # via
# dm-control # dm-control
# dm-env # dm-env
# dm-tree # dm-tree
# labmaze # labmaze
# mujoco # mujoco
accelerate==1.13.0 # tensorboard
accelerate==1.11.0
# via # via
# lerobot # lerobot
# peft # peft
aiohappyeyeballs==2.6.1 aiohappyeyeballs==2.6.1
# via aiohttp # via aiohttp
aiohttp==3.13.3 aiohttp==3.13.1
# via fsspec # via fsspec
aiosignal==1.4.0 aiosignal==1.4.0
# via aiohttp # via aiohttp
annotated-doc==0.0.4
# via
# fastapi
# typer
annotated-types==0.7.0 annotated-types==0.7.0
# via pydantic # via pydantic
anyio==4.12.1 antlr4-python3-runtime==4.9.3
# via
# hydra-core
# omegaconf
anyio==4.11.0
# via # via
# httpx
# starlette # starlette
# watchfiles # watchfiles
asttokens==3.0.1 asttokens==3.0.0
# via stack-data # via stack-data
async-timeout==5.0.1
# via aiohttp
attrs==25.4.0 attrs==25.4.0
# via # via
# aiohttp # aiohttp
# dm-tree # dm-tree
# jsonlines # jsonlines
# jsonschema
# referencing
# rerun-sdk # rerun-sdk
av==15.1.0 av==15.1.0
# via lerobot
bddl==1.0.1
# via libero
certifi==2025.10.5
# via # via
# lerobot
# qwen-vl-utils
certifi==2026.2.25
# via
# httpcore
# httpx
# requests # requests
# sentry-sdk # sentry-sdk
cffi==2.0.0 cffi==2.0.0
# via pymunk # via pymunk
cfgv==3.5.0 cfgv==3.4.0
# via pre-commit # via pre-commit
charset-normalizer==3.4.5 charset-normalizer==3.4.4
# via requests # via requests
click==8.3.1 click==8.3.0
# via # via
# typer
# uvicorn # uvicorn
# wandb # wandb
cloudpickle==3.1.2 cloudpickle==3.1.1
# via gymnasium # via
cmake==4.1.3 # gymnasium
# libero
cmake==4.1.0
# via lerobot # via lerobot
cmeel==0.59.0 cmeel==0.57.3
# via # via
# cmeel-assimp # cmeel-assimp
# cmeel-boost # cmeel-boost
@@ -105,17 +108,15 @@ cmeel-zlib==1.3.1
# via cmeel-assimp # via cmeel-assimp
coal-library==3.0.1 coal-library==3.0.1
# via pin # via pin
contourpy==1.3.3 contourpy==1.3.2
# via # via matplotlib
# lerobot coverage[toml]==7.11.0
# matplotlib
coverage[toml]==7.13.4
# via pytest-cov # via pytest-cov
cycler==0.12.1 cycler==0.12.1
# via matplotlib # via matplotlib
datasets==4.6.1 datasets==4.1.1
# via lerobot # via lerobot
debugpy==1.8.20 debugpy==1.8.17
# via lerobot # via lerobot
decorator==5.2.1 decorator==5.2.1
# via ipython # via ipython
@@ -129,7 +130,7 @@ dill==0.4.0
# multiprocess # multiprocess
distlib==0.4.0 distlib==0.4.0
# via virtualenv # via virtualenv
dm-control==1.0.37 dm-control==1.0.34
# via gym-aloha # via gym-aloha
dm-env==1.6 dm-env==1.6
# via dm-control # via dm-control
@@ -137,55 +138,69 @@ dm-tree==0.1.9
# via # via
# dm-control # dm-control
# dm-env # dm-env
# lerobot
docopt==0.6.2 docopt==0.6.2
# via num2words # via num2words
draccus==0.10.0 draccus==0.10.0
# via lerobot # via lerobot
dynamixel-sdk==3.8.4 dynamixel-sdk==3.8.4
# via lerobot # via lerobot
easydict==1.13
# via libero
egl-probe @ git+https://github.com/huggingface/egl_probe.git
# via
# libero
# robomimic
eigenpy==3.10.3 eigenpy==3.10.3
# via coal-library # via coal-library
einops==0.8.2 einops==0.8.1
# via lerobot
eiquadprog==1.2.9
# via placo
etils[epath,epy]==1.14.0
# via mujoco
executing==2.2.1
# via stack-data
faker==34.0.2
# via lerobot
farama-notifications==0.0.4
# via gymnasium
fastapi==0.135.1
# via # via
# lerobot # lerobot
# teleop # libero
eiquadprog==1.2.9
# via placo
etils[epath,epy]==1.13.0
# via mujoco
exceptiongroup==1.3.0
# via
# anyio
# ipython
# pytest
executing==2.2.1
# via stack-data
farama-notifications==0.0.4
# via gymnasium
fastapi==0.119.1
# via teleop
fastjsonschema==2.21.2
# via nbformat
feetech-servo-sdk==1.0.0 feetech-servo-sdk==1.0.0
# via lerobot # via lerobot
filelock==3.25.0 filelock==3.20.0
# via # via
# datasets # datasets
# diffusers # diffusers
# huggingface-hub # huggingface-hub
# python-discovery
# torch # torch
# transformers
# virtualenv # virtualenv
fonttools==4.61.1 fonttools==4.60.1
# via matplotlib # via matplotlib
frozenlist==1.8.0 frozenlist==1.8.0
# via # via
# aiohttp # aiohttp
# aiosignal # aiosignal
fsspec[http]==2026.2.0 fsspec[http]==2025.9.0
# via # via
# datasets # datasets
# etils # etils
# huggingface-hub # huggingface-hub
# torch # torch
future==1.0.0
# via libero
gitdb==4.0.12 gitdb==4.0.12
# via gitpython # via gitpython
gitpython==3.1.46 gitpython==3.1.45
# via wandb # via wandb
glfw==2.10.0 glfw==2.10.0
# via # via
@@ -197,6 +212,7 @@ grpcio==1.73.1
# lerobot # lerobot
# reachy2-sdk # reachy2-sdk
# reachy2-sdk-api # reachy2-sdk-api
# tensorboard
grpcio-tools==1.73.1 grpcio-tools==1.73.1
# via # via
# lerobot # lerobot
@@ -207,67 +223,71 @@ gym-hil==0.1.13
# via lerobot # via lerobot
gym-pusht==0.1.6 gym-pusht==0.1.6
# via lerobot # via lerobot
gymnasium==1.2.3 gymnasium==1.2.1
# via # via
# gym-aloha # gym-aloha
# gym-hil # gym-hil
# gym-pusht # gym-pusht
# lerobot # lerobot
# libero
# metaworld # metaworld
h11==0.16.0 h11==0.16.0
# via # via uvicorn
# httpcore h5py==3.15.1
# uvicorn # via robomimic
hebi-py==2.11.0 hebi-py==2.11.0
# via lerobot # via lerobot
hf-xet==1.3.2 hf-transfer==0.1.9
# via huggingface-hub
hf-xet==1.1.10
# via huggingface-hub # via huggingface-hub
hidapi==0.14.0.post4 hidapi==0.14.0.post4
# via # via
# gym-hil # gym-hil
# lerobot # lerobot
httpcore==1.0.9
# via httpx
httptools==0.7.1 httptools==0.7.1
# via uvicorn # via uvicorn
httpx==0.28.1 huggingface-hub[cli,hf-transfer]==0.35.3
# via
# datasets
# huggingface-hub
huggingface-hub==1.6.0
# via # via
# accelerate # accelerate
# datasets # datasets
# diffusers # diffusers
# lerobot # lerobot
# peft # peft
# timm
# tokenizers # tokenizers
# transformers # transformers
identify==2.6.17 hydra-core==1.3.2
# via libero
identify==2.6.15
# via pre-commit # via pre-commit
idna==3.11 idna==3.11
# via # via
# anyio # anyio
# httpx
# requests # requests
# yarl # yarl
imageio[ffmpeg]==2.37.2 imageio[ffmpeg]==2.37.0
# via # via
# gym-aloha # gym-aloha
# gym-hil # gym-hil
# lerobot # lerobot
# metaworld # metaworld
# robomimic
# scikit-image # scikit-image
imageio-ffmpeg==0.6.0 imageio-ffmpeg==0.6.0
# via imageio # via
importlib-metadata==8.7.1 # imageio
# robomimic
importlib-metadata==8.7.0
# via diffusers # via diffusers
importlib-resources==6.5.2
# via etils
iniconfig==2.3.0 iniconfig==2.3.0
# via pytest # via pytest
ipython==9.11.0 inquirerpy==0.3.4
# via huggingface-hub
ipython==8.37.0
# via meshcat # via meshcat
ipython-pygments-lexers==1.1.1
# via ipython
ischedule==1.2.7 ischedule==1.2.7
# via placo # via placo
jedi==0.19.2 jedi==0.19.2
@@ -276,24 +296,44 @@ jinja2==3.1.6
# via torch # via torch
jsonlines==4.0.0 jsonlines==4.0.0
# via lerobot # via lerobot
jsonschema==4.25.1
# via nbformat
jsonschema-specifications==2025.9.1
# via jsonschema
jupyter-core==5.9.1
# via nbformat
jupytext==1.18.1
# via bddl
kiwisolver==1.4.9 kiwisolver==1.4.9
# via matplotlib # via matplotlib
labmaze==1.0.6 labmaze==1.0.6
# via dm-control # via dm-control
lazy-loader==0.5 lazy-loader==0.4
# via scikit-image # via scikit-image
librt==0.8.1 libero @ git+https://github.com/huggingface/lerobot-libero.git@main
# via mypy # via lerobot
llvmlite==0.45.1
# via numba
lxml==6.0.2 lxml==6.0.2
# via dm-control # via dm-control
markdown==3.9
# via tensorboard
markdown-it-py==4.0.0 markdown-it-py==4.0.0
# via rich # via
# jupytext
# mdit-py-plugins
markupsafe==3.0.3 markupsafe==3.0.3
# via jinja2 # via
matplotlib==3.10.8 # jinja2
# via lerobot # werkzeug
matplotlib==3.10.7
# via
# lerobot
# libero
matplotlib-inline==0.2.1 matplotlib-inline==0.2.1
# via ipython # via ipython
mdit-py-plugins==0.5.0
# via jupytext
mdurl==0.1.2 mdurl==0.1.2
# via markdown-it-py # via markdown-it-py
mergedeep==1.3.4 mergedeep==1.3.4
@@ -306,35 +346,41 @@ mock-serial==0.0.1
# via lerobot # via lerobot
mpmath==1.3.0 mpmath==1.3.0
# via sympy # via sympy
mujoco==3.5.0 mujoco==3.3.7
# via # via
# dm-control # dm-control
# gym-aloha # gym-aloha
# gym-hil # gym-hil
# libero
# metaworld # metaworld
multidict==6.7.1 # robosuite
multidict==6.7.0
# via # via
# aiohttp # aiohttp
# yarl # yarl
multiprocess==0.70.18 multiprocess==0.70.16
# via datasets # via datasets
mypy==1.19.1
# via lerobot
mypy-extensions==1.1.0 mypy-extensions==1.1.0
# via typing-inspect
nbformat==5.10.4
# via jupytext
networkx==3.4.2
# via # via
# mypy # bddl
# typing-inspect
networkx==3.6.1
# via
# scikit-image # scikit-image
# torch # torch
nodeenv==1.10.0 ninja==1.13.0
# via lerobot
nodeenv==1.9.1
# via pre-commit # via pre-commit
num2words==0.5.14 num2words==0.5.14
# via lerobot # via lerobot
numba==0.62.1
# via robosuite
numpy==2.2.6 numpy==2.2.6
# via # via
# accelerate # accelerate
# bddl
# cmeel-boost # cmeel-boost
# contourpy # contourpy
# datasets # datasets
@@ -343,14 +389,16 @@ numpy==2.2.6
# dm-env # dm-env
# dm-tree # dm-tree
# gymnasium # gymnasium
# h5py
# hebi-py # hebi-py
# imageio # imageio
# labmaze # labmaze
# lerobot # libero
# matplotlib # matplotlib
# meshcat # meshcat
# metaworld # metaworld
# mujoco # mujoco
# numba
# opencv-python # opencv-python
# opencv-python-headless # opencv-python-headless
# pandas # pandas
@@ -358,18 +406,26 @@ numpy==2.2.6
# pyquaternion # pyquaternion
# reachy2-sdk # reachy2-sdk
# rerun-sdk # rerun-sdk
# robomimic
# robosuite
# scikit-image # scikit-image
# scipy # scipy
# shapely # shapely
# teleop # teleop
# tensorboard
# tensorboardx
# tifffile # tifffile
# torchvision # torchvision
# transformers # transformers
# transforms3d # transforms3d
opencv-python==4.13.0.92 omegaconf==2.3.0
# via hydra-core
opencv-python==4.12.0.88
# via # via
# gym-pusht # gym-pusht
# libero
# reachy2-sdk # reachy2-sdk
# robosuite
opencv-python-headless==4.12.0.88 opencv-python-headless==4.12.0.88
# via lerobot # via lerobot
orderly-set==5.5.0 orderly-set==5.5.0
@@ -379,87 +435,97 @@ packaging==25.0
# accelerate # accelerate
# datasets # datasets
# huggingface-hub # huggingface-hub
# hydra-core
# jupytext
# lazy-loader # lazy-loader
# lerobot # lerobot
# matplotlib # matplotlib
# peft # peft
# pytest # pytest
# qwen-vl-utils
# reachy2-sdk # reachy2-sdk
# scikit-image # scikit-image
# tensorboard
# tensorboardx
# transformers # transformers
# wandb # wandb
pandas==2.3.3 pandas==2.3.3
# via # via
# datasets # datasets
# lerobot # lerobot
parso==0.8.6 parso==0.8.5
# via jedi # via jedi
pathspec==1.0.4 peft==0.17.1
# via mypy
peft==0.18.1
# via lerobot # via lerobot
pexpect==4.9.0 pexpect==4.9.0
# via ipython # via ipython
pillow==12.1.1 pfzy==0.3.4
# via inquirerpy
pillow==12.0.0
# via # via
# diffusers # diffusers
# imageio # imageio
# lerobot
# matplotlib # matplotlib
# meshcat # meshcat
# qwen-vl-utils
# rerun-sdk # rerun-sdk
# robosuite
# scikit-image # scikit-image
# tensorboard
# torchvision # torchvision
pin==3.4.0 pin==3.4.0
# via placo # via placo
placo==0.9.16 placo==0.9.14
# via lerobot # via lerobot
platformdirs==4.9.4 platformdirs==4.5.0
# via # via
# python-discovery # jupyter-core
# virtualenv # virtualenv
# wandb # wandb
pluggy==1.6.0 pluggy==1.6.0
# via # via
# pytest # pytest
# pytest-cov # pytest-cov
pre-commit==4.5.1 pre-commit==4.3.0
# via lerobot # via lerobot
prompt-toolkit==3.0.52 prompt-toolkit==3.0.52
# via ipython # via
# inquirerpy
# ipython
propcache==0.4.1 propcache==0.4.1
# via # via
# aiohttp # aiohttp
# yarl # yarl
protobuf==6.31.1 protobuf==6.31.0
# via # via
# dm-control # dm-control
# grpcio-tools # grpcio-tools
# lerobot # lerobot
# reachy2-sdk # reachy2-sdk
# reachy2-sdk-api # reachy2-sdk-api
# tensorboard
# tensorboardx
# wandb # wandb
psutil==7.2.2 psutil==7.1.1
# via # via
# accelerate # accelerate
# imageio # imageio
# peft # peft
# robomimic
ptyprocess==0.7.0 ptyprocess==0.7.0
# via pexpect # via pexpect
pure-eval==0.2.3 pure-eval==0.2.3
# via stack-data # via stack-data
pyarrow==23.0.1 pyarrow==21.0.0
# via # via
# datasets # datasets
# rerun-sdk # rerun-sdk
pycparser==3.0 pycparser==2.23
# via cffi # via cffi
pydantic==2.12.5 pydantic==2.12.3
# via # via
# fastapi # fastapi
# wandb # wandb
pydantic-core==2.41.5 pydantic-core==2.41.4
# via pydantic # via pydantic
pygame==2.6.1 pygame==2.6.1
# via # via
@@ -469,35 +535,33 @@ pygame==2.6.1
pygments==2.19.2 pygments==2.19.2
# via # via
# ipython # ipython
# ipython-pygments-lexers
# pytest # pytest
# rich
pymunk==6.11.1 pymunk==6.11.1
# via # via
# gym-pusht # gym-pusht
# lerobot # lerobot
pyngrok==7.5.1 pyngrok==7.4.1
# via meshcat # via meshcat
pynput==1.8.1 pynput==1.8.1
# via # via
# gym-hil # gym-hil
# lerobot # lerobot
pyobjc-core==12.1 pyobjc-core==12.0
# via # via
# pyobjc-framework-applicationservices # pyobjc-framework-applicationservices
# pyobjc-framework-cocoa # pyobjc-framework-cocoa
# pyobjc-framework-coretext # pyobjc-framework-coretext
# pyobjc-framework-quartz # pyobjc-framework-quartz
pyobjc-framework-applicationservices==12.1 pyobjc-framework-applicationservices==12.0
# via pynput # via pynput
pyobjc-framework-cocoa==12.1 pyobjc-framework-cocoa==12.0
# via # via
# pyobjc-framework-applicationservices # pyobjc-framework-applicationservices
# pyobjc-framework-coretext # pyobjc-framework-coretext
# pyobjc-framework-quartz # pyobjc-framework-quartz
pyobjc-framework-coretext==12.1 pyobjc-framework-coretext==12.0
# via pyobjc-framework-applicationservices # via pyobjc-framework-applicationservices
pyobjc-framework-quartz==12.1 pyobjc-framework-quartz==12.0
# via # via
# pynput # pynput
# pyobjc-framework-applicationservices # pyobjc-framework-applicationservices
@@ -506,13 +570,13 @@ pyopengl==3.1.10
# via # via
# dm-control # dm-control
# mujoco # mujoco
pyparsing==3.3.2 pyparsing==3.2.5
# via # via
# dm-control # dm-control
# matplotlib # matplotlib
pyquaternion==0.9.9 pyquaternion==0.9.9
# via reachy2-sdk # via reachy2-sdk
pyrealsense2-macosx==2.56.5 pyrealsense2-macosx==2.54.2
# via lerobot # via lerobot
pyserial==3.5 pyserial==3.5
# via # via
@@ -521,6 +585,7 @@ pyserial==3.5
# lerobot # lerobot
pytest==8.4.2 pytest==8.4.2
# via # via
# bddl
# lerobot # lerobot
# pytest-cov # pytest-cov
# pytest-timeout # pytest-timeout
@@ -531,14 +596,11 @@ pytest-timeout==2.4.0
# via lerobot # via lerobot
python-dateutil==2.9.0.post0 python-dateutil==2.9.0.post0
# via # via
# faker
# matplotlib # matplotlib
# pandas # pandas
python-discovery==1.1.1 python-dotenv==1.1.1
# via virtualenv
python-dotenv==1.2.2
# via uvicorn # via uvicorn
pytz==2026.1.post1 pytz==2025.2
# via pandas # via pandas
pyyaml==6.0.3 pyyaml==6.0.3
# via # via
@@ -547,10 +609,13 @@ pyyaml==6.0.3
# draccus # draccus
# hebi-py # hebi-py
# huggingface-hub # huggingface-hub
# jupytext
# omegaconf
# peft # peft
# pre-commit # pre-commit
# pyngrok # pyngrok
# pyyaml-include # pyyaml-include
# timm
# transformers # transformers
# uvicorn # uvicorn
# wandb # wandb
@@ -560,13 +625,15 @@ pyzmq==27.1.0
# via # via
# lerobot # lerobot
# meshcat # meshcat
qwen-vl-utils==0.0.14 reachy2-sdk==1.0.14
# via lerobot
reachy2-sdk==1.0.15
# via lerobot # via lerobot
reachy2-sdk-api==1.0.21 reachy2-sdk-api==1.0.21
# via reachy2-sdk # via reachy2-sdk
regex==2026.2.28 referencing==0.37.0
# via
# jsonschema
# jsonschema-specifications
regex==2025.10.23
# via # via
# diffusers # diffusers
# transformers # transformers
@@ -575,150 +642,184 @@ requests==2.32.5
# datasets # datasets
# diffusers # diffusers
# dm-control # dm-control
# qwen-vl-utils # huggingface-hub
# teleop # teleop
# transformers
# wandb # wandb
rerun-sdk==0.26.2 rerun-sdk==0.26.1
# via lerobot # via lerobot
rhoban-cmeel-jsoncpp==1.9.4.9 rhoban-cmeel-jsoncpp==1.9.4.9
# via placo # via placo
rich==14.3.3 robomimic==0.2.0
# via typer # via libero
safetensors==0.7.0 robosuite==1.4.0
# via libero
rpds-py==0.28.0
# via
# jsonschema
# referencing
safetensors==0.6.2
# via # via
# accelerate # accelerate
# diffusers # diffusers
# lerobot # lerobot
# peft # peft
# timm
# transformers # transformers
scikit-image==0.25.2 scikit-image==0.25.2
# via # via
# gym-pusht # gym-pusht
# lerobot # lerobot
scipy==1.17.1 scipy==1.15.3
# via # via
# dm-control # dm-control
# lerobot
# metaworld # metaworld
# robosuite
# scikit-image # scikit-image
# torchdiffeq sentry-sdk==2.42.1
sentry-sdk==2.54.0
# via wandb # via wandb
shapely==2.1.2 shapely==2.1.2
# via gym-pusht # via gym-pusht
shellingham==1.5.4
# via typer
six==1.17.0 six==1.17.0
# via # via
# pynput # pynput
# python-dateutil # python-dateutil
smmap==5.0.3 smmap==5.0.2
# via gitdb # via gitdb
sniffio==1.3.1
# via anyio
stack-data==0.6.3 stack-data==0.6.3
# via ipython # via ipython
starlette==0.52.1 starlette==0.48.0
# via fastapi # via fastapi
sympy==1.14.0 sympy==1.14.0
# via torch # via torch
teleop==0.1.4 teleop==0.1.2
# via lerobot # via lerobot
termcolor==3.3.0 tensorboard==2.20.0
# via lerobot # via robomimic
tifffile==2026.3.3 tensorboard-data-server==0.7.2
# via tensorboard
tensorboardx==2.6.4
# via robomimic
termcolor==3.1.0
# via
# lerobot
# robomimic
thop==0.1.1.post2209072238
# via libero
tifffile==2025.5.10
# via scikit-image # via scikit-image
tokenizers==0.22.2 timm==1.0.20
# via lerobot
tokenizers==0.22.1
# via transformers # via transformers
toml==0.10.2 toml==0.10.2
# via draccus # via draccus
torch==2.10.0 tomli==2.3.0
# via
# cmeel
# coverage
# jupytext
# pytest
torch==2.7.1
# via # via
# accelerate # accelerate
# lerobot # lerobot
# peft # peft
# torchdiffeq # robomimic
# thop
# timm
# torchvision # torchvision
torchcodec==0.10.0 torchcodec==0.5
# via lerobot # via lerobot
torchdiffeq==0.2.5 torchvision==0.22.1
# via lerobot # via
torchvision==0.25.0 # lerobot
# via lerobot # robomimic
tornado==6.5.4 # timm
tornado==6.5.2
# via meshcat # via meshcat
tqdm==4.67.3 tqdm==4.67.1
# via # via
# datasets # datasets
# dm-control # dm-control
# huggingface-hub # huggingface-hub
# peft # peft
# robomimic
# transformers # transformers
traitlets==5.14.3 traitlets==5.14.3
# via # via
# ipython # ipython
# jupyter-core
# matplotlib-inline # matplotlib-inline
transformers==5.3.0 # nbformat
transformers==4.57.1
# via # via
# lerobot # lerobot
# libero
# peft # peft
transforms3d==0.4.2 transforms3d==0.4.2
# via teleop # via teleop
typer==0.24.1
# via
# huggingface-hub
# transformers
typing-extensions==4.15.0 typing-extensions==4.15.0
# via # via
# aiosignal # aiosignal
# anyio # anyio
# etils # etils
# faker # exceptiongroup
# fastapi # fastapi
# gymnasium # gymnasium
# huggingface-hub # huggingface-hub
# mypy # ipython
# multidict
# pydantic # pydantic
# pydantic-core # pydantic-core
# referencing
# rerun-sdk # rerun-sdk
# starlette # starlette
# torch # torch
# typing-inspect # typing-inspect
# typing-inspection # typing-inspection
# uvicorn
# virtualenv
# wandb # wandb
typing-inspect==0.9.0 typing-inspect==0.9.0
# via draccus # via draccus
typing-inspection==0.4.2 typing-inspection==0.4.2
# via # via pydantic
# fastapi tzdata==2025.2
# pydantic
tzdata==2025.3
# via pandas # via pandas
u-msgpack-python==2.8.0 u-msgpack-python==2.8.0
# via meshcat # via meshcat
urllib3==2.6.3 urllib3==2.5.0
# via # via
# requests # requests
# sentry-sdk # sentry-sdk
uvicorn[standard]==0.41.0 uvicorn[standard]==0.38.0
# via teleop # via teleop
uvloop==0.22.1 uvloop==0.22.1
# via uvicorn # via uvicorn
virtualenv==21.1.0 virtualenv==20.35.3
# via pre-commit # via pre-commit
wandb==0.24.2 wandb==0.21.4
# via lerobot # via
# lerobot
# libero
watchfiles==1.1.1 watchfiles==1.1.1
# via uvicorn # via uvicorn
wcwidth==0.6.0 wcwidth==0.2.14
# via prompt-toolkit # via prompt-toolkit
websocket-client==1.9.0 websocket-client==1.9.0
# via teleop # via teleop
websockets==16.0 websockets==15.0.1
# via uvicorn # via uvicorn
wrapt==2.1.2 werkzeug==3.1.3
# via tensorboard
wrapt==2.0.0
# via dm-tree # via dm-tree
xxhash==3.6.0 xxhash==3.6.0
# via datasets # via datasets
yarl==1.23.0 yarl==1.22.0
# via aiohttp # via aiohttp
zipp==3.23.0 zipp==3.23.0
# via # via
+187 -208
View File
@@ -1,12 +1,12 @@
# #
# This file is autogenerated by pip-compile with Python 3.12 # This file is autogenerated by pip-compile with Python 3.10
# by the following command: # by the following command:
# #
# pip-compile --output-file=requirements-ubuntu.txt requirements.in # pip-compile --output-file=requirements-ubuntu.txt requirements.in
# #
-e .[all] -e .[all]
# via -[all] # via -[all]
absl-py==2.4.0 absl-py==2.3.1
# via # via
# dm-control # dm-control
# dm-env # dm-env
@@ -14,33 +14,30 @@ absl-py==2.4.0
# labmaze # labmaze
# mujoco # mujoco
# tensorboard # tensorboard
accelerate==1.13.0 accelerate==1.11.0
# via # via
# lerobot # lerobot
# peft # peft
aiohappyeyeballs==2.6.1 aiohappyeyeballs==2.6.1
# via aiohttp # via aiohttp
aiohttp==3.13.3 aiohttp==3.13.1
# via fsspec # via fsspec
aiosignal==1.4.0 aiosignal==1.4.0
# via aiohttp # via aiohttp
annotated-doc==0.0.4
# via
# fastapi
# typer
annotated-types==0.7.0 annotated-types==0.7.0
# via pydantic # via pydantic
antlr4-python3-runtime==4.9.3 antlr4-python3-runtime==4.9.3
# via # via
# hydra-core # hydra-core
# omegaconf # omegaconf
anyio==4.12.1 anyio==4.11.0
# via # via
# httpx
# starlette # starlette
# watchfiles # watchfiles
asttokens==3.0.1 asttokens==3.0.0
# via stack-data # via stack-data
async-timeout==5.0.1
# via aiohttp
attrs==25.4.0 attrs==25.4.0
# via # via
# aiohttp # aiohttp
@@ -50,35 +47,30 @@ attrs==25.4.0
# referencing # referencing
# rerun-sdk # rerun-sdk
av==15.1.0 av==15.1.0
# via # via lerobot
# lerobot
# qwen-vl-utils
bddl==1.0.1 bddl==1.0.1
# via hf-libero # via libero
certifi==2026.2.25 certifi==2025.10.5
# via # via
# httpcore
# httpx
# requests # requests
# sentry-sdk # sentry-sdk
cffi==2.0.0 cffi==2.0.0
# via pymunk # via pymunk
cfgv==3.5.0 cfgv==3.4.0
# via pre-commit # via pre-commit
charset-normalizer==3.4.5 charset-normalizer==3.4.4
# via requests # via requests
click==8.3.1 click==8.3.0
# via # via
# typer
# uvicorn # uvicorn
# wandb # wandb
cloudpickle==3.1.2 cloudpickle==3.1.1
# via # via
# gymnasium # gymnasium
# hf-libero # libero
cmake==4.1.3 cmake==4.1.0
# via lerobot # via lerobot
cmeel==0.59.0 cmeel==0.57.3
# via # via
# cmeel-assimp # cmeel-assimp
# cmeel-boost # cmeel-boost
@@ -116,24 +108,20 @@ cmeel-zlib==1.3.1
# via cmeel-assimp # via cmeel-assimp
coal-library==3.0.1 coal-library==3.0.1
# via pin # via pin
contourpy==1.3.3 contourpy==1.3.2
# via # via matplotlib
# lerobot coverage[toml]==7.11.0
# matplotlib
coverage[toml]==7.13.4
# via pytest-cov # via pytest-cov
cuda-bindings==12.9.4
# via torch
cuda-pathfinder==1.4.1
# via cuda-bindings
cycler==0.12.1 cycler==0.12.1
# via matplotlib # via matplotlib
datasets==4.6.1 datasets==4.1.1
# via lerobot # via lerobot
debugpy==1.8.20 debugpy==1.8.17
# via lerobot # via lerobot
decorator==5.2.1 decorator==5.2.1
# via ipython # via ipython
decord==0.6.0
# via lerobot
deepdiff==8.6.1 deepdiff==8.6.1
# via lerobot # via lerobot
diffusers==0.35.2 diffusers==0.35.2
@@ -144,7 +132,7 @@ dill==0.4.0
# multiprocess # multiprocess
distlib==0.4.0 distlib==0.4.0
# via virtualenv # via virtualenv
dm-control==1.0.37 dm-control==1.0.34
# via gym-aloha # via gym-aloha
dm-env==1.6 dm-env==1.6
# via dm-control # via dm-control
@@ -152,6 +140,7 @@ dm-tree==0.1.9
# via # via
# dm-control # dm-control
# dm-env # dm-env
# lerobot
docopt==0.6.2 docopt==0.6.2
# via num2words # via num2words
draccus==0.10.0 draccus==0.10.0
@@ -159,60 +148,66 @@ draccus==0.10.0
dynamixel-sdk==3.8.4 dynamixel-sdk==3.8.4
# via lerobot # via lerobot
easydict==1.13 easydict==1.13
# via hf-libero # via libero
egl-probe==1.0.2 egl-probe @ git+https://github.com/huggingface/egl_probe.git
# via robomimic # via
# libero
# robomimic
eigenpy==3.10.3 eigenpy==3.10.3
# via coal-library # via coal-library
einops==0.8.2 einops==0.8.1
# via # via
# hf-libero # flash-attn
# lerobot # lerobot
# libero
eiquadprog==1.2.9 eiquadprog==1.2.9
# via placo # via placo
etils[epath,epy]==1.14.0 etils[epath,epy]==1.13.0
# via mujoco # via mujoco
evdev==1.9.3 evdev==1.9.2
# via pynput # via pynput
exceptiongroup==1.3.0
# via
# anyio
# ipython
# pytest
executing==2.2.1 executing==2.2.1
# via stack-data # via stack-data
faker==34.0.2
# via lerobot
farama-notifications==0.0.4 farama-notifications==0.0.4
# via gymnasium # via gymnasium
fastapi==0.135.1 fastapi==0.119.1
# via # via teleop
# lerobot
# teleop
fastjsonschema==2.21.2 fastjsonschema==2.21.2
# via nbformat # via nbformat
feetech-servo-sdk==1.0.0 feetech-servo-sdk==1.0.0
# via lerobot # via lerobot
filelock==3.25.0 filelock==3.20.0
# via # via
# datasets # datasets
# diffusers # diffusers
# huggingface-hub # huggingface-hub
# python-discovery
# torch # torch
# transformers
# virtualenv # virtualenv
fonttools==4.61.1 flash-attn==2.8.3
# via lerobot
fonttools==4.60.1
# via matplotlib # via matplotlib
frozenlist==1.8.0 frozenlist==1.8.0
# via # via
# aiohttp # aiohttp
# aiosignal # aiosignal
fsspec[http]==2026.2.0 fsspec[http]==2025.9.0
# via # via
# datasets # datasets
# etils # etils
# huggingface-hub # huggingface-hub
# torch # torch
future==1.0.0 future==1.0.0
# via hf-libero # via libero
gitdb==4.0.12 gitdb==4.0.12
# via gitpython # via gitpython
gitpython==3.1.46 gitpython==3.1.45
# via wandb # via wandb
glfw==2.10.0 glfw==2.10.0
# via # via
@@ -235,60 +230,50 @@ gym-hil==0.1.13
# via lerobot # via lerobot
gym-pusht==0.1.6 gym-pusht==0.1.6
# via lerobot # via lerobot
gymnasium==1.2.3 gymnasium==1.2.1
# via # via
# gym-aloha # gym-aloha
# gym-hil # gym-hil
# gym-pusht # gym-pusht
# hf-libero
# lerobot # lerobot
# libero
# metaworld # metaworld
h11==0.16.0 h11==0.16.0
# via # via uvicorn
# httpcore h5py==3.15.1
# uvicorn
h5py==3.16.0
# via robomimic # via robomimic
hebi-py==2.11.0 hebi-py==2.11.0
# via lerobot # via lerobot
hf-egl-probe==1.0.2 hf-transfer==0.1.9
# via hf-libero # via huggingface-hub
hf-libero==0.1.3 hf-xet==1.1.10
# via lerobot
hf-xet==1.3.2
# via huggingface-hub # via huggingface-hub
hidapi==0.14.0.post4 hidapi==0.14.0.post4
# via # via
# gym-hil # gym-hil
# lerobot # lerobot
httpcore==1.0.9
# via httpx
httptools==0.7.1 httptools==0.7.1
# via uvicorn # via uvicorn
httpx==0.28.1 huggingface-hub[cli,hf-transfer]==0.35.3
# via
# datasets
# huggingface-hub
huggingface-hub==1.6.0
# via # via
# accelerate # accelerate
# datasets # datasets
# diffusers # diffusers
# lerobot # lerobot
# peft # peft
# timm
# tokenizers # tokenizers
# transformers # transformers
hydra-core==1.3.2 hydra-core==1.3.2
# via hf-libero # via libero
identify==2.6.17 identify==2.6.15
# via pre-commit # via pre-commit
idna==3.11 idna==3.11
# via # via
# anyio # anyio
# httpx
# requests # requests
# yarl # yarl
imageio[ffmpeg]==2.37.2 imageio[ffmpeg]==2.37.0
# via # via
# gym-aloha # gym-aloha
# gym-hil # gym-hil
@@ -300,14 +285,16 @@ imageio-ffmpeg==0.6.0
# via # via
# imageio # imageio
# robomimic # robomimic
importlib-metadata==8.7.1 importlib-metadata==8.7.0
# via diffusers # via diffusers
importlib-resources==6.5.2
# via etils
iniconfig==2.3.0 iniconfig==2.3.0
# via pytest # via pytest
ipython==9.11.0 inquirerpy==0.3.4
# via huggingface-hub
ipython==8.37.0
# via meshcat # via meshcat
ipython-pygments-lexers==1.1.1
# via ipython
ischedule==1.2.7 ischedule==1.2.7
# via placo # via placo
jedi==0.19.2 jedi==0.19.2
@@ -316,41 +303,40 @@ jinja2==3.1.6
# via torch # via torch
jsonlines==4.0.0 jsonlines==4.0.0
# via lerobot # via lerobot
jsonschema==4.26.0 jsonschema==4.25.1
# via nbformat # via nbformat
jsonschema-specifications==2025.9.1 jsonschema-specifications==2025.9.1
# via jsonschema # via jsonschema
jupyter-core==5.9.1 jupyter-core==5.9.1
# via nbformat # via nbformat
jupytext==1.19.1 jupytext==1.18.1
# via bddl # via bddl
kiwisolver==1.4.9 kiwisolver==1.4.9
# via matplotlib # via matplotlib
labmaze==1.0.6 labmaze==1.0.6
# via dm-control # via dm-control
lazy-loader==0.5 lazy-loader==0.4
# via scikit-image # via scikit-image
librt==0.8.1 libero @ git+https://github.com/huggingface/lerobot-libero.git@main
# via mypy # via lerobot
llvmlite==0.46.0 llvmlite==0.45.1
# via numba # via numba
lxml==6.0.2 lxml==6.0.2
# via dm-control # via dm-control
markdown==3.10.2 markdown==3.9
# via tensorboard # via tensorboard
markdown-it-py==4.0.0 markdown-it-py==4.0.0
# via # via
# jupytext # jupytext
# mdit-py-plugins # mdit-py-plugins
# rich
markupsafe==3.0.3 markupsafe==3.0.3
# via # via
# jinja2 # jinja2
# werkzeug # werkzeug
matplotlib==3.10.8 matplotlib==3.10.7
# via # via
# hf-libero
# lerobot # lerobot
# libero
matplotlib-inline==0.2.1 matplotlib-inline==0.2.1
# via ipython # via ipython
mdit-py-plugins==0.5.0 mdit-py-plugins==0.5.0
@@ -367,38 +353,36 @@ mock-serial==0.0.1
# via lerobot # via lerobot
mpmath==1.3.0 mpmath==1.3.0
# via sympy # via sympy
mujoco==3.5.0 mujoco==3.3.7
# via # via
# dm-control # dm-control
# gym-aloha # gym-aloha
# gym-hil # gym-hil
# hf-libero # libero
# metaworld # metaworld
# robosuite # robosuite
multidict==6.7.1 multidict==6.7.0
# via # via
# aiohttp # aiohttp
# yarl # yarl
multiprocess==0.70.18 multiprocess==0.70.16
# via datasets # via datasets
mypy==1.19.1
# via lerobot
mypy-extensions==1.1.0 mypy-extensions==1.1.0
# via # via typing-inspect
# mypy
# typing-inspect
nbformat==5.10.4 nbformat==5.10.4
# via jupytext # via jupytext
networkx==3.6.1 networkx==3.4.2
# via # via
# bddl # bddl
# scikit-image # scikit-image
# torch # torch
nodeenv==1.10.0 ninja==1.13.0
# via lerobot
nodeenv==1.9.1
# via pre-commit # via pre-commit
num2words==0.5.14 num2words==0.5.14
# via lerobot # via lerobot
numba==0.64.0 numba==0.62.1
# via robosuite # via robosuite
numpy==2.2.6 numpy==2.2.6
# via # via
@@ -407,6 +391,7 @@ numpy==2.2.6
# cmeel-boost # cmeel-boost
# contourpy # contourpy
# datasets # datasets
# decord
# diffusers # diffusers
# dm-control # dm-control
# dm-env # dm-env
@@ -414,10 +399,9 @@ numpy==2.2.6
# gymnasium # gymnasium
# h5py # h5py
# hebi-py # hebi-py
# hf-libero
# imageio # imageio
# labmaze # labmaze
# lerobot # libero
# matplotlib # matplotlib
# meshcat # meshcat
# metaworld # metaworld
@@ -442,51 +426,49 @@ numpy==2.2.6
# torchvision # torchvision
# transformers # transformers
# transforms3d # transforms3d
nvidia-cublas-cu12==12.8.4.1 nvidia-cublas-cu12==12.6.4.1
# via # via
# nvidia-cudnn-cu12 # nvidia-cudnn-cu12
# nvidia-cusolver-cu12 # nvidia-cusolver-cu12
# torch # torch
nvidia-cuda-cupti-cu12==12.8.90 nvidia-cuda-cupti-cu12==12.6.80
# via torch # via torch
nvidia-cuda-nvrtc-cu12==12.8.93 nvidia-cuda-nvrtc-cu12==12.6.77
# via torch # via torch
nvidia-cuda-runtime-cu12==12.8.90 nvidia-cuda-runtime-cu12==12.6.77
# via torch # via torch
nvidia-cudnn-cu12==9.10.2.21 nvidia-cudnn-cu12==9.5.1.17
# via torch # via torch
nvidia-cufft-cu12==11.3.3.83 nvidia-cufft-cu12==11.3.0.4
# via torch # via torch
nvidia-cufile-cu12==1.13.1.3 nvidia-cufile-cu12==1.11.1.6
# via torch # via torch
nvidia-curand-cu12==10.3.9.90 nvidia-curand-cu12==10.3.7.77
# via torch # via torch
nvidia-cusolver-cu12==11.7.3.90 nvidia-cusolver-cu12==11.7.1.2
# via torch # via torch
nvidia-cusparse-cu12==12.5.8.93 nvidia-cusparse-cu12==12.5.4.2
# via # via
# nvidia-cusolver-cu12 # nvidia-cusolver-cu12
# torch # torch
nvidia-cusparselt-cu12==0.7.1 nvidia-cusparselt-cu12==0.6.3
# via torch # via torch
nvidia-nccl-cu12==2.27.5 nvidia-nccl-cu12==2.26.2
# via torch # via torch
nvidia-nvjitlink-cu12==12.8.93 nvidia-nvjitlink-cu12==12.6.85
# via # via
# nvidia-cufft-cu12 # nvidia-cufft-cu12
# nvidia-cusolver-cu12 # nvidia-cusolver-cu12
# nvidia-cusparse-cu12 # nvidia-cusparse-cu12
# torch # torch
nvidia-nvshmem-cu12==3.4.5 nvidia-nvtx-cu12==12.6.77
# via torch
nvidia-nvtx-cu12==12.8.90
# via torch # via torch
omegaconf==2.3.0 omegaconf==2.3.0
# via hydra-core # via hydra-core
opencv-python==4.13.0.92 opencv-python==4.12.0.88
# via # via
# gym-pusht # gym-pusht
# hf-libero # libero
# reachy2-sdk # reachy2-sdk
# robosuite # robosuite
opencv-python-headless==4.12.0.88 opencv-python-headless==4.12.0.88
@@ -505,7 +487,6 @@ packaging==25.0
# matplotlib # matplotlib
# peft # peft
# pytest # pytest
# qwen-vl-utils
# reachy2-sdk # reachy2-sdk
# scikit-image # scikit-image
# tensorboard # tensorboard
@@ -516,21 +497,21 @@ pandas==2.3.3
# via # via
# datasets # datasets
# lerobot # lerobot
parso==0.8.6 parso==0.8.5
# via jedi # via jedi
pathspec==1.0.4 peft==0.17.1
# via mypy
peft==0.18.1
# via lerobot # via lerobot
pexpect==4.9.0 pexpect==4.9.0
# via ipython # via ipython
pillow==12.1.1 pfzy==0.3.4
# via inquirerpy
pillow==12.0.0
# via # via
# diffusers # diffusers
# imageio # imageio
# lerobot
# matplotlib # matplotlib
# meshcat # meshcat
# qwen-vl-utils
# rerun-sdk # rerun-sdk
# robosuite # robosuite
# scikit-image # scikit-image
@@ -538,27 +519,28 @@ pillow==12.1.1
# torchvision # torchvision
pin==3.4.0 pin==3.4.0
# via placo # via placo
placo==0.9.16 placo==0.9.14
# via lerobot # via lerobot
platformdirs==4.9.4 platformdirs==4.5.0
# via # via
# jupyter-core # jupyter-core
# python-discovery
# virtualenv # virtualenv
# wandb # wandb
pluggy==1.6.0 pluggy==1.6.0
# via # via
# pytest # pytest
# pytest-cov # pytest-cov
pre-commit==4.5.1 pre-commit==4.3.0
# via lerobot # via lerobot
prompt-toolkit==3.0.52 prompt-toolkit==3.0.52
# via ipython # via
# inquirerpy
# ipython
propcache==0.4.1 propcache==0.4.1
# via # via
# aiohttp # aiohttp
# yarl # yarl
protobuf==6.31.1 protobuf==6.31.0
# via # via
# dm-control # dm-control
# grpcio-tools # grpcio-tools
@@ -568,7 +550,7 @@ protobuf==6.31.1
# tensorboard # tensorboard
# tensorboardx # tensorboardx
# wandb # wandb
psutil==7.2.2 psutil==7.1.1
# via # via
# accelerate # accelerate
# imageio # imageio
@@ -578,17 +560,17 @@ ptyprocess==0.7.0
# via pexpect # via pexpect
pure-eval==0.2.3 pure-eval==0.2.3
# via stack-data # via stack-data
pyarrow==23.0.1 pyarrow==21.0.0
# via # via
# datasets # datasets
# rerun-sdk # rerun-sdk
pycparser==3.0 pycparser==2.23
# via cffi # via cffi
pydantic==2.12.5 pydantic==2.12.3
# via # via
# fastapi # fastapi
# wandb # wandb
pydantic-core==2.41.5 pydantic-core==2.41.4
# via pydantic # via pydantic
pygame==2.6.1 pygame==2.6.1
# via # via
@@ -598,14 +580,12 @@ pygame==2.6.1
pygments==2.19.2 pygments==2.19.2
# via # via
# ipython # ipython
# ipython-pygments-lexers
# pytest # pytest
# rich
pymunk==6.11.1 pymunk==6.11.1
# via # via
# gym-pusht # gym-pusht
# lerobot # lerobot
pyngrok==7.5.1 pyngrok==7.4.1
# via meshcat # via meshcat
pynput==1.8.1 pynput==1.8.1
# via # via
@@ -615,7 +595,7 @@ pyopengl==3.1.10
# via # via
# dm-control # dm-control
# mujoco # mujoco
pyparsing==3.3.2 pyparsing==3.2.5
# via # via
# dm-control # dm-control
# matplotlib # matplotlib
@@ -641,16 +621,13 @@ pytest-timeout==2.4.0
# via lerobot # via lerobot
python-dateutil==2.9.0.post0 python-dateutil==2.9.0.post0
# via # via
# faker
# matplotlib # matplotlib
# pandas # pandas
python-discovery==1.1.1 python-dotenv==1.1.1
# via virtualenv
python-dotenv==1.2.2
# via uvicorn # via uvicorn
python-xlib==0.33 python-xlib==0.33
# via pynput # via pynput
pytz==2026.1.post1 pytz==2025.2
# via pandas # via pandas
pyyaml==6.0.3 pyyaml==6.0.3
# via # via
@@ -665,6 +642,7 @@ pyyaml==6.0.3
# pre-commit # pre-commit
# pyngrok # pyngrok
# pyyaml-include # pyyaml-include
# timm
# transformers # transformers
# uvicorn # uvicorn
# wandb # wandb
@@ -674,9 +652,7 @@ pyzmq==27.1.0
# via # via
# lerobot # lerobot
# meshcat # meshcat
qwen-vl-utils==0.0.14 reachy2-sdk==1.0.14
# via lerobot
reachy2-sdk==1.0.15
# via lerobot # via lerobot
reachy2-sdk-api==1.0.21 reachy2-sdk-api==1.0.21
# via reachy2-sdk # via reachy2-sdk
@@ -684,7 +660,7 @@ referencing==0.37.0
# via # via
# jsonschema # jsonschema
# jsonschema-specifications # jsonschema-specifications
regex==2026.2.28 regex==2025.10.23
# via # via
# diffusers # diffusers
# transformers # transformers
@@ -693,62 +669,60 @@ requests==2.32.5
# datasets # datasets
# diffusers # diffusers
# dm-control # dm-control
# qwen-vl-utils # huggingface-hub
# teleop # teleop
# transformers
# wandb # wandb
rerun-sdk==0.26.2 rerun-sdk==0.26.1
# via lerobot # via lerobot
rhoban-cmeel-jsoncpp==1.9.4.9 rhoban-cmeel-jsoncpp==1.9.4.9
# via placo # via placo
rich==14.3.3
# via typer
robomimic==0.2.0 robomimic==0.2.0
# via hf-libero # via libero
robosuite==1.4.0 robosuite==1.4.0
# via hf-libero # via libero
rpds-py==0.30.0 rpds-py==0.28.0
# via # via
# jsonschema # jsonschema
# referencing # referencing
safetensors==0.7.0 safetensors==0.6.2
# via # via
# accelerate # accelerate
# diffusers # diffusers
# lerobot # lerobot
# peft # peft
# timm
# transformers # transformers
scikit-image==0.25.2 scikit-image==0.25.2
# via # via
# gym-pusht # gym-pusht
# lerobot # lerobot
scipy==1.17.1 scipy==1.15.3
# via # via
# dm-control # dm-control
# lerobot
# metaworld # metaworld
# robosuite # robosuite
# scikit-image # scikit-image
# torchdiffeq sentry-sdk==2.42.1
sentry-sdk==2.54.0
# via wandb # via wandb
shapely==2.1.2 shapely==2.1.2
# via gym-pusht # via gym-pusht
shellingham==1.5.4
# via typer
six==1.17.0 six==1.17.0
# via # via
# pynput # pynput
# python-dateutil # python-dateutil
# python-xlib # python-xlib
smmap==5.0.3 smmap==5.0.2
# via gitdb # via gitdb
sniffio==1.3.1
# via anyio
stack-data==0.6.3 stack-data==0.6.3
# via ipython # via ipython
starlette==0.52.1 starlette==0.48.0
# via fastapi # via fastapi
sympy==1.14.0 sympy==1.14.0
# via torch # via torch
teleop==0.1.4 teleop==0.1.2
# via lerobot # via lerobot
tensorboard==2.20.0 tensorboard==2.20.0
# via robomimic # via robomimic
@@ -756,38 +730,46 @@ tensorboard-data-server==0.7.2
# via tensorboard # via tensorboard
tensorboardx==2.6.4 tensorboardx==2.6.4
# via robomimic # via robomimic
termcolor==3.3.0 termcolor==3.1.0
# via # via
# lerobot # lerobot
# robomimic # robomimic
thop==0.1.1.post2209072238 thop==0.1.1.post2209072238
# via hf-libero # via libero
tifffile==2026.3.3 tifffile==2025.5.10
# via scikit-image # via scikit-image
tokenizers==0.22.2 timm==1.0.20
# via lerobot
tokenizers==0.22.1
# via transformers # via transformers
toml==0.10.2 toml==0.10.2
# via draccus # via draccus
torch==2.10.0 tomli==2.3.0
# via
# cmeel
# coverage
# jupytext
# pytest
torch==2.7.1
# via # via
# accelerate # accelerate
# flash-attn
# lerobot # lerobot
# peft # peft
# robomimic # robomimic
# thop # thop
# torchdiffeq # timm
# torchvision # torchvision
torchcodec==0.10.0 torchcodec==0.5
# via lerobot # via lerobot
torchdiffeq==0.2.5 torchvision==0.22.1
# via lerobot
torchvision==0.25.0
# via # via
# lerobot # lerobot
# robomimic # robomimic
tornado==6.5.4 # timm
tornado==6.5.2
# via meshcat # via meshcat
tqdm==4.67.3 tqdm==4.67.1
# via # via
# datasets # datasets
# dm-control # dm-control
@@ -801,29 +783,26 @@ traitlets==5.14.3
# jupyter-core # jupyter-core
# matplotlib-inline # matplotlib-inline
# nbformat # nbformat
transformers==5.3.0 transformers==4.57.1
# via # via
# hf-libero
# lerobot # lerobot
# libero
# peft # peft
transforms3d==0.4.2 transforms3d==0.4.2
# via teleop # via teleop
triton==3.6.0 triton==3.3.1
# via torch # via torch
typer==0.24.1
# via
# huggingface-hub
# transformers
typing-extensions==4.15.0 typing-extensions==4.15.0
# via # via
# aiosignal # aiosignal
# anyio # anyio
# etils # etils
# faker # exceptiongroup
# fastapi # fastapi
# gymnasium # gymnasium
# huggingface-hub # huggingface-hub
# mypy # ipython
# multidict
# pydantic # pydantic
# pydantic-core # pydantic-core
# referencing # referencing
@@ -832,46 +811,46 @@ typing-extensions==4.15.0
# torch # torch
# typing-inspect # typing-inspect
# typing-inspection # typing-inspection
# uvicorn
# virtualenv
# wandb # wandb
typing-inspect==0.9.0 typing-inspect==0.9.0
# via draccus # via draccus
typing-inspection==0.4.2 typing-inspection==0.4.2
# via # via pydantic
# fastapi tzdata==2025.2
# pydantic
tzdata==2025.3
# via pandas # via pandas
u-msgpack-python==2.8.0 u-msgpack-python==2.8.0
# via meshcat # via meshcat
urllib3==2.6.3 urllib3==2.5.0
# via # via
# requests # requests
# sentry-sdk # sentry-sdk
uvicorn[standard]==0.41.0 uvicorn[standard]==0.38.0
# via teleop # via teleop
uvloop==0.22.1 uvloop==0.22.1
# via uvicorn # via uvicorn
virtualenv==21.1.0 virtualenv==20.35.3
# via pre-commit # via pre-commit
wandb==0.24.2 wandb==0.21.4
# via # via
# hf-libero
# lerobot # lerobot
# libero
watchfiles==1.1.1 watchfiles==1.1.1
# via uvicorn # via uvicorn
wcwidth==0.6.0 wcwidth==0.2.14
# via prompt-toolkit # via prompt-toolkit
websocket-client==1.9.0 websocket-client==1.9.0
# via teleop # via teleop
websockets==16.0 websockets==15.0.1
# via uvicorn # via uvicorn
werkzeug==3.1.6 werkzeug==3.1.3
# via tensorboard # via tensorboard
wrapt==2.1.2 wrapt==2.0.0
# via dm-tree # via dm-tree
xxhash==3.6.0 xxhash==3.6.0
# via datasets # via datasets
yarl==1.23.0 yarl==1.22.0
# via aiohttp # via aiohttp
zipp==3.23.0 zipp==3.23.0
# via # via
+4 -4
View File
@@ -1,9 +1,9 @@
# requirements.in # requirements.in
# requirements-macos.txt was generated on macOS and is platform-specific (macOS 26.3.1 25D2128 arm64). # requirements-macos.txt was generated on macOS and is platform-specific (macOS 26.0.1 25A362 arm64).
# Darwin MacBook-Pro.local 25.3.0 Darwin Kernel Version 25.3.0: Wed Jan 28 20:54:55 PST 2026; root:xnu-12377.91.3~2/RELEASE_ARM64_T8132 arm64 # Darwin MacBook-Pro.local 25.0.0 Darwin Kernel Version 25.0.0: Wed Sep 17 21:42:08 PDT 2025; root:xnu-12377.1.9~141/RELEASE_ARM64_T8132 arm64
# requirements-ubuntu.txt was generated on Linux and is platform-specific (Ubuntu 24.04.4 LTS x86_64). # requirements-ubuntu.txt was generated on Linux and is platform-specific (Ubuntu 24.04.3 LTS x86_64).
# Linux lerobot-linux 6.17.0-14-generic #14~24.04.1-Ubuntu SMP PREEMPT_DYNAMIC Thu Jan 15 15:52:10 UTC 2 x86_64 x86_64 x86_64 GNU/Linux # Linux mlerobot-linux 6.14.0-33-generic #33~24.04.1-Ubuntu SMP PREEMPT_DYNAMIC Fri Sep 19 17:02:30 UTC 2 x86_64 x86_64 x86_64 GNU/Linux
-e .[all] -e .[all]
+3 -5
View File
@@ -63,9 +63,9 @@ from lerobot.transport import (
services_pb2_grpc, # type: ignore services_pb2_grpc, # type: ignore
) )
from lerobot.transport.utils import grpc_channel_options, send_bytes_in_chunks from lerobot.transport.utils import grpc_channel_options, send_bytes_in_chunks
from lerobot.utils.import_utils import register_third_party_plugins
from .configs import RobotClientConfig from .configs import RobotClientConfig
from .constants import SUPPORTED_ROBOTS
from .helpers import ( from .helpers import (
Action, Action,
FPSTracker, FPSTracker,
@@ -485,9 +485,8 @@ class RobotClient:
def async_client(cfg: RobotClientConfig): def async_client(cfg: RobotClientConfig):
logging.info(pformat(asdict(cfg))) logging.info(pformat(asdict(cfg)))
# TODO: Assert if checking robot support is still needed with the plugin system if cfg.robot.type not in SUPPORTED_ROBOTS:
# if cfg.robot.type not in SUPPORTED_ROBOTS: raise ValueError(f"Robot {cfg.robot.type} not yet supported!")
# raise ValueError(f"Robot {cfg.robot.type} not yet supported!")
client = RobotClient(cfg) client = RobotClient(cfg)
@@ -513,5 +512,4 @@ def async_client(cfg: RobotClientConfig):
if __name__ == "__main__": if __name__ == "__main__":
register_third_party_plugins()
async_client() # run the client async_client() # run the client
+1 -1
View File
@@ -181,7 +181,7 @@ class ZMQCamera(Camera):
try: try:
message = self.socket.recv_string() message = self.socket.recv_string()
except Exception as e: except Exception as e:
# zmq is lazy-imported in connect(), so check by name to avoid a top-level import # Check for ZMQ timeout (EAGAIN/Again) without requiring global zmq import
if type(e).__name__ == "Again": if type(e).__name__ == "Again":
raise TimeoutError(f"{self} timeout after {self.timeout_ms}ms") from e raise TimeoutError(f"{self} timeout after {self.timeout_ms}ms") from e
raise raise
+4 -72
View File
@@ -23,7 +23,6 @@ import base64
import contextlib import contextlib
import json import json
import logging import logging
import threading
import time import time
from collections import deque from collections import deque
@@ -43,57 +42,10 @@ def encode_image(image: np.ndarray, quality: int = 80) -> str:
return base64.b64encode(buffer).decode("utf-8") return base64.b64encode(buffer).decode("utf-8")
class CameraCaptureThread:
"""Background thread that continuously captures and encodes frames from a camera."""
def __init__(self, camera: OpenCVCamera, name: str):
self.camera = camera
self.name = name
self.latest_encoded: str | None = None # Pre-encoded JPEG as base64
self.latest_timestamp: float = 0.0
self.frame_lock = threading.Lock()
self.running = False
self.thread: threading.Thread | None = None
def start(self):
"""Start the capture thread."""
self.running = True
self.thread = threading.Thread(target=self._capture_loop, daemon=True)
self.thread.start()
def stop(self):
"""Stop the capture thread."""
self.running = False
if self.thread:
self.thread.join(timeout=1.0)
def _capture_loop(self):
"""Continuously capture and encode frames at the camera's native rate."""
while self.running:
try:
frame = self.camera.read() # Blocks at camera's native rate
timestamp = time.time()
# Encode immediately in capture thread (this is the slow part)
encoded = encode_image(frame)
with self.frame_lock:
self.latest_encoded = encoded
self.latest_timestamp = timestamp
except Exception as e:
logger.warning(f"Camera {self.name} capture error: {e}")
time.sleep(0.01)
def get_latest(self) -> tuple[str | None, float]:
"""Get the latest encoded frame and its timestamp."""
with self.frame_lock:
return self.latest_encoded, self.latest_timestamp
class ImageServer: class ImageServer:
def __init__(self, config: dict, port: int = 5555): def __init__(self, config: dict, port: int = 5555):
# fps controls the publish loop rate (how often frames are sent over ZMQ), not the camera capture rate
self.fps = config.get("fps", 30) self.fps = config.get("fps", 30)
self.cameras: dict[str, OpenCVCamera] = {} self.cameras: dict[str, OpenCVCamera] = {}
self.capture_threads: dict[str, CameraCaptureThread] = {}
for name, cfg in config.get("cameras", {}).items(): for name, cfg in config.get("cameras", {}).items():
shape = cfg.get("shape", [480, 640]) shape = cfg.get("shape", [480, 640])
@@ -109,10 +61,6 @@ class ImageServer:
self.cameras[name] = camera self.cameras[name] = camera
logger.info(f"Camera {name}: {shape[1]}x{shape[0]}") logger.info(f"Camera {name}: {shape[1]}x{shape[0]}")
# Create capture thread for this camera
capture_thread = CameraCaptureThread(camera, name)
self.capture_threads[name] = capture_thread
# ZMQ PUB socket # ZMQ PUB socket
self.context = zmq.Context() self.context = zmq.Context()
self.socket = self.context.socket(zmq.PUB) self.socket = self.context.socket(zmq.PUB)
@@ -125,18 +73,6 @@ class ImageServer:
def run(self): def run(self):
frame_count = 0 frame_count = 0
frame_times = deque(maxlen=60) frame_times = deque(maxlen=60)
last_published_ts: dict[str, float] = {}
# Start all capture threads
for capture_thread in self.capture_threads.values():
capture_thread.start()
# Wait for first frames to be captured and encoded
logger.info("Waiting for cameras to start capturing...")
for name, capture_thread in self.capture_threads.items():
while capture_thread.get_latest()[0] is None:
time.sleep(0.01)
logger.info(f"Camera {name} ready (capture + encode in background)")
try: try:
while True: while True:
@@ -144,12 +80,10 @@ class ImageServer:
# Build message # Build message
message = {"timestamps": {}, "images": {}} message = {"timestamps": {}, "images": {}}
for name, capture_thread in self.capture_threads.items(): for name, cam in self.cameras.items():
encoded, timestamp = capture_thread.get_latest() frame = cam.read() # Returns RGB
if encoded is not None and timestamp > last_published_ts.get(name, 0.0): message["timestamps"][name] = time.time()
message["timestamps"][name] = timestamp message["images"][name] = encode_image(frame)
message["images"][name] = encoded
last_published_ts[name] = timestamp
# Send as JSON string (suppress if buffer full) # Send as JSON string (suppress if buffer full)
with contextlib.suppress(zmq.Again): with contextlib.suppress(zmq.Again):
@@ -168,8 +102,6 @@ class ImageServer:
except KeyboardInterrupt: except KeyboardInterrupt:
pass pass
finally: finally:
for capture_thread in self.capture_threads.values():
capture_thread.stop()
for cam in self.cameras.values(): for cam in self.cameras.values():
cam.disconnect() cam.disconnect()
self.socket.close() self.socket.close()
+6 -27
View File
@@ -16,13 +16,18 @@
from dataclasses import dataclass, field from dataclasses import dataclass, field
from lerobot.datasets.transforms import DatasetTransformStepConfig, ImageTransformsConfig from lerobot.datasets.transforms import ImageTransformsConfig
from lerobot.datasets.video_utils import get_safe_default_codec from lerobot.datasets.video_utils import get_safe_default_codec
@dataclass @dataclass
class DatasetConfig: class DatasetConfig:
# You may provide a list of datasets here. `train.py` creates them all and concatenates them. Note: only data
# keys common between the datasets are kept. Each dataset gets and additional transform that inserts the
# "dataset_index" into the returned item. The index mapping is made according to the order in which the
# datasets are provided.
repo_id: str repo_id: str
# Root directory where the dataset will be stored (e.g. 'dataset/path').
root: str | None = None root: str | None = None
episodes: list[int] | None = None episodes: list[int] | None = None
image_transforms: ImageTransformsConfig = field(default_factory=ImageTransformsConfig) image_transforms: ImageTransformsConfig = field(default_factory=ImageTransformsConfig)
@@ -32,32 +37,6 @@ class DatasetConfig:
streaming: bool = False streaming: bool = False
@dataclass
class SubDatasetConfig:
"""Configuration for a single dataset within a MultiDatasetConfig."""
repo_id: str
root: str | None = None
episodes: list[int] | None = None
revision: str | None = None
video_backend: str = field(default_factory=get_safe_default_codec)
weight: float = 1.0
# Maps dataset-local feature keys to unified policy keys.
# Keys not listed pass through unchanged.
feature_map: dict[str, str] = field(default_factory=dict)
# Per-dataset transforms applied after feature renaming, before cross-dataset padding.
transforms: list[DatasetTransformStepConfig] | None = None
@dataclass
class MultiDatasetConfig:
"""Configuration for training on multiple datasets jointly."""
datasets: list[SubDatasetConfig] = field(default_factory=list)
image_transforms: ImageTransformsConfig = field(default_factory=ImageTransformsConfig)
use_imagenet_stats: bool = True
@dataclass @dataclass
class WandBConfig: class WandBConfig:
enable: bool = False enable: bool = False
+6 -9
View File
@@ -24,7 +24,7 @@ from huggingface_hub.errors import HfHubHTTPError
from lerobot import envs from lerobot import envs
from lerobot.configs import parser from lerobot.configs import parser
from lerobot.configs.default import DatasetConfig, EvalConfig, MultiDatasetConfig, PeftConfig, WandBConfig from lerobot.configs.default import DatasetConfig, EvalConfig, PeftConfig, WandBConfig
from lerobot.configs.policies import PreTrainedConfig from lerobot.configs.policies import PreTrainedConfig
from lerobot.optim import OptimizerConfig from lerobot.optim import OptimizerConfig
from lerobot.optim.schedulers import LRSchedulerConfig from lerobot.optim.schedulers import LRSchedulerConfig
@@ -35,7 +35,7 @@ TRAIN_CONFIG_NAME = "train_config.json"
@dataclass @dataclass
class TrainPipelineConfig(HubMixin): class TrainPipelineConfig(HubMixin):
dataset: DatasetConfig | MultiDatasetConfig dataset: DatasetConfig
env: envs.EnvConfig | None = None env: envs.EnvConfig | None = None
policy: PreTrainedConfig | None = None policy: PreTrainedConfig | None = None
# Set `dir` to where you would like to save all of the run outputs. If you run another training session # Set `dir` to where you would like to save all of the run outputs. If you run another training session
@@ -50,9 +50,6 @@ class TrainPipelineConfig(HubMixin):
# `seed` is used for training (eg: model initialization, dataset shuffling) # `seed` is used for training (eg: model initialization, dataset shuffling)
# AND for the evaluation environments. # AND for the evaluation environments.
seed: int | None = 1000 seed: int | None = 1000
# Set to True to use deterministic cuDNN algorithms for reproducibility.
# This disables cudnn.benchmark and may reduce training speed by ~10-20%.
cudnn_deterministic: bool = False
# Number of workers for the dataloader. # Number of workers for the dataloader.
num_workers: int = 4 num_workers: int = 4
batch_size: int = 8 batch_size: int = 8
@@ -129,9 +126,8 @@ class TrainPipelineConfig(HubMixin):
train_dir = f"{now:%Y-%m-%d}/{now:%H-%M-%S}_{self.job_name}" train_dir = f"{now:%Y-%m-%d}/{now:%H-%M-%S}_{self.job_name}"
self.output_dir = Path("outputs/train") / train_dir self.output_dir = Path("outputs/train") / train_dir
if isinstance(self.dataset, MultiDatasetConfig): if isinstance(self.dataset.repo_id, list):
if len(self.dataset.datasets) < 1: raise NotImplementedError("LeRobotMultiDataset is not currently implemented.")
raise ValueError("MultiDatasetConfig.datasets must contain at least one sub-dataset.")
if not self.use_policy_training_preset and (self.optimizer is None or self.scheduler is None): if not self.use_policy_training_preset and (self.optimizer is None or self.scheduler is None):
raise ValueError("Optimizer and Scheduler must be set when the policy presets are not used.") raise ValueError("Optimizer and Scheduler must be set when the policy presets are not used.")
@@ -144,7 +140,8 @@ class TrainPipelineConfig(HubMixin):
"'policy.repo_id' argument missing. Please specify it to push the model to the hub." "'policy.repo_id' argument missing. Please specify it to push the model to the hub."
) )
if self.use_rabc and not self.rabc_progress_path and isinstance(self.dataset, DatasetConfig): if self.use_rabc and not self.rabc_progress_path:
# Auto-detect from dataset path
repo_id = self.dataset.repo_id repo_id = self.dataset.repo_id
if self.dataset.root: if self.dataset.root:
self.rabc_progress_path = str(Path(self.dataset.root) / "sarm_progress.parquet") self.rabc_progress_path = str(Path(self.dataset.root) / "sarm_progress.parquet")
+1 -3
View File
@@ -289,9 +289,7 @@ def aggregate_datasets(
logging.info("Find all tasks") logging.info("Find all tasks")
unique_tasks = pd.concat([m.tasks for m in all_metadata]).index.unique() unique_tasks = pd.concat([m.tasks for m in all_metadata]).index.unique()
dst_meta.tasks = pd.DataFrame( dst_meta.tasks = pd.DataFrame({"task_index": range(len(unique_tasks))}, index=unique_tasks)
{"task_index": range(len(unique_tasks))}, index=pd.Index(unique_tasks, name="task")
)
meta_idx = {"chunk": 0, "file": 0} meta_idx = {"chunk": 0, "file": 0}
data_idx = {"chunk": 0, "file": 0} data_idx = {"chunk": 0, "file": 0}
-7
View File
@@ -7,13 +7,6 @@
This dataset was created using [LeRobot](https://github.com/huggingface/lerobot). This dataset was created using [LeRobot](https://github.com/huggingface/lerobot).
{% if repo_id is defined and repo_id %}
<a class="flex" href="https://huggingface.co/spaces/lerobot/visualize_dataset?path={{ repo_id }}">
<img class="block dark:hidden" src="https://huggingface.co/datasets/huggingface/badges/resolve/main/visualize-this-dataset-xl.svg"/>
<img class="hidden dark:block" src="https://huggingface.co/datasets/huggingface/badges/resolve/main/visualize-this-dataset-xl-dark.svg"/>
</a>
{% endif %}
## Dataset Description ## Dataset Description
{{ dataset_description | default("", true) }} {{ dataset_description | default("", true) }}
+33 -41
View File
@@ -89,8 +89,8 @@ def delete_episodes(
Args: Args:
dataset: The source LeRobotDataset. dataset: The source LeRobotDataset.
episode_indices: List of episode indices to delete. episode_indices: List of episode indices to delete.
output_dir: Root directory where the edited dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id. Equivalent to new_root in EditDatasetConfig. output_dir: Directory to save the new dataset. If None, uses default location.
repo_id: Edited dataset identifier. Equivalent to new_repo_id in EditDatasetConfig. repo_id: Repository ID for the new dataset. If None, appends "_modified" to original.
""" """
if not episode_indices: if not episode_indices:
raise ValueError("No episodes to delete") raise ValueError("No episodes to delete")
@@ -152,7 +152,7 @@ def split_dataset(
dataset: The source LeRobotDataset to split. dataset: The source LeRobotDataset to split.
splits: Either a dict mapping split names to episode indices, or a dict mapping splits: Either a dict mapping split names to episode indices, or a dict mapping
split names to fractions (must sum to <= 1.0). split names to fractions (must sum to <= 1.0).
output_dir: Root directory where the split datasets will be stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id. output_dir: Base directory for output datasets. If None, uses default location.
Examples: Examples:
Split by specific episodes Split by specific episodes
@@ -243,8 +243,8 @@ def merge_datasets(
Args: Args:
datasets: List of LeRobotDatasets to merge. datasets: List of LeRobotDatasets to merge.
output_repo_id: Merged dataset identifier. output_repo_id: Repository ID for the merged dataset.
output_dir: Root directory where the merged dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/output_repo_id. output_dir: Directory to save the merged dataset. If None, uses default location.
""" """
if not datasets: if not datasets:
raise ValueError("No datasets to merge") raise ValueError("No datasets to merge")
@@ -288,8 +288,8 @@ def modify_features(
dataset: The source LeRobotDataset. dataset: The source LeRobotDataset.
add_features: Optional dict mapping feature names to (feature_values, feature_info) tuples. add_features: Optional dict mapping feature names to (feature_values, feature_info) tuples.
remove_features: Optional feature name(s) to remove. Can be a single string or list. remove_features: Optional feature name(s) to remove. Can be a single string or list.
output_dir: Root directory where the edited dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id. Equivalent to new_root in EditDatasetConfig. output_dir: Directory to save the new dataset. If None, uses default location.
repo_id: Edited dataset identifier. Equivalent to new_repo_id in EditDatasetConfig. repo_id: Repository ID for the new dataset. If None, appends "_modified" to original.
Returns: Returns:
New dataset with features modified. New dataset with features modified.
@@ -390,8 +390,8 @@ def add_features(
Args: Args:
dataset: The source LeRobotDataset. dataset: The source LeRobotDataset.
features: Dictionary mapping feature names to (feature_values, feature_info) tuples. features: Dictionary mapping feature names to (feature_values, feature_info) tuples.
output_dir: Root directory where the edited dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id. Equivalent to new_root in EditDatasetConfig. output_dir: Directory to save the new dataset. If None, uses default location.
repo_id: Edited dataset identifier. Equivalent to new_repo_id in EditDatasetConfig. repo_id: Repository ID for the new dataset. If None, appends "_modified" to original.
Returns: Returns:
New dataset with all features added. New dataset with all features added.
@@ -427,8 +427,8 @@ def remove_feature(
Args: Args:
dataset: The source LeRobotDataset. dataset: The source LeRobotDataset.
feature_names: Name(s) of features to remove. Can be a single string or list. feature_names: Name(s) of features to remove. Can be a single string or list.
output_dir: Root directory where the edited dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id. Equivalent to new_root in EditDatasetConfig. output_dir: Directory to save the new dataset. If None, uses default location.
repo_id: Edited dataset identifier. Equivalent to new_repo_id in EditDatasetConfig. repo_id: Repository ID for the new dataset. If None, appends "_modified" to original.
Returns: Returns:
New dataset with features removed. New dataset with features removed.
@@ -567,22 +567,20 @@ def _copy_and_reindex_data(
def _keep_episodes_from_video_with_av( def _keep_episodes_from_video_with_av(
input_path: Path, input_path: Path,
output_path: Path, output_path: Path,
episodes_to_keep: list[tuple[int, int]], episodes_to_keep: list[tuple[float, float]],
fps: float, fps: float,
vcodec: str = "libsvtav1", vcodec: str = "libsvtav1",
pix_fmt: str = "yuv420p", pix_fmt: str = "yuv420p",
) -> None: ) -> None:
"""Keep only specified episodes from a video file using PyAV. """Keep only specified episodes from a video file using PyAV.
This function decodes frames from specified frame ranges and re-encodes them with This function decodes frames from specified time ranges and re-encodes them with
properly reset timestamps to ensure monotonic progression. properly reset timestamps to ensure monotonic progression.
Args: Args:
input_path: Source video file path. input_path: Source video file path.
output_path: Destination video file path. output_path: Destination video file path.
episodes_to_keep: List of (start_frame, end_frame) tuples for episodes to keep. episodes_to_keep: List of (start_time, end_time) tuples for episodes to keep.
Ranges are half-open intervals: [start_frame, end_frame), where start_frame
is inclusive and end_frame is exclusive.
fps: Frame rate of the video. fps: Frame rate of the video.
vcodec: Video codec to use for encoding. vcodec: Video codec to use for encoding.
pix_fmt: Pixel format for output video. pix_fmt: Pixel format for output video.
@@ -624,10 +622,9 @@ def _keep_episodes_from_video_with_av(
# Create set of (start, end) ranges for fast lookup. # Create set of (start, end) ranges for fast lookup.
# Convert to a sorted list for efficient checking. # Convert to a sorted list for efficient checking.
frame_ranges = sorted(episodes_to_keep) time_ranges = sorted(episodes_to_keep)
# Track frame index for setting PTS and current range being processed. # Track frame index for setting PTS and current range being processed.
src_frame_count = 0
frame_count = 0 frame_count = 0
range_idx = 0 range_idx = 0
@@ -637,20 +634,21 @@ def _keep_episodes_from_video_with_av(
if frame is None: if frame is None:
continue continue
# Check if frame is in any of our desired frame ranges. # Get frame timestamp.
frame_time = float(frame.pts * frame.time_base) if frame.pts is not None else 0.0
# Check if frame is in any of our desired time ranges.
# Skip ranges that have already passed. # Skip ranges that have already passed.
while range_idx < len(frame_ranges) and src_frame_count >= frame_ranges[range_idx][1]: while range_idx < len(time_ranges) and frame_time >= time_ranges[range_idx][1]:
range_idx += 1 range_idx += 1
# If we've passed all ranges, stop processing. # If we've passed all ranges, stop processing.
if range_idx >= len(frame_ranges): if range_idx >= len(time_ranges):
break break
# Check if frame is in current range. # Check if frame is in current range.
start_frame = frame_ranges[range_idx][0] start_ts, end_ts = time_ranges[range_idx]
if frame_time < start_ts:
if src_frame_count < start_frame:
src_frame_count += 1
continue continue
# Frame is in range - create a new frame with reset timestamps. # Frame is in range - create a new frame with reset timestamps.
@@ -663,7 +661,6 @@ def _keep_episodes_from_video_with_av(
for pkt in v_out.encode(new_frame): for pkt in v_out.encode(new_frame):
out.mux(pkt) out.mux(pkt)
src_frame_count += 1
frame_count += 1 frame_count += 1
# Flush encoder. # Flush encoder.
@@ -752,17 +749,15 @@ def _copy_and_reindex_videos(
f"videos/{video_key}/to_timestamp" f"videos/{video_key}/to_timestamp"
] ]
else: else:
# Build list of frame ranges to keep, in sorted order. # Build list of time ranges to keep, in sorted order.
sorted_keep_episodes = sorted(episodes_in_file, key=lambda x: episode_mapping[x]) sorted_keep_episodes = sorted(episodes_in_file, key=lambda x: episode_mapping[x])
episodes_to_keep_ranges: list[tuple[int, int]] = [] episodes_to_keep_ranges: list[tuple[float, float]] = []
for old_idx in sorted_keep_episodes: for old_idx in sorted_keep_episodes:
src_ep = src_dataset.meta.episodes[old_idx] src_ep = src_dataset.meta.episodes[old_idx]
from_frame = round(src_ep[f"videos/{video_key}/from_timestamp"] * src_dataset.meta.fps) from_ts = src_ep[f"videos/{video_key}/from_timestamp"]
to_frame = round(src_ep[f"videos/{video_key}/to_timestamp"] * src_dataset.meta.fps) to_ts = src_ep[f"videos/{video_key}/to_timestamp"]
assert src_ep["length"] == to_frame - from_frame, ( episodes_to_keep_ranges.append((from_ts, to_ts))
f"Episode length mismatch: {src_ep['length']} vs {to_frame - from_frame}"
)
episodes_to_keep_ranges.append((from_frame, to_frame))
# Use PyAV filters to efficiently re-encode only the desired segments. # Use PyAV filters to efficiently re-encode only the desired segments.
assert src_dataset.meta.video_path is not None assert src_dataset.meta.video_path is not None
@@ -1475,9 +1470,7 @@ def modify_tasks(
# Collect all unique tasks and create new task mapping # Collect all unique tasks and create new task mapping
unique_tasks = sorted(set(episode_to_task.values())) unique_tasks = sorted(set(episode_to_task.values()))
new_task_df = pd.DataFrame( new_task_df = pd.DataFrame({"task_index": list(range(len(unique_tasks)))}, index=unique_tasks)
{"task_index": list(range(len(unique_tasks)))}, index=pd.Index(unique_tasks, name="task")
)
task_to_index = {task: idx for idx, task in enumerate(unique_tasks)} task_to_index = {task: idx for idx, task in enumerate(unique_tasks)}
logging.info(f"Modifying tasks in {dataset.repo_id}") logging.info(f"Modifying tasks in {dataset.repo_id}")
@@ -1531,7 +1524,7 @@ def modify_tasks(
def convert_image_to_video_dataset( def convert_image_to_video_dataset(
dataset: LeRobotDataset, dataset: LeRobotDataset,
output_dir: Path | None = None, output_dir: Path,
repo_id: str | None = None, repo_id: str | None = None,
vcodec: str = "libsvtav1", vcodec: str = "libsvtav1",
pix_fmt: str = "yuv420p", pix_fmt: str = "yuv420p",
@@ -1550,8 +1543,8 @@ def convert_image_to_video_dataset(
Args: Args:
dataset: The source LeRobot dataset with images dataset: The source LeRobot dataset with images
output_dir: Root directory where the edited dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id. Equivalent to new_root in EditDatasetConfig. output_dir: Directory to save the new video dataset
repo_id: Edited dataset identifier. Equivalent to new_repo_id in EditDatasetConfig. repo_id: Repository ID for the new dataset (default: original_id + "_video")
vcodec: Video codec (default: libsvtav1) vcodec: Video codec (default: libsvtav1)
pix_fmt: Pixel format (default: yuv420p) pix_fmt: Pixel format (default: yuv420p)
g: Group of pictures size (default: 2) g: Group of pictures size (default: 2)
@@ -1602,7 +1595,6 @@ def convert_image_to_video_dataset(
# Video info will be updated after episodes are encoded # Video info will be updated after episodes are encoded
# Create new metadata for video dataset # Create new metadata for video dataset
output_dir = Path(output_dir) if output_dir is not None else HF_LEROBOT_HOME / repo_id
new_meta = LeRobotDatasetMetadata.create( new_meta = LeRobotDatasetMetadata.create(
repo_id=repo_id, repo_id=repo_id,
fps=dataset.meta.fps, fps=dataset.meta.fps,
+50 -66
View File
@@ -18,14 +18,13 @@ from pprint import pformat
import torch import torch
from lerobot.configs.default import DatasetConfig, MultiDatasetConfig
from lerobot.configs.policies import PreTrainedConfig from lerobot.configs.policies import PreTrainedConfig
from lerobot.configs.train import TrainPipelineConfig from lerobot.configs.train import TrainPipelineConfig
from lerobot.datasets.lerobot_dataset import ( from lerobot.datasets.lerobot_dataset import (
LeRobotDataset, LeRobotDataset,
LeRobotDatasetMetadata, LeRobotDatasetMetadata,
MultiLeRobotDataset,
) )
from lerobot.datasets.multi_dataset import NewMultiLeRobotDataset
from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset
from lerobot.datasets.transforms import ImageTransforms from lerobot.datasets.transforms import ImageTransforms
from lerobot.utils.constants import ACTION, OBS_PREFIX, REWARD from lerobot.utils.constants import ACTION, OBS_PREFIX, REWARD
@@ -69,81 +68,66 @@ def resolve_delta_timestamps(
return delta_timestamps return delta_timestamps
def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | NewMultiLeRobotDataset: def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDataset:
"""Create a single or multi-dataset depending on the config type. """Handles the logic of setting up delta timestamps and image transforms before creating a dataset.
Args:
cfg (TrainPipelineConfig): A TrainPipelineConfig config which contains a DatasetConfig and a PreTrainedConfig.
Raises:
NotImplementedError: The MultiLeRobotDataset is currently deactivated.
Returns: Returns:
LeRobotDataset | NewMultiLeRobotDataset LeRobotDataset | MultiLeRobotDataset
""" """
if isinstance(cfg.dataset, MultiDatasetConfig):
return _make_multi_dataset(cfg)
return _make_single_dataset(cfg)
def _make_single_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset:
ds_cfg: DatasetConfig = cfg.dataset # type: ignore[assignment]
image_transforms = ( image_transforms = (
ImageTransforms(ds_cfg.image_transforms) if ds_cfg.image_transforms.enable else None ImageTransforms(cfg.dataset.image_transforms) if cfg.dataset.image_transforms.enable else None
) )
ds_meta = LeRobotDatasetMetadata(ds_cfg.repo_id, root=ds_cfg.root, revision=ds_cfg.revision)
delta_timestamps = resolve_delta_timestamps(cfg.policy, ds_meta)
if not ds_cfg.streaming: if isinstance(cfg.dataset.repo_id, str):
dataset = LeRobotDataset( ds_meta = LeRobotDatasetMetadata(
ds_cfg.repo_id, cfg.dataset.repo_id, root=cfg.dataset.root, revision=cfg.dataset.revision
root=ds_cfg.root,
episodes=ds_cfg.episodes,
delta_timestamps=delta_timestamps,
image_transforms=image_transforms,
revision=ds_cfg.revision,
video_backend=ds_cfg.video_backend,
tolerance_s=cfg.tolerance_s,
) )
delta_timestamps = resolve_delta_timestamps(cfg.policy, ds_meta)
if not cfg.dataset.streaming:
dataset = LeRobotDataset(
cfg.dataset.repo_id,
root=cfg.dataset.root,
episodes=cfg.dataset.episodes,
delta_timestamps=delta_timestamps,
image_transforms=image_transforms,
revision=cfg.dataset.revision,
video_backend=cfg.dataset.video_backend,
tolerance_s=cfg.tolerance_s,
)
else:
dataset = StreamingLeRobotDataset(
cfg.dataset.repo_id,
root=cfg.dataset.root,
episodes=cfg.dataset.episodes,
delta_timestamps=delta_timestamps,
image_transforms=image_transforms,
revision=cfg.dataset.revision,
max_num_shards=cfg.num_workers,
tolerance_s=cfg.tolerance_s,
)
else: else:
dataset = StreamingLeRobotDataset( raise NotImplementedError("The MultiLeRobotDataset isn't supported for now.")
ds_cfg.repo_id, dataset = MultiLeRobotDataset(
root=ds_cfg.root, cfg.dataset.repo_id,
episodes=ds_cfg.episodes, # TODO(aliberts): add proper support for multi dataset
delta_timestamps=delta_timestamps, # delta_timestamps=delta_timestamps,
image_transforms=image_transforms, image_transforms=image_transforms,
revision=ds_cfg.revision, video_backend=cfg.dataset.video_backend,
max_num_shards=cfg.num_workers, )
tolerance_s=cfg.tolerance_s, logging.info(
"Multiple datasets were provided. Applied the following index mapping to the provided datasets: "
f"{pformat(dataset.repo_id_to_index, indent=2)}"
) )
if ds_cfg.use_imagenet_stats: if cfg.dataset.use_imagenet_stats:
for key in dataset.meta.camera_keys: for key in dataset.meta.camera_keys:
for stats_type, stats_val in IMAGENET_STATS.items(): for stats_type, stats in IMAGENET_STATS.items():
dataset.meta.stats[key][stats_type] = torch.tensor(stats_val, dtype=torch.float32) dataset.meta.stats[key][stats_type] = torch.tensor(stats, dtype=torch.float32)
return dataset
def _make_multi_dataset(cfg: TrainPipelineConfig) -> NewMultiLeRobotDataset:
multi_cfg: MultiDatasetConfig = cfg.dataset # type: ignore[assignment]
image_transforms = (
ImageTransforms(multi_cfg.image_transforms) if multi_cfg.image_transforms.enable else None
)
dataset = NewMultiLeRobotDataset(
configs=multi_cfg.datasets,
image_transforms=image_transforms,
tolerance_s=cfg.tolerance_s,
)
logging.info(
"MultiLeRobotDataset created with %d sub-datasets:\n%s",
len(multi_cfg.datasets),
pformat(
{i: c.repo_id for i, c in enumerate(multi_cfg.datasets)},
indent=2,
),
)
if multi_cfg.use_imagenet_stats:
for key in dataset.meta.camera_keys:
for stats_type, stats_val in IMAGENET_STATS.items():
dataset.meta.stats[key][stats_type] = torch.tensor(stats_val, dtype=torch.float32)
return dataset return dataset
+13 -14
View File
@@ -314,7 +314,7 @@ class LeRobotDatasetMetadata:
if self.tasks is None: if self.tasks is None:
new_tasks = tasks new_tasks = tasks
task_indices = range(len(tasks)) task_indices = range(len(tasks))
self.tasks = pd.DataFrame({"task_index": task_indices}, index=pd.Index(tasks, name="task")) self.tasks = pd.DataFrame({"task_index": task_indices}, index=tasks)
else: else:
new_tasks = [task for task in tasks if task not in self.tasks.index] new_tasks = [task for task in tasks if task not in self.tasks.index]
new_task_indices = range(len(self.tasks), len(self.tasks) + len(new_tasks)) new_task_indices = range(len(self.tasks), len(self.tasks) + len(new_tasks))
@@ -664,11 +664,11 @@ class LeRobotDataset(torch.utils.data.Dataset):
for the README). for the README).
Args: Args:
repo_id (str): This is the repo id that will be used to fetch the dataset. repo_id (str): This is the repo id that will be used to fetch the dataset. Locally, the dataset
root (Path | None, optional): Local directory where the dataset will be downloaded and will be stored under root/repo_id.
stored. If set, all dataset files will be stored directly under this path. If not set, the root (Path | None, optional): Local directory to use for downloading/writing files. You can also
dataset files will be stored under $HF_LEROBOT_HOME/repo_id (configurable via the set the HF_LEROBOT_HOME environment variable to point to a different location. Defaults to
HF_LEROBOT_HOME environment variable). '~/.cache/huggingface/lerobot'.
episodes (list[int] | None, optional): If specified, this will only load episodes specified by episodes (list[int] | None, optional): If specified, this will only load episodes specified by
their episode_index in this list. Defaults to None. their episode_index in this list. Defaults to None.
image_transforms (Callable | None, optional): You can pass standard v2 image transforms from image_transforms (Callable | None, optional): You can pass standard v2 image transforms from
@@ -747,7 +747,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
# Check if cached dataset contains all requested episodes # Check if cached dataset contains all requested episodes
if not self._check_cached_episodes_sufficient(): if not self._check_cached_episodes_sufficient():
raise FileNotFoundError("Cached dataset doesn't contain all requested episodes") raise FileNotFoundError("Cached dataset doesn't contain all requested episodes")
except (FileNotFoundError, NotADirectoryError): except (AssertionError, FileNotFoundError, NotADirectoryError):
if is_valid_version(self.revision): if is_valid_version(self.revision):
self.revision = get_safe_version(self.repo_id, self.revision) self.revision = get_safe_version(self.repo_id, self.revision)
self.download(download_videos) self.download(download_videos)
@@ -839,7 +839,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
hub_api.upload_folder(**upload_kwargs) hub_api.upload_folder(**upload_kwargs)
card = create_lerobot_dataset_card( card = create_lerobot_dataset_card(
tags=tags, dataset_info=self.meta.info, license=license, repo_id=self.repo_id, **card_kwargs tags=tags, dataset_info=self.meta.info, license=license, **card_kwargs
) )
card.push_to_hub(repo_id=self.repo_id, repo_type="dataset", revision=branch) card.push_to_hub(repo_id=self.repo_id, repo_type="dataset", revision=branch)
@@ -1771,12 +1771,11 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
) )
for repo_id, ds in zip(self.repo_ids, self._datasets, strict=True): for repo_id, ds in zip(self.repo_ids, self._datasets, strict=True):
extra_keys = set(ds.features).difference(intersection_features) extra_keys = set(ds.features).difference(intersection_features)
if extra_keys: logging.warning(
logging.warning( f"keys {extra_keys} of {repo_id} were disabled as they are not contained in all the "
f"keys {extra_keys} of {repo_id} were disabled as they are not contained in all the " "other datasets."
"other datasets." )
) self.disabled_features.update(extra_keys)
self.disabled_features.update(extra_keys)
self.image_transforms = image_transforms self.image_transforms = image_transforms
self.delta_timestamps = delta_timestamps self.delta_timestamps = delta_timestamps
-364
View File
@@ -1,364 +0,0 @@
"""MultiLeRobotDataset: joint training over heterogeneous LeRobot datasets.
Supports:
- Per-dataset feature mapping (rename keys to a unified namespace)
- Automatic zero-padding for features missing in some datasets
- Per-dataset transform pipelines
- Weighted sampling via dataset weights
- Aggregated stats across all sub-datasets
- A ``meta`` shim compatible with EpisodeAwareSampler and make_policy
"""
from __future__ import annotations
import logging
from collections.abc import Callable
import numpy as np
import torch
import torch.utils.data
from lerobot.configs.default import SubDatasetConfig
from lerobot.datasets.compute_stats import aggregate_stats
from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.datasets.transforms import DatasetTransformPipeline
class MultiDatasetMeta:
"""Lightweight metadata shim that exposes the same interface as ``LeRobotDatasetMetadata``.
Built by aggregating the metadata of multiple sub-datasets after their
feature keys have been mapped to a unified namespace.
"""
def __init__(
self,
datasets: list[LeRobotDataset],
feature_maps: list[dict[str, str]],
):
self._datasets = datasets
self._feature_maps = feature_maps
self._unified_features = self._build_unified_features()
self._episodes = self._build_episodes()
self._stats = self._build_stats()
# ------------------------------------------------------------------
# Feature union
# ------------------------------------------------------------------
def _build_unified_features(self) -> dict[str, dict]:
"""Build feature dict as the *union* of all mapped feature keys."""
unified: dict[str, dict] = {}
for ds, fmap in zip(self._datasets, self._feature_maps):
for original_key, feat_info in ds.meta.features.items():
mapped_key = fmap.get(original_key, original_key)
if mapped_key not in unified:
unified[mapped_key] = dict(feat_info)
else:
existing_shape = tuple(unified[mapped_key]["shape"])
new_shape = tuple(feat_info["shape"])
if existing_shape != new_shape and unified[mapped_key]["dtype"] == feat_info["dtype"]:
logging.warning(
"Feature '%s' has shape %s in one dataset but %s in another. "
"The larger shape will be used (padding applied automatically).",
mapped_key,
existing_shape,
new_shape,
)
if np.prod(new_shape) > np.prod(existing_shape):
unified[mapped_key] = dict(feat_info)
return unified
# ------------------------------------------------------------------
# Episode metadata (global flat indexing)
# ------------------------------------------------------------------
def _build_episodes(self) -> dict[str, list]:
"""Concatenate episode boundaries across sub-datasets with frame offsets.
Produces the same column structure as ``load_episodes()`` so that
``EpisodeAwareSampler`` and ``WeightedEpisodeAwareSampler`` can consume it.
"""
from_indices: list[int] = []
to_indices: list[int] = []
dataset_source: list[int] = []
frame_offset = 0
for ds_idx, ds in enumerate(self._datasets):
eps = ds.meta.episodes
for ep in eps:
from_indices.append(ep["dataset_from_index"] + frame_offset)
to_indices.append(ep["dataset_to_index"] + frame_offset)
dataset_source.append(ds_idx)
frame_offset += ds.num_frames
return {
"dataset_from_index": from_indices,
"dataset_to_index": to_indices,
"dataset_source": dataset_source,
}
# ------------------------------------------------------------------
# Stats aggregation
# ------------------------------------------------------------------
def _build_stats(self) -> dict[str, dict[str, np.ndarray]]:
"""Aggregate stats across sub-datasets using mapped feature keys."""
mapped_stats_list: list[dict[str, dict]] = []
for ds, fmap in zip(self._datasets, self._feature_maps):
reverse_map = {v: k for k, v in fmap.items()}
mapped: dict[str, dict] = {}
for unified_key in self._unified_features:
original_key = reverse_map.get(unified_key, unified_key)
if original_key in ds.meta.stats:
mapped[unified_key] = ds.meta.stats[original_key]
mapped_stats_list.append(mapped)
return aggregate_stats(mapped_stats_list)
# ------------------------------------------------------------------
# Properties matching LeRobotDatasetMetadata API
# ------------------------------------------------------------------
@property
def features(self) -> dict[str, dict]:
return self._unified_features
@property
def image_keys(self) -> list[str]:
return [k for k, f in self._unified_features.items() if f["dtype"] == "image"]
@property
def video_keys(self) -> list[str]:
return [k for k, f in self._unified_features.items() if f["dtype"] == "video"]
@property
def camera_keys(self) -> list[str]:
return [k for k, f in self._unified_features.items() if f["dtype"] in ("video", "image")]
@property
def names(self) -> dict[str, list | dict]:
return {k: f["names"] for k, f in self._unified_features.items()}
@property
def shapes(self) -> dict[str, tuple]:
return {k: tuple(f["shape"]) for k, f in self._unified_features.items()}
@property
def fps(self) -> int:
fps_values = {ds.meta.fps for ds in self._datasets}
if len(fps_values) > 1:
logging.warning("Sub-datasets have different FPS values: %s. Using the first.", fps_values)
return self._datasets[0].meta.fps
@property
def stats(self) -> dict[str, dict[str, np.ndarray]]:
return self._stats
@stats.setter
def stats(self, value: dict):
self._stats = value
@property
def episodes(self) -> dict[str, list]:
return self._episodes
@property
def total_episodes(self) -> int:
return sum(ds.meta.total_episodes for ds in self._datasets)
@property
def total_frames(self) -> int:
return sum(ds.meta.total_frames for ds in self._datasets)
@property
def total_tasks(self) -> int:
return sum(ds.meta.total_tasks for ds in self._datasets)
@property
def info(self) -> dict:
return {
"fps": self.fps,
"features": self._unified_features,
"total_episodes": self.total_episodes,
"total_frames": self.total_frames,
"total_tasks": self.total_tasks,
"codebase_version": "v3.0",
}
class NewMultiLeRobotDataset(torch.utils.data.Dataset):
"""Dataset that wraps multiple ``LeRobotDataset`` instances with feature mapping and padding.
Each sub-dataset can have different feature names and shapes. A per-dataset
``feature_map`` renames keys into a shared namespace. Features that a given
sub-dataset does not provide are zero-padded so every ``__getitem__`` returns
the full unified feature set.
"""
def __init__(
self,
configs: list[SubDatasetConfig],
image_transforms: Callable | None = None,
delta_timestamps: dict[str, list[float]] | None = None,
tolerance_s: float = 1e-4,
):
super().__init__()
self._configs = configs
self.image_transforms = image_transforms
self._datasets: list[LeRobotDataset] = []
self._feature_maps: list[dict[str, str]] = []
self._transform_pipelines: list[DatasetTransformPipeline | None] = []
self._weights: list[float] = []
for cfg in configs:
ds = LeRobotDataset(
repo_id=cfg.repo_id,
root=cfg.root,
episodes=cfg.episodes,
image_transforms=image_transforms,
delta_timestamps=delta_timestamps,
tolerance_s=tolerance_s,
revision=cfg.revision,
video_backend=cfg.video_backend,
)
self._datasets.append(ds)
self._feature_maps.append(cfg.feature_map or {})
self._transform_pipelines.append(
DatasetTransformPipeline(cfg.transforms) if cfg.transforms else None
)
self._weights.append(cfg.weight)
self._meta = MultiDatasetMeta(self._datasets, self._feature_maps)
# Pre-compute cumulative frame counts for fast index mapping.
self._cumulative_frames: list[int] = []
total = 0
for ds in self._datasets:
total += ds.num_frames
self._cumulative_frames.append(total)
# Build reverse maps (unified_key -> original_key) per dataset for padding.
self._reverse_maps: list[dict[str, str]] = []
for fmap in self._feature_maps:
self._reverse_maps.append({v: k for k, v in fmap.items()})
logging.info(
"MultiLeRobotDataset: %d sub-datasets, %d total frames, %d total episodes, "
"%d unified features",
len(self._datasets),
self.num_frames,
self.num_episodes,
len(self._meta.features),
)
# ------------------------------------------------------------------
# Public interface
# ------------------------------------------------------------------
@property
def meta(self) -> MultiDatasetMeta:
return self._meta
@property
def dataset_weights(self) -> list[float]:
return self._weights
@property
def num_frames(self) -> int:
return self._cumulative_frames[-1] if self._cumulative_frames else 0
@property
def num_episodes(self) -> int:
return sum(ds.num_episodes for ds in self._datasets)
@property
def episodes(self) -> list[int] | None:
return None
@property
def fps(self) -> int:
return self._meta.fps
@property
def features(self) -> dict[str, dict]:
return self._meta.features
@property
def camera_keys(self) -> list[str]:
return self._meta.camera_keys
# ------------------------------------------------------------------
# Indexing
# ------------------------------------------------------------------
def _locate(self, idx: int) -> tuple[int, int]:
"""Map a global frame index to (dataset_index, local_index)."""
for ds_idx, cum in enumerate(self._cumulative_frames):
if idx < cum:
local = idx - (self._cumulative_frames[ds_idx - 1] if ds_idx > 0 else 0)
return ds_idx, local
raise IndexError(f"Index {idx} out of range (total {self.num_frames})")
def __len__(self) -> int:
return self.num_frames
def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
ds_idx, local_idx = self._locate(idx)
item = self._datasets[ds_idx][local_idx]
# 1. Rename keys according to feature_map.
fmap = self._feature_maps[ds_idx]
if fmap:
renamed: dict[str, torch.Tensor] = {}
for key, value in item.items():
renamed[fmap.get(key, key)] = value
item = renamed
# 2. Apply per-dataset transform pipeline.
pipeline = self._transform_pipelines[ds_idx]
if pipeline is not None:
item = pipeline(item)
# 3. Pad missing features with zeros.
reverse_map = self._reverse_maps[ds_idx]
ds_features = self._datasets[ds_idx].meta.features
for unified_key, feat_info in self._meta.features.items():
if unified_key in item:
continue
original_key = reverse_map.get(unified_key, unified_key)
if original_key in ds_features:
continue
shape = tuple(feat_info["shape"])
dtype = feat_info["dtype"]
if dtype in ("video", "image"):
# Camera tensors are (C, H, W) after transforms.
c, h, w = (shape[2], shape[0], shape[1]) if len(shape) == 3 else (3, shape[0], shape[1])
item[unified_key] = torch.zeros(c, h, w, dtype=torch.float32)
elif dtype in ("float32", "float64"):
item[unified_key] = torch.zeros(shape, dtype=torch.float32)
elif dtype in ("int32", "int64"):
item[unified_key] = torch.zeros(shape, dtype=torch.int64)
elif dtype == "bool":
item[unified_key] = torch.zeros(shape, dtype=torch.bool)
else:
item[unified_key] = torch.zeros(shape, dtype=torch.float32)
item[f"{unified_key}_is_pad"] = torch.tensor(True)
# 4. Tag which dataset this sample came from.
item["dataset_index"] = torch.tensor(ds_idx)
return item
def __repr__(self) -> str:
repo_ids = [c.repo_id for c in self._configs]
return (
f"NewMultiLeRobotDataset(\n"
f" repo_ids={repo_ids},\n"
f" num_frames={self.num_frames},\n"
f" num_episodes={self.num_episodes},\n"
f" unified_features={list(self._meta.features.keys())},\n"
f" weights={self._weights},\n"
f")"
)
-77
View File
@@ -59,80 +59,3 @@ class EpisodeAwareSampler:
def __len__(self) -> int: def __len__(self) -> int:
return len(self.indices) return len(self.indices)
class WeightedEpisodeAwareSampler:
"""Sampler that draws frames from multiple datasets according to per-dataset weights.
Each iteration first selects a sub-dataset proportionally to its weight, then
uniformly samples a frame from that sub-dataset's valid index set. Episode
boundary information is respected so that dropped frames are excluded.
Args:
dataset_from_indices: Start index for each episode (global, flat).
dataset_to_indices: End index (exclusive) for each episode (global, flat).
dataset_membership: Which sub-dataset each episode belongs to (integer id).
dataset_weights: Relative sampling weight per sub-dataset.
episode_indices_to_use: If given, only episodes in this set are used.
drop_n_first_frames: Frames to skip at the start of each episode.
drop_n_last_frames: Frames to skip at the end of each episode.
shuffle: Whether to shuffle within each epoch.
num_samples: How many samples per epoch. Defaults to total valid frames.
generator: Optional torch.Generator for reproducibility.
"""
def __init__(
self,
dataset_from_indices: list[int],
dataset_to_indices: list[int],
dataset_membership: list[int],
dataset_weights: list[float],
episode_indices_to_use: list | None = None,
drop_n_first_frames: int = 0,
drop_n_last_frames: int = 0,
shuffle: bool = False,
num_samples: int | None = None,
generator: torch.Generator | None = None,
):
n_datasets = max(dataset_membership) + 1 if dataset_membership else 0
self._per_dataset_indices: list[list[int]] = [[] for _ in range(n_datasets)]
episodes_to_use = set(episode_indices_to_use) if episode_indices_to_use is not None else None
for ep_idx, (start, end, ds_id) in enumerate(
zip(dataset_from_indices, dataset_to_indices, dataset_membership, strict=True)
):
if episodes_to_use is not None and ep_idx not in episodes_to_use:
continue
frame_range = range(start + drop_n_first_frames, end - drop_n_last_frames)
self._per_dataset_indices[ds_id].extend(frame_range)
# Normalise weights (only over datasets that actually have frames).
raw_weights = list(dataset_weights[:n_datasets])
self._weights = torch.zeros(n_datasets)
for i, w in enumerate(raw_weights):
if len(self._per_dataset_indices[i]) > 0:
self._weights[i] = w
total_w = self._weights.sum()
if total_w > 0:
self._weights /= total_w
self._total_frames = sum(len(idx) for idx in self._per_dataset_indices)
self._num_samples = num_samples if num_samples is not None else self._total_frames
self.shuffle = shuffle
self._generator = generator
def __iter__(self) -> Iterator[int]:
if not self.shuffle:
for ds_indices in self._per_dataset_indices:
yield from ds_indices
return
for _ in range(self._num_samples):
ds_id = int(torch.multinomial(self._weights, 1, generator=self._generator).item())
indices = self._per_dataset_indices[ds_id]
local_idx = int(torch.randint(len(indices), (1,), generator=self._generator).item())
yield indices[local_idx]
def __len__(self) -> int:
return self._num_samples
-113
View File
@@ -14,13 +14,11 @@
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
import collections import collections
import logging
from collections.abc import Callable, Sequence from collections.abc import Callable, Sequence
from dataclasses import dataclass, field from dataclasses import dataclass, field
from typing import Any from typing import Any
import torch import torch
import torch.nn.functional as F_nn
from torchvision.transforms import v2 from torchvision.transforms import v2
from torchvision.transforms.v2 import ( from torchvision.transforms.v2 import (
Transform, Transform,
@@ -260,114 +258,3 @@ class ImageTransforms(Transform):
def forward(self, *inputs: Any) -> Any: def forward(self, *inputs: Any) -> Any:
return self.tf(*inputs) return self.tf(*inputs)
# Per-dataset transform pipeline (used by MultiLeRobotDataset)
@dataclass
class DatasetTransformStepConfig:
"""Config for a single per-dataset transform step."""
type: str
kwargs: dict[str, Any] = field(default_factory=dict)
_DATASET_TRANSFORM_REGISTRY: dict[str, type["DatasetTransformStep"]] = {}
def register_dataset_transform(name: str):
"""Decorator to register a DatasetTransformStep by name."""
def decorator(cls: type["DatasetTransformStep"]) -> type["DatasetTransformStep"]:
_DATASET_TRANSFORM_REGISTRY[name] = cls
return cls
return decorator
class DatasetTransformStep:
"""Base class for a single per-dataset transform applied to a sample dict."""
def __call__(self, sample: dict) -> dict:
raise NotImplementedError
@register_dataset_transform("pad_action")
class PadAction(DatasetTransformStep):
"""Zero-pad the ``action`` tensor to *target_dim* along the last axis."""
def __init__(self, target_dim: int):
self.target_dim = target_dim
def __call__(self, sample: dict) -> dict:
action = sample.get("action")
if action is None:
return sample
current = action.shape[-1]
if current < self.target_dim:
sample["action"] = F_nn.pad(action, (0, self.target_dim - current))
return sample
@register_dataset_transform("pad_state")
class PadState(DatasetTransformStep):
"""Zero-pad ``observation.state`` to *target_dim* along the last axis."""
def __init__(self, target_dim: int):
self.target_dim = target_dim
def __call__(self, sample: dict) -> dict:
state = sample.get("observation.state")
if state is None:
return sample
current = state.shape[-1]
if current < self.target_dim:
sample["observation.state"] = F_nn.pad(state, (0, self.target_dim - current))
return sample
@register_dataset_transform("resize_images")
class ResizeImages(DatasetTransformStep):
"""Resize all image/video camera tensors to (height, width)."""
def __init__(self, height: int, width: int):
self.size = (height, width)
def __call__(self, sample: dict) -> dict:
for key in list(sample.keys()):
if not key.startswith("observation.images."):
continue
img = sample[key]
if not isinstance(img, torch.Tensor) or img.ndim < 3:
continue
sample[key] = F.resize(img, self.size, antialias=True)
return sample
class DatasetTransformPipeline:
"""Sequential pipeline of DatasetTransformStep instances."""
def __init__(self, configs: list[DatasetTransformStepConfig] | None = None):
self.steps: list[DatasetTransformStep] = []
if configs:
for cfg in configs:
self.steps.append(self._build(cfg))
@staticmethod
def _build(cfg: DatasetTransformStepConfig) -> DatasetTransformStep:
cls = _DATASET_TRANSFORM_REGISTRY.get(cfg.type)
if cls is None:
raise ValueError(
f"Unknown dataset transform '{cfg.type}'. "
f"Available: {list(_DATASET_TRANSFORM_REGISTRY)}"
)
return cls(**cfg.kwargs)
def __call__(self, sample: dict) -> dict:
for step in self.steps:
sample = step(sample)
return sample
def __repr__(self) -> str:
return f"DatasetTransformPipeline(steps={self.steps})"
+4 -3
View File
@@ -21,7 +21,7 @@ from collections import deque
from collections.abc import Iterable, Iterator from collections.abc import Iterable, Iterator
from pathlib import Path from pathlib import Path
from pprint import pformat from pprint import pformat
from typing import Any from typing import Any, Generic, TypeVar
import datasets import datasets
import numpy as np import numpy as np
@@ -78,6 +78,8 @@ DEFAULT_FEATURES = {
"task_index": {"dtype": "int64", "shape": (1,), "names": None}, "task_index": {"dtype": "int64", "shape": (1,), "names": None},
} }
T = TypeVar("T")
def get_parquet_file_size_in_mb(parquet_path: str | Path) -> float: def get_parquet_file_size_in_mb(parquet_path: str | Path) -> float:
metadata = pq.read_metadata(parquet_path) metadata = pq.read_metadata(parquet_path)
@@ -339,7 +341,6 @@ def write_tasks(tasks: pandas.DataFrame, local_dir: Path) -> None:
def load_tasks(local_dir: Path) -> pandas.DataFrame: def load_tasks(local_dir: Path) -> pandas.DataFrame:
tasks = pd.read_parquet(local_dir / DEFAULT_TASKS_PATH) tasks = pd.read_parquet(local_dir / DEFAULT_TASKS_PATH)
tasks.index.name = "task"
return tasks return tasks
@@ -1232,7 +1233,7 @@ class LookAheadError(Exception):
pass pass
class Backtrackable[T]: class Backtrackable(Generic[T]):
""" """
Wrap any iterator/iterable so you can step back up to `history` items Wrap any iterator/iterable so you can step back up to `history` items
and look ahead up to `lookahead` items. and look ahead up to `lookahead` items.
@@ -36,11 +36,8 @@ Convert a local dataset (works in place):
```bash ```bash
python src/lerobot/datasets/v30/convert_dataset_v21_to_v30.py \ python src/lerobot/datasets/v30/convert_dataset_v21_to_v30.py \
--repo-id=lerobot/pusht \ --repo-id=lerobot/pusht \
--root=/path/to/local/dataset/directory \ --root=/path/to/local/dataset/directory
--push-to-hub=false --push-to-hub=false
N.B. Path semantics (v2): --root is the exact dataset folder containing
meta/, data/, videos/. When omitted, defaults to $HF_LEROBOT_HOME/{repo_id}.
``` ```
""" """
@@ -108,7 +105,7 @@ episodes.jsonl
{"episode_index": 1, "tasks": ["Put the blue block in the green bowl"], "length": 266} {"episode_index": 1, "tasks": ["Put the blue block in the green bowl"], "length": 266}
NEW NEW
meta/episodes/chunk-000/file_000.parquet meta/episodes/chunk-000/episodes_000.parquet
episode_index | video_chunk_index | video_file_index | data_chunk_index | data_file_index | tasks | length episode_index | video_chunk_index | video_file_index | data_chunk_index | data_file_index | tasks | length
------------------------- -------------------------
OLD OLD
@@ -116,16 +113,15 @@ tasks.jsonl
{"task_index": 1, "task": "Put the blue block in the green bowl"} {"task_index": 1, "task": "Put the blue block in the green bowl"}
NEW NEW
meta/tasks.parquet meta/tasks/chunk-000/file_000.parquet
task_index | task task_index | task
------------------------- -------------------------
OLD OLD
episodes_stats.jsonl episodes_stats.jsonl
{"episode_index": 1, "stats": {"feature_name": {"min": ..., "max": ..., "mean": ..., "std": ..., "count": ...}}}
NEW NEW
meta/episodes/chunk-000/file_000.parquet meta/episodes_stats/chunk-000/file_000.parquet
episode_index | feature_name/min | feature_name/max | feature_name/mean | feature_name/std | feature_name/count episode_index | mean | std | min | max
------------------------- -------------------------
UPDATE UPDATE
meta/info.json meta/info.json
@@ -174,7 +170,7 @@ def convert_tasks(root, new_root):
tasks, _ = legacy_load_tasks(root) tasks, _ = legacy_load_tasks(root)
task_indices = tasks.keys() task_indices = tasks.keys()
task_strings = tasks.values() task_strings = tasks.values()
df_tasks = pd.DataFrame({"task_index": task_indices}, index=pd.Index(task_strings, name="task")) df_tasks = pd.DataFrame({"task_index": task_indices}, index=task_strings)
write_tasks(df_tasks, new_root) write_tasks(df_tasks, new_root)
@@ -205,6 +201,7 @@ def convert_data(root: Path, new_root: Path, data_file_size_in_mb: int):
image_keys = get_image_keys(root) image_keys = get_image_keys(root)
ep_idx = 0
chunk_idx = 0 chunk_idx = 0
file_idx = 0 file_idx = 0
size_in_mb = 0 size_in_mb = 0
@@ -214,23 +211,9 @@ def convert_data(root: Path, new_root: Path, data_file_size_in_mb: int):
logging.info(f"Converting data files from {len(ep_paths)} episodes") logging.info(f"Converting data files from {len(ep_paths)} episodes")
for ep_idx, ep_path in enumerate(tqdm.tqdm(ep_paths, desc="convert data files")): for ep_path in tqdm.tqdm(ep_paths, desc="convert data files"):
ep_size_in_mb = get_parquet_file_size_in_mb(ep_path) ep_size_in_mb = get_parquet_file_size_in_mb(ep_path)
ep_num_frames = get_parquet_num_frames(ep_path) ep_num_frames = get_parquet_num_frames(ep_path)
# Check if we need to start a new file BEFORE creating metadata
if size_in_mb + ep_size_in_mb >= data_file_size_in_mb and len(paths_to_cat) > 0:
# Write the accumulated data files
concat_data_files(paths_to_cat, new_root, chunk_idx, file_idx, image_keys)
# Move to next file
chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, DEFAULT_CHUNK_SIZE)
# Reset for the next file
size_in_mb = 0
paths_to_cat = []
# Now create metadata with correct chunk/file indices
ep_metadata = { ep_metadata = {
"episode_index": ep_idx, "episode_index": ep_idx,
"data/chunk_index": chunk_idx, "data/chunk_index": chunk_idx,
@@ -241,7 +224,20 @@ def convert_data(root: Path, new_root: Path, data_file_size_in_mb: int):
size_in_mb += ep_size_in_mb size_in_mb += ep_size_in_mb
num_frames += ep_num_frames num_frames += ep_num_frames
episodes_metadata.append(ep_metadata) episodes_metadata.append(ep_metadata)
paths_to_cat.append(ep_path) ep_idx += 1
if size_in_mb < data_file_size_in_mb:
paths_to_cat.append(ep_path)
continue
if paths_to_cat:
concat_data_files(paths_to_cat, new_root, chunk_idx, file_idx, image_keys)
# Reset for the next file
size_in_mb = ep_size_in_mb
paths_to_cat = [ep_path]
chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, DEFAULT_CHUNK_SIZE)
# Write remaining data if any # Write remaining data if any
if paths_to_cat: if paths_to_cat:
@@ -473,7 +469,7 @@ def convert_dataset(
# Set root based on whether local dataset path is provided # Set root based on whether local dataset path is provided
use_local_dataset = False use_local_dataset = False
root = HF_LEROBOT_HOME / repo_id if root is None else Path(root) root = HF_LEROBOT_HOME / repo_id if root is None else Path(root) / repo_id
if root.exists(): if root.exists():
validate_local_dataset_version(root) validate_local_dataset_version(root)
use_local_dataset = True use_local_dataset = True
@@ -557,7 +553,7 @@ if __name__ == "__main__":
"--root", "--root",
type=str, type=str,
default=None, default=None,
help="Local directory to use for downloading/writing the dataset. Defaults to $HF_LEROBOT_HOME/repo_id.", help="Local directory to use for downloading/writing the dataset.",
) )
parser.add_argument( parser.add_argument(
"--push-to-hub", "--push-to-hub",
+20 -26
View File
@@ -227,17 +227,16 @@ def decode_video_frames_torchvision(
min_, argmin_ = dist.min(1) min_, argmin_ = dist.min(1)
is_within_tol = min_ < tolerance_s is_within_tol = min_ < tolerance_s
if not is_within_tol.all(): assert is_within_tol.all(), (
raise FrameTimestampError( f"One or several query timestamps unexpectedly violate the tolerance ({min_[~is_within_tol]} > {tolerance_s=})."
f"One or several query timestamps unexpectedly violate the tolerance ({min_[~is_within_tol]} > {tolerance_s=})." "It means that the closest frame that can be loaded from the video is too far away in time."
" It means that the closest frame that can be loaded from the video is too far away in time." "This might be due to synchronization issues with timestamps during data collection."
" This might be due to synchronization issues with timestamps during data collection." "To be safe, we advise to ignore this item during training."
" To be safe, we advise to ignore this item during training." f"\nqueried timestamps: {query_ts}"
f"\nqueried timestamps: {query_ts}" f"\nloaded timestamps: {loaded_ts}"
f"\nloaded timestamps: {loaded_ts}" f"\nvideo: {video_path}"
f"\nvideo: {video_path}" f"\nbackend: {backend}"
f"\nbackend: {backend}" )
)
# get closest frames to the query timestamps # get closest frames to the query timestamps
closest_frames = torch.stack([loaded_frames[idx] for idx in argmin_]) closest_frames = torch.stack([loaded_frames[idx] for idx in argmin_])
@@ -249,11 +248,7 @@ def decode_video_frames_torchvision(
# convert to the pytorch format which is float32 in [0,1] range (and channel first) # convert to the pytorch format which is float32 in [0,1] range (and channel first)
closest_frames = closest_frames.type(torch.float32) / 255 closest_frames = closest_frames.type(torch.float32) / 255
if len(timestamps) != len(closest_frames): assert len(timestamps) == len(closest_frames)
raise FrameTimestampError(
f"Number of retrieved frames ({len(closest_frames)}) does not match "
f"number of queried timestamps ({len(timestamps)})"
)
return closest_frames return closest_frames
@@ -358,16 +353,15 @@ def decode_video_frames_torchcodec(
min_, argmin_ = dist.min(1) min_, argmin_ = dist.min(1)
is_within_tol = min_ < tolerance_s is_within_tol = min_ < tolerance_s
if not is_within_tol.all(): assert is_within_tol.all(), (
raise FrameTimestampError( f"One or several query timestamps unexpectedly violate the tolerance ({min_[~is_within_tol]} > {tolerance_s=})."
f"One or several query timestamps unexpectedly violate the tolerance ({min_[~is_within_tol]} > {tolerance_s=})." "It means that the closest frame that can be loaded from the video is too far away in time."
" It means that the closest frame that can be loaded from the video is too far away in time." "This might be due to synchronization issues with timestamps during data collection."
" This might be due to synchronization issues with timestamps during data collection." "To be safe, we advise to ignore this item during training."
" To be safe, we advise to ignore this item during training." f"\nqueried timestamps: {query_ts}"
f"\nqueried timestamps: {query_ts}" f"\nloaded timestamps: {loaded_ts}"
f"\nloaded timestamps: {loaded_ts}" f"\nvideo: {video_path}"
f"\nvideo: {video_path}" )
)
# get closest frames to the query timestamps # get closest frames to the query timestamps
closest_frames = torch.stack([loaded_frames[idx] for idx in argmin_]) closest_frames = torch.stack([loaded_frames[idx] for idx in argmin_])
-99
View File
@@ -346,105 +346,6 @@ class LiberoEnv(EnvConfig):
return kwargs return kwargs
@EnvConfig.register_subclass("libero_plus")
@dataclass
class LiberoPlusEnv(LiberoEnv):
"""Alias config for LIBERO-plus benchmarks.
LIBERO-plus keeps the same Python package/module names as LIBERO, so this
config reuses the existing LIBERO env implementation while making intent explicit
in experiment configs (`env.type=libero_plus`).
"""
task: str = "libero_spatial"
@EnvConfig.register_subclass("robocasa")
@dataclass
class RoboCasaEnv(EnvConfig):
"""RoboCasa kitchen composite-task environments.
Wraps ``robocasa.wrappers.gym_wrapper.RoboCasaGymEnv`` with a flat 12-D Box
action space and a structured pixel + state observation dict.
Selected benchmark tasks (3 short + 2 long):
Short: PickPlaceCounterToCabinet, PrepareToast, CoffeeSetupMug
Long: PrepareCoffee, RestockPantry
"""
task: str = "PickPlaceCounterToCabinet"
tasks: list[str] | None = None # multi-task: list of task names (without robocasa/ prefix)
fps: int = 20
episode_length: int = 500
image_size: int = 128
split: str = "target" # "pretrain" or "target"
features: dict[str, PolicyFeature] = field(
default_factory=lambda: {
ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(12,)),
}
)
features_map: dict[str, str] = field(
default_factory=lambda: {
ACTION: ACTION,
"agentview_left": f"{OBS_IMAGES}.agentview_left",
"agentview_right": f"{OBS_IMAGES}.agentview_right",
"eye_in_hand": f"{OBS_IMAGES}.eye_in_hand",
"robot_state": OBS_STATE,
}
)
def __post_init__(self):
for cam in ("agentview_left", "agentview_right", "eye_in_hand"):
self.features[cam] = PolicyFeature(
type=FeatureType.VISUAL, shape=(self.image_size, self.image_size, 3)
)
self.features["robot_state"] = PolicyFeature(type=FeatureType.STATE, shape=(16,))
@property
def gym_kwargs(self) -> dict:
return {"split": self.split}
@EnvConfig.register_subclass("robomme")
@dataclass
class RoboMMEEnv(EnvConfig):
"""RoboMME memory-augmented manipulation benchmark (ManiSkill/SAPIEN).
16 tasks across 4 suites: Counting, Permanence, Reference, Imitation.
Uses BenchmarkEnvBuilder from the robomme package.
"""
task: str = "PickXtimes"
fps: int = 10
episode_length: int = 300
action_space: str = "joint_angle"
dataset_split: str = "test"
task_ids: list[int] | None = None
features: dict[str, PolicyFeature] = field(
default_factory=lambda: {
ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(8,)),
"front_rgb": PolicyFeature(type=FeatureType.VISUAL, shape=(256, 256, 3)),
"wrist_rgb": PolicyFeature(type=FeatureType.VISUAL, shape=(256, 256, 3)),
OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(8,)),
}
)
features_map: dict[str, str] = field(
default_factory=lambda: {
ACTION: ACTION,
"front_rgb": f"{OBS_IMAGES}.front",
"wrist_rgb": f"{OBS_IMAGES}.wrist",
OBS_STATE: OBS_STATE,
}
)
@property
def gym_kwargs(self) -> dict:
return {
"action_space": self.action_space,
"dataset": self.dataset_split,
}
@EnvConfig.register_subclass("metaworld") @EnvConfig.register_subclass("metaworld")
@dataclass @dataclass
class MetaworldEnv(EnvConfig): class MetaworldEnv(EnvConfig):
+3 -50
View File
@@ -20,21 +20,11 @@ import gymnasium as gym
from gymnasium.envs.registration import registry as gym_registry from gymnasium.envs.registration import registry as gym_registry
from lerobot.configs.policies import PreTrainedConfig from lerobot.configs.policies import PreTrainedConfig
from lerobot.envs.configs import ( from lerobot.envs.configs import AlohaEnv, EnvConfig, HubEnvConfig, IsaaclabArenaEnv, LiberoEnv, PushtEnv
AlohaEnv,
EnvConfig,
HubEnvConfig,
IsaaclabArenaEnv,
LiberoEnv,
LiberoPlusEnv,
PushtEnv,
RoboCasaEnv,
RoboMMEEnv,
)
from lerobot.envs.utils import _call_make_env, _download_hub_file, _import_hub_module, _normalize_hub_result from lerobot.envs.utils import _call_make_env, _download_hub_file, _import_hub_module, _normalize_hub_result
from lerobot.policies.xvla.configuration_xvla import XVLAConfig from lerobot.policies.xvla.configuration_xvla import XVLAConfig
from lerobot.processor import ProcessorStep from lerobot.processor import ProcessorStep
from lerobot.processor.env_processor import IsaaclabArenaProcessorStep, LiberoProcessorStep, RoboCasaProcessorStep from lerobot.processor.env_processor import IsaaclabArenaProcessorStep, LiberoProcessorStep
from lerobot.processor.pipeline import PolicyProcessorPipeline from lerobot.processor.pipeline import PolicyProcessorPipeline
@@ -45,12 +35,6 @@ def make_env_config(env_type: str, **kwargs) -> EnvConfig:
return PushtEnv(**kwargs) return PushtEnv(**kwargs)
elif env_type == "libero": elif env_type == "libero":
return LiberoEnv(**kwargs) return LiberoEnv(**kwargs)
elif env_type == "libero_plus":
return LiberoPlusEnv(**kwargs)
elif env_type == "robocasa":
return RoboCasaEnv(**kwargs)
elif env_type == "robomme":
return RoboMMEEnv(**kwargs)
else: else:
raise ValueError(f"Policy type '{env_type}' is not available.") raise ValueError(f"Policy type '{env_type}' is not available.")
@@ -86,13 +70,9 @@ def make_env_pre_post_processors(
return make_xvla_libero_pre_post_processors() return make_xvla_libero_pre_post_processors()
# For LIBERO environments, add the LiberoProcessorStep to preprocessor # For LIBERO environments, add the LiberoProcessorStep to preprocessor
if isinstance(env_cfg, (LiberoEnv, LiberoPlusEnv)) or "libero" in env_cfg.type: if isinstance(env_cfg, LiberoEnv) or "libero" in env_cfg.type:
preprocessor_steps.append(LiberoProcessorStep()) preprocessor_steps.append(LiberoProcessorStep())
# For RoboCasa environments, add the RoboCasaProcessorStep to preprocessor
if isinstance(env_cfg, RoboCasaEnv) or "robocasa" in env_cfg.type:
preprocessor_steps.append(RoboCasaProcessorStep())
# For Isaaclab Arena environments, add the IsaaclabArenaProcessorStep # For Isaaclab Arena environments, add the IsaaclabArenaProcessorStep
if isinstance(env_cfg, IsaaclabArenaEnv) or "isaaclab_arena" in env_cfg.type: if isinstance(env_cfg, IsaaclabArenaEnv) or "isaaclab_arena" in env_cfg.type:
# Parse comma-separated keys (handle None for state-based policies) # Parse comma-separated keys (handle None for state-based policies)
@@ -201,33 +181,6 @@ def make_env(
control_mode=cfg.control_mode, control_mode=cfg.control_mode,
episode_length=cfg.episode_length, episode_length=cfg.episode_length,
) )
elif "robocasa" in cfg.type:
from lerobot.envs.robocasa import create_robocasa_envs
tasks = cfg.tasks if cfg.tasks else [cfg.task]
return create_robocasa_envs(
tasks=tasks,
n_envs=n_envs,
image_size=cfg.image_size,
split=cfg.split,
episode_length=cfg.episode_length,
gym_kwargs=cfg.gym_kwargs,
env_cls=env_cls,
)
elif "robomme" in cfg.type:
from lerobot.envs.robomme import create_robomme_envs
return create_robomme_envs(
task=cfg.task,
n_envs=n_envs,
action_space_type=cfg.action_space,
dataset=cfg.dataset_split,
episode_length=cfg.episode_length,
task_ids=cfg.task_ids,
env_cls=env_cls,
)
elif "metaworld" in cfg.type: elif "metaworld" in cfg.type:
from lerobot.envs.metaworld import create_metaworld_envs from lerobot.envs.metaworld import create_metaworld_envs
+2 -8
View File
@@ -26,14 +26,8 @@ import gymnasium as gym
import numpy as np import numpy as np
import torch import torch
from gymnasium import spaces from gymnasium import spaces
from libero.libero import benchmark, get_libero_path
try: from libero.libero.envs import OffScreenRenderEnv
from libero.libero import benchmark, get_libero_path
from libero.libero.envs import OffScreenRenderEnv
except ImportError:
# LIBERO-plus may be installed from source with an extra nested package level.
from libero.libero.libero import benchmark, get_libero_path
from libero.libero.libero.envs import OffScreenRenderEnv
from lerobot.processor import RobotObservation from lerobot.processor import RobotObservation
-273
View File
@@ -1,273 +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.
from __future__ import annotations
from collections import defaultdict
from collections.abc import Callable, Sequence
from functools import partial
from typing import Any
import gymnasium as gym
import numpy as np
from gymnasium import spaces
# Action layout (flat 12D, normalized to [-1, 1]):
# [0:3] end_effector_position (delta x, y, z)
# [3:6] end_effector_rotation (delta roll, pitch, yaw)
# [6:7] gripper_close (open=-1, close=+1)
# [7:11] base_motion (x, y, theta, torso_height)
# [11:12] control_mode (arm=-1, base=+1)
ACTION_DIM = 12
ACTION_LOW = -1.0
ACTION_HIGH = 1.0
# Proprioceptive state layout (flat 16D):
# [0:2] gripper_qpos
# [2:5] base_position
# [5:9] base_rotation (quaternion)
# [9:12] end_effector_position_relative
# [12:16] end_effector_rotation_relative (quaternion)
STATE_DIM = 16
# Obs dict keys from RoboCasaGymEnv.get_observation()
_CAM_KEYS = (
"video.robot0_agentview_left",
"video.robot0_agentview_right",
"video.robot0_eye_in_hand",
)
_STATE_KEYS_ORDERED = (
"state.gripper_qpos", # (2,)
"state.base_position", # (3,)
"state.base_rotation", # (4,)
"state.end_effector_position_relative", # (3,)
"state.end_effector_rotation_relative", # (4,)
)
# Mapping from video.* key → short image name used in features_map
CAM_KEY_TO_NAME = {
"video.robot0_agentview_left": "agentview_left",
"video.robot0_agentview_right": "agentview_right",
"video.robot0_eye_in_hand": "eye_in_hand",
}
def _flat_to_action_dict(flat: np.ndarray) -> dict[str, np.ndarray]:
"""Convert a 12D flat action array to the Dict format expected by RoboCasaGymEnv."""
return {
"action.end_effector_position": flat[0:3],
"action.end_effector_rotation": flat[3:6],
"action.gripper_close": flat[6:7],
"action.base_motion": flat[7:11],
"action.control_mode": flat[11:12],
}
class RoboCasaEnv(gym.Env):
"""Thin wrapper around RoboCasaGymEnv that provides a flat Box action space
and a structured observation dict compatible with LeRobot policies.
Observations returned by step/reset:
{
"pixels": {
"agentview_left": (H, W, 3) uint8,
"agentview_right": (H, W, 3) uint8,
"eye_in_hand": (H, W, 3) uint8,
},
"robot_state": (16,) float32,
}
Actions: flat float32 ndarray of shape (12,), normalized to [-1, 1].
"""
metadata = {"render_modes": ["rgb_array"], "render_fps": 20}
def __init__(
self,
task: str,
split: str = "target",
image_size: int = 128,
render_mode: str = "rgb_array",
episode_length: int = 500,
**gym_kwargs: Any,
):
super().__init__()
# Lazy import — robocasa is optional
import robocasa.environments # noqa: F401 — registers all gym envs
self.task = task
self.render_mode = render_mode
self.image_size = image_size
self._max_episode_steps = episode_length
self._step_count = 0
self._env = gym.make(
f"robocasa/{task}",
split=split,
camera_widths=image_size,
camera_heights=image_size,
**gym_kwargs,
)
# Flat 12D Box action space
self.action_space = spaces.Box(
low=ACTION_LOW,
high=ACTION_HIGH,
shape=(ACTION_DIM,),
dtype=np.float32,
)
images = {
name: spaces.Box(low=0, high=255, shape=(image_size, image_size, 3), dtype=np.uint8)
for name in CAM_KEY_TO_NAME.values()
}
self.observation_space = spaces.Dict(
{
"pixels": spaces.Dict(images),
"robot_state": spaces.Box(
low=-np.inf, high=np.inf, shape=(STATE_DIM,), dtype=np.float32
),
}
)
def _format_obs(self, raw_obs: dict) -> dict:
pixels = {
CAM_KEY_TO_NAME[k]: raw_obs[k]
for k in _CAM_KEYS
if k in raw_obs
}
state_parts = [
np.asarray(raw_obs[k], dtype=np.float32)
for k in _STATE_KEYS_ORDERED
if k in raw_obs
]
robot_state = np.concatenate(state_parts) if state_parts else np.zeros(STATE_DIM, dtype=np.float32)
return {"pixels": pixels, "robot_state": robot_state}
def reset(self, seed: int | None = None, **kwargs) -> tuple[dict, dict]:
super().reset(seed=seed)
self._step_count = 0
raw_obs, info = self._env.reset(seed=seed)
info.setdefault("is_success", False)
info["task"] = self.task
return self._format_obs(raw_obs), info
def step(self, action: np.ndarray) -> tuple[dict, float, bool, bool, dict]:
if action.ndim != 1 or action.shape[0] != ACTION_DIM:
raise ValueError(
f"Expected 1-D action of shape ({ACTION_DIM},), got {action.shape}"
)
action_dict = _flat_to_action_dict(action)
raw_obs, reward, terminated, truncated, info = self._env.step(action_dict)
self._step_count += 1
is_success = bool(info.get("success", False))
terminated = terminated or is_success
if self._step_count >= self._max_episode_steps:
truncated = True
info.update({"task": self.task, "is_success": is_success})
obs = self._format_obs(raw_obs)
if terminated or truncated:
info["final_info"] = {"task": self.task, "is_success": is_success}
return obs, reward, terminated, truncated, info
def render(self) -> np.ndarray | None:
if self.render_mode == "rgb_array":
return self._env.render()
return None
def close(self) -> None:
self._env.close()
def _make_env_fns(
*,
task: str,
n_envs: int,
image_size: int,
split: str,
episode_length: int,
gym_kwargs: dict[str, Any],
) -> list[Callable[[], RoboCasaEnv]]:
"""Build n_envs factory callables for a single task."""
def _make(episode_index: int) -> RoboCasaEnv: # noqa: ARG001
return RoboCasaEnv(
task=task,
split=split,
image_size=image_size,
episode_length=episode_length,
**gym_kwargs,
)
return [partial(_make, i) for i in range(n_envs)]
def create_robocasa_envs(
tasks: str | Sequence[str],
n_envs: int,
image_size: int = 128,
split: str = "target",
episode_length: int = 500,
gym_kwargs: dict[str, Any] | None = None,
env_cls: Callable[[Sequence[Callable[[], Any]]], Any] | None = None,
) -> dict[str, dict[int, Any]]:
"""Create vectorized RoboCasa environments.
Args:
tasks: A single task name or list of task names (without "robocasa/" prefix).
E.g. "PickPlaceCounterToCabinet" or ["BoilPot", "PrepareCoffee"].
n_envs: Number of parallel envs per task.
image_size: Square image resolution for all cameras.
split: RoboCasa dataset split — "pretrain" or "target".
episode_length: Max steps per episode before truncation.
gym_kwargs: Extra kwargs forwarded to each RoboCasaEnv.
env_cls: Callable to wrap list of factory fns (SyncVectorEnv or AsyncVectorEnv).
Returns:
dict[task_name][task_id=0] -> vec_env
"""
if env_cls is None or not callable(env_cls):
raise ValueError("env_cls must be a callable wrapping a list of env factory callables.")
if not isinstance(n_envs, int) or n_envs <= 0:
raise ValueError(f"n_envs must be a positive int; got {n_envs}.")
if isinstance(tasks, str):
task_list = [t.strip() for t in tasks.split(",") if t.strip()]
else:
task_list = [str(t).strip() for t in tasks if str(t).strip()]
if not task_list:
raise ValueError("`tasks` must contain at least one task name.")
gym_kwargs = dict(gym_kwargs or {})
out: dict[str, dict[int, Any]] = defaultdict(dict)
print(f"Creating RoboCasa envs | tasks={task_list} | n_envs(per task)={n_envs} | split={split}")
for task in task_list:
fns = _make_env_fns(
task=task,
n_envs=n_envs,
image_size=image_size,
split=split,
episode_length=episode_length,
gym_kwargs=gym_kwargs,
)
out["robocasa"][len(out["robocasa"])] = env_cls(fns)
print(f" Built vec env | task={task} | n_envs={n_envs}")
return {suite: dict(task_map) for suite, task_map in out.items()}
-154
View File
@@ -1,154 +0,0 @@
"""RoboMME environment wrapper for LeRobot evaluation.
Wraps the RoboMME ``BenchmarkEnvBuilder`` into a Gymnasium-compatible
``VectorEnv`` suitable for ``lerobot_eval``.
RoboMME tasks:
Counting: BinFill, PickXtimes, SwingXtimes, StopCube
Permanence: VideoUnmask, VideoUnmaskSwap, ButtonUnmask, ButtonUnmaskSwap
Reference: PickHighlight, VideoRepick, VideoPlaceButton, VideoPlaceOrder
Imitation: MoveCube, InsertPeg, PatternLock, RouteStick
Install: pip install robomme (or from source: https://github.com/RoboMME/robomme_benchmark)
"""
from __future__ import annotations
from typing import Any
import gymnasium as gym
import numpy as np
from gymnasium import spaces
ROBOMME_TASKS = [
"BinFill", "PickXtimes", "SwingXtimes", "StopCube",
"VideoUnmask", "VideoUnmaskSwap", "ButtonUnmask", "ButtonUnmaskSwap",
"PickHighlight", "VideoRepick", "VideoPlaceButton", "VideoPlaceOrder",
"MoveCube", "InsertPeg", "PatternLock", "RouteStick",
]
class RoboMMEGymEnv(gym.Env):
"""Thin Gymnasium wrapper around a single RoboMME episode env."""
metadata = {"render_modes": ["rgb_array"]}
def __init__(
self,
task: str = "PickXtimes",
action_space_type: str = "joint_angle",
dataset: str = "test",
episode_idx: int = 0,
max_steps: int = 300,
):
super().__init__()
from robomme.env_record_wrapper import BenchmarkEnvBuilder
self._task = task
self._action_space_type = action_space_type
self._dataset = dataset
self._episode_idx = episode_idx
self._max_steps = max_steps
self._builder = BenchmarkEnvBuilder(
env_id=task,
dataset=dataset,
action_space=action_space_type,
gui_render=False,
max_steps=max_steps,
)
self._env = None
action_dim = 8 if action_space_type == "joint_angle" else 7
self.action_space = spaces.Box(low=-1.0, high=1.0, shape=(action_dim,), dtype=np.float32)
self.observation_space = spaces.Dict({
"front_rgb": spaces.Box(0, 255, shape=(256, 256, 3), dtype=np.uint8),
"wrist_rgb": spaces.Box(0, 255, shape=(256, 256, 3), dtype=np.uint8),
"state": spaces.Box(-np.inf, np.inf, shape=(8,), dtype=np.float32),
})
def reset(self, *, seed=None, options=None):
super().reset(seed=seed)
self._env = self._builder.make_env_for_episode(
episode_idx=self._episode_idx, max_steps=self._max_steps,
)
obs, info = self._env.reset()
return self._convert_obs(obs), self._convert_info(info)
def step(self, action):
obs, reward, terminated, truncated, info = self._env.step(action)
terminated_bool = bool(terminated.item()) if hasattr(terminated, "item") else bool(terminated)
truncated_bool = bool(truncated.item()) if hasattr(truncated, "item") else bool(truncated)
status = info.get("status", "ongoing")
is_success = status == "success"
conv_info = self._convert_info(info)
conv_info["is_success"] = is_success
return self._convert_obs(obs), float(reward), terminated_bool, truncated_bool, conv_info
def _convert_obs(self, obs: dict) -> dict:
front_rgb = obs["front_rgb_list"][-1] if isinstance(obs["front_rgb_list"], list) else obs["front_rgb_list"]
wrist_rgb = obs["wrist_rgb_list"][-1] if isinstance(obs["wrist_rgb_list"], list) else obs["wrist_rgb_list"]
joint_state = obs["joint_state_list"][-1] if isinstance(obs["joint_state_list"], list) else obs["joint_state_list"]
gripper_state = obs["gripper_state_list"][-1] if isinstance(obs["gripper_state_list"], list) else obs["gripper_state_list"]
front_rgb = np.asarray(front_rgb, dtype=np.uint8)
wrist_rgb = np.asarray(wrist_rgb, dtype=np.uint8)
joint = np.asarray(joint_state, dtype=np.float32).flatten()[:7]
gripper = np.asarray(gripper_state, dtype=np.float32).flatten()[:1]
state = np.concatenate([joint, gripper])
return {
"front_rgb": front_rgb,
"wrist_rgb": wrist_rgb,
"state": state,
}
def _convert_info(self, info: dict) -> dict:
return {
"status": info.get("status", "ongoing"),
"task_goal": info.get("task_goal", ""),
}
def create_robomme_envs(
task: str,
n_envs: int = 1,
action_space_type: str = "joint_angle",
dataset: str = "test",
episode_length: int = 300,
task_ids: list[int] | None = None,
env_cls=None,
) -> dict[str, dict[int, gym.vector.VectorEnv]]:
"""Create vectorized RoboMME environments for evaluation.
Returns {suite_name: {task_id: VectorEnv}} matching lerobot's expected format.
"""
if env_cls is None:
env_cls = gym.vector.SyncVectorEnv
if task_ids is None:
task_ids = [0]
suite_name = "robomme"
envs_by_task = {}
for task_id in task_ids:
def _make_one(ep_idx=task_id):
return RoboMMEGymEnv(
task=task,
action_space_type=action_space_type,
dataset=dataset,
episode_idx=ep_idx,
max_steps=episode_length,
)
vec = env_cls(
[_make_one for _ in range(n_envs)],
autoreset_mode=gym.vector.AutoresetMode.SAME_STEP,
)
envs_by_task[task_id] = vec
return {suite_name: envs_by_task}
+4 -4
View File
@@ -29,7 +29,7 @@ from dataclasses import dataclass
from enum import Enum from enum import Enum
from functools import cached_property from functools import cached_property
from pprint import pformat from pprint import pformat
from typing import Protocol from typing import Protocol, TypeAlias
import serial import serial
from deepdiff import DeepDiff from deepdiff import DeepDiff
@@ -38,8 +38,8 @@ from tqdm import tqdm
from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
from lerobot.utils.utils import enter_pressed, move_cursor_up from lerobot.utils.utils import enter_pressed, move_cursor_up
type NameOrID = str | int NameOrID: TypeAlias = str | int
type Value = int | float Value: TypeAlias = int | float
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -1277,4 +1277,4 @@ class SerialMotorsBus(MotorsBusBase):
# Backward compatibility alias # Backward compatibility alias
MotorsBus = SerialMotorsBus MotorsBus: TypeAlias = SerialMotorsBus
@@ -55,16 +55,10 @@ class DiffusionConfig(PreTrainedConfig):
normalization_mapping: A dictionary that maps from a str value of FeatureType (e.g., "STATE", "VISUAL") to normalization_mapping: A dictionary that maps from a str value of FeatureType (e.g., "STATE", "VISUAL") to
a corresponding NormalizationMode (e.g., NormalizationMode.MIN_MAX) a corresponding NormalizationMode (e.g., NormalizationMode.MIN_MAX)
vision_backbone: Name of the torchvision resnet backbone to use for encoding images. vision_backbone: Name of the torchvision resnet backbone to use for encoding images.
resize_shape: (H, W) shape to resize images to as a preprocessing step for the vision crop_shape: (H, W) shape to crop images to as a preprocessing step for the vision backbone. Must fit
backbone. If None, no resizing is done and the original image resolution is used. within the image size. If None, no cropping is done.
crop_ratio: Ratio in (0, 1] used to derive the crop size from resize_shape crop_is_random: Whether the crop should be random at training time (it's always a center crop in eval
(crop_h = int(resize_shape[0] * crop_ratio), likewise for width). mode).
Set to 1.0 to disable cropping. Only takes effect when resize_shape is not None.
crop_shape: (H, W) shape to crop images to. When resize_shape is set and crop_ratio < 1.0,
this is computed automatically. Can also be set directly for legacy configs that use
crop-only (without resize). If None and no derivation applies, no cropping is done.
crop_is_random: Whether the crop should be random at training time (it's always a center
crop in eval mode).
pretrained_backbone_weights: Pretrained weights from torchvision to initialize the backbone. pretrained_backbone_weights: Pretrained weights from torchvision to initialize the backbone.
`None` means no pretrained weights. `None` means no pretrained weights.
use_group_norm: Whether to replace batch normalization with group normalization in the backbone. use_group_norm: Whether to replace batch normalization with group normalization in the backbone.
@@ -120,9 +114,7 @@ class DiffusionConfig(PreTrainedConfig):
# Architecture / modeling. # Architecture / modeling.
# Vision backbone. # Vision backbone.
vision_backbone: str = "resnet18" vision_backbone: str = "resnet18"
resize_shape: tuple[int, int] | None = None crop_shape: tuple[int, int] | None = (84, 84)
crop_ratio: float = 1.0
crop_shape: tuple[int, int] | None = None
crop_is_random: bool = True crop_is_random: bool = True
pretrained_backbone_weights: str | None = None pretrained_backbone_weights: str | None = None
use_group_norm: bool = True use_group_norm: bool = True
@@ -147,10 +139,6 @@ class DiffusionConfig(PreTrainedConfig):
# Inference # Inference
num_inference_steps: int | None = None num_inference_steps: int | None = None
# Optimization
compile_model: bool = False
compile_mode: str = "reduce-overhead"
# Loss computation # Loss computation
do_mask_loss_for_padding: bool = False do_mask_loss_for_padding: bool = False
@@ -183,25 +171,6 @@ class DiffusionConfig(PreTrainedConfig):
f"Got {self.noise_scheduler_type}." f"Got {self.noise_scheduler_type}."
) )
if self.resize_shape is not None and (
len(self.resize_shape) != 2 or any(d <= 0 for d in self.resize_shape)
):
raise ValueError(f"`resize_shape` must be a pair of positive integers. Got {self.resize_shape}.")
if not (0 < self.crop_ratio <= 1.0):
raise ValueError(f"`crop_ratio` must be in (0, 1]. Got {self.crop_ratio}.")
if self.resize_shape is not None:
if self.crop_ratio < 1.0:
self.crop_shape = (
int(self.resize_shape[0] * self.crop_ratio),
int(self.resize_shape[1] * self.crop_ratio),
)
else:
# Explicitly disable cropping for resize+ratio path when crop_ratio == 1.0.
self.crop_shape = None
if self.crop_shape is not None and (self.crop_shape[0] <= 0 or self.crop_shape[1] <= 0):
raise ValueError(f"`crop_shape` must have positive dimensions. Got {self.crop_shape}.")
# Check that the horizon size and U-Net downsampling is compatible. # Check that the horizon size and U-Net downsampling is compatible.
# U-Net downsamples by 2 with each stage. # U-Net downsamples by 2 with each stage.
downsampling_factor = 2 ** len(self.down_dims) downsampling_factor = 2 ** len(self.down_dims)
@@ -229,12 +198,13 @@ class DiffusionConfig(PreTrainedConfig):
if len(self.image_features) == 0 and self.env_state_feature is None: if len(self.image_features) == 0 and self.env_state_feature is None:
raise ValueError("You must provide at least one image or the environment state among the inputs.") raise ValueError("You must provide at least one image or the environment state among the inputs.")
if self.resize_shape is None and self.crop_shape is not None: if self.crop_shape is not None:
for key, image_ft in self.image_features.items(): for key, image_ft in self.image_features.items():
if self.crop_shape[0] > image_ft.shape[1] or self.crop_shape[1] > image_ft.shape[2]: if self.crop_shape[0] > image_ft.shape[1] or self.crop_shape[1] > image_ft.shape[2]:
raise ValueError( raise ValueError(
f"`crop_shape` should fit within the image shapes. Got {self.crop_shape} " f"`crop_shape` should fit within the images shapes. Got {self.crop_shape} "
f"for `crop_shape` and {image_ft.shape} for `{key}`." f"for `crop_shape` and {image_ft.shape} for "
f"`{key}`."
) )
# Check that all input images have the same shape. # Check that all input images have the same shape.
@@ -142,9 +142,6 @@ class DiffusionPolicy(PreTrainedPolicy):
"""Run the batch through the model and compute the loss for training or validation.""" """Run the batch through the model and compute the loss for training or validation."""
if self.config.image_features: if self.config.image_features:
batch = dict(batch) # shallow copy so that adding a key doesn't modify the original batch = dict(batch) # shallow copy so that adding a key doesn't modify the original
for key in self.config.image_features:
if self.config.n_obs_steps == 1 and batch[key].ndim == 4:
batch[key] = batch[key].unsqueeze(1)
batch[OBS_IMAGES] = torch.stack([batch[key] for key in self.config.image_features], dim=-4) batch[OBS_IMAGES] = torch.stack([batch[key] for key in self.config.image_features], dim=-4)
loss = self.diffusion.compute_loss(batch) loss = self.diffusion.compute_loss(batch)
# no output_dict so returning None # no output_dict so returning None
@@ -185,11 +182,6 @@ class DiffusionModel(nn.Module):
self.unet = DiffusionConditionalUnet1d(config, global_cond_dim=global_cond_dim * config.n_obs_steps) self.unet = DiffusionConditionalUnet1d(config, global_cond_dim=global_cond_dim * config.n_obs_steps)
if config.compile_model:
# Compile the U-Net. "reduce-overhead" is preferred for the small-batch repetitive loops
# common in diffusion inference.
self.unet = torch.compile(self.unet, mode=config.compile_mode)
self.noise_scheduler = _make_noise_scheduler( self.noise_scheduler = _make_noise_scheduler(
config.noise_scheduler_type, config.noise_scheduler_type,
num_train_timesteps=config.num_train_timesteps, num_train_timesteps=config.num_train_timesteps,
@@ -454,18 +446,12 @@ class DiffusionRgbEncoder(nn.Module):
def __init__(self, config: DiffusionConfig): def __init__(self, config: DiffusionConfig):
super().__init__() super().__init__()
# Set up optional preprocessing. # Set up optional preprocessing.
if config.resize_shape is not None: if config.crop_shape is not None:
self.resize = torchvision.transforms.Resize(config.resize_shape)
else:
self.resize = None
crop_shape = config.crop_shape
if crop_shape is not None:
self.do_crop = True self.do_crop = True
# Always use center crop for eval # Always use center crop for eval
self.center_crop = torchvision.transforms.CenterCrop(crop_shape) self.center_crop = torchvision.transforms.CenterCrop(config.crop_shape)
if config.crop_is_random: if config.crop_is_random:
self.maybe_random_crop = torchvision.transforms.RandomCrop(crop_shape) self.maybe_random_crop = torchvision.transforms.RandomCrop(config.crop_shape)
else: else:
self.maybe_random_crop = self.center_crop self.maybe_random_crop = self.center_crop
else: else:
@@ -491,16 +477,13 @@ class DiffusionRgbEncoder(nn.Module):
# Set up pooling and final layers. # Set up pooling and final layers.
# Use a dry run to get the feature map shape. # Use a dry run to get the feature map shape.
# The dummy shape mirrors the runtime preprocessing order: resize -> crop. # The dummy input should take the number of image channels from `config.image_features` and it should
# use the height and width from `config.crop_shape` if it is provided, otherwise it should use the
# height and width from `config.image_features`.
# Note: we have a check in the config class to make sure all images have the same shape. # Note: we have a check in the config class to make sure all images have the same shape.
images_shape = next(iter(config.image_features.values())).shape images_shape = next(iter(config.image_features.values())).shape
if config.crop_shape is not None: dummy_shape_h_w = config.crop_shape if config.crop_shape is not None else images_shape[1:]
dummy_shape_h_w = config.crop_shape
elif config.resize_shape is not None:
dummy_shape_h_w = config.resize_shape
else:
dummy_shape_h_w = images_shape[1:]
dummy_shape = (1, images_shape[0], *dummy_shape_h_w) dummy_shape = (1, images_shape[0], *dummy_shape_h_w)
feature_map_shape = get_output_shape(self.backbone, dummy_shape)[1:] feature_map_shape = get_output_shape(self.backbone, dummy_shape)[1:]
@@ -516,10 +499,7 @@ class DiffusionRgbEncoder(nn.Module):
Returns: Returns:
(B, D) image feature. (B, D) image feature.
""" """
# Preprocess: resize if configured, then crop if configured. # Preprocess: maybe crop (if it was set up in the __init__).
if self.resize is not None:
x = self.resize(x)
if self.do_crop: if self.do_crop:
if self.training: # noqa: SIM108 if self.training: # noqa: SIM108
x = self.maybe_random_crop(x) x = self.maybe_random_crop(x)
+2 -1
View File
@@ -18,9 +18,10 @@ from __future__ import annotations
import importlib import importlib
import logging import logging
from typing import Any, TypedDict, Unpack from typing import Any, TypedDict
import torch import torch
from typing_extensions import Unpack
from lerobot.configs.policies import PreTrainedConfig from lerobot.configs.policies import PreTrainedConfig
from lerobot.configs.types import FeatureType from lerobot.configs.types import FeatureType
@@ -4,9 +4,10 @@
# Licensed under The MIT License [see LICENSE for details] # Licensed under The MIT License [see LICENSE for details]
# -------------------------------------------------------- # --------------------------------------------------------
from __future__ import annotations
# copy from https://github.com/huggingface/transformers/blob/main/src/transformers/models/llava_onevision/image_processing_llava_onevision_fast.py # copy from https://github.com/huggingface/transformers/blob/main/src/transformers/models/llava_onevision/image_processing_llava_onevision_fast.py
from typing import Optional
from transformers.image_processing_utils import ( from transformers.image_processing_utils import (
BatchFeature, BatchFeature,
get_patch_output_size, get_patch_output_size,
@@ -164,11 +165,11 @@ class Eagle25VLImageProcessorFast(BaseImageProcessorFast):
def _resize_for_patching( def _resize_for_patching(
self, self,
image: torch.Tensor, image: "torch.Tensor",
target_resolution: tuple, target_resolution: tuple,
interpolation: F.InterpolationMode, interpolation: "F.InterpolationMode",
input_data_format: ChannelDimension, input_data_format: ChannelDimension,
) -> torch.Tensor: ) -> "torch.Tensor":
""" """
Resizes an image to a target resolution while maintaining aspect ratio. Resizes an image to a target resolution while maintaining aspect ratio.
@@ -218,8 +219,8 @@ class Eagle25VLImageProcessorFast(BaseImageProcessorFast):
return best_ratio return best_ratio
def _pad_for_patching( def _pad_for_patching(
self, image: torch.Tensor, target_resolution: tuple, input_data_format: ChannelDimension self, image: "torch.Tensor", target_resolution: tuple, input_data_format: ChannelDimension
) -> torch.Tensor: ) -> "torch.Tensor":
""" """
Pad an image to a target resolution while maintaining aspect ratio. Pad an image to a target resolution while maintaining aspect ratio.
""" """
@@ -235,15 +236,15 @@ class Eagle25VLImageProcessorFast(BaseImageProcessorFast):
def _get_image_patches( def _get_image_patches(
self, self,
image: torch.Tensor, image: "torch.Tensor",
min_num: int, min_num: int,
max_num: int, max_num: int,
size: tuple, size: tuple,
tile_size: int, tile_size: int,
use_thumbnail: bool, use_thumbnail: bool,
interpolation: F.InterpolationMode, interpolation: "F.InterpolationMode",
pad_during_tiling: bool, pad_during_tiling: bool,
) -> list[torch.Tensor]: ) -> list["torch.Tensor"]:
image_size = get_image_size(image, channel_dim=ChannelDimension.FIRST) image_size = get_image_size(image, channel_dim=ChannelDimension.FIRST)
orig_height, orig_width = image_size orig_height, orig_width = image_size
aspect_ratio = orig_width / orig_height aspect_ratio = orig_width / orig_height
@@ -304,8 +305,8 @@ class Eagle25VLImageProcessorFast(BaseImageProcessorFast):
def _pad_for_batching( def _pad_for_batching(
self, self,
pixel_values: list[torch.Tensor], pixel_values: list["torch.Tensor"],
) -> list[torch.Tensor]: ) -> list["torch.Tensor"]:
""" """
Pads images on the `num_of_patches` dimension with zeros to form a batch of same number of patches. Pads images on the `num_of_patches` dimension with zeros to form a batch of same number of patches.
@@ -326,14 +327,14 @@ class Eagle25VLImageProcessorFast(BaseImageProcessorFast):
def _preprocess( def _preprocess(
self, self,
images: list[torch.Tensor], images: list["torch.Tensor"],
do_resize: bool, do_resize: bool,
size: SizeDict, size: SizeDict,
max_dynamic_tiles: int, max_dynamic_tiles: int,
min_dynamic_tiles: int, min_dynamic_tiles: int,
use_thumbnail: bool, use_thumbnail: bool,
pad_during_tiling: bool, pad_during_tiling: bool,
interpolation: F.InterpolationMode | None, interpolation: Optional["F.InterpolationMode"],
do_center_crop: bool, do_center_crop: bool,
crop_size: SizeDict, crop_size: SizeDict,
do_rescale: bool, do_rescale: bool,
+2 -1
View File
@@ -20,11 +20,12 @@ import logging
import math import math
from collections import deque from collections import deque
from pathlib import Path from pathlib import Path
from typing import TYPE_CHECKING, Literal, TypedDict, Unpack from typing import TYPE_CHECKING, Literal, TypedDict
import torch import torch
import torch.nn.functional as F # noqa: N812 import torch.nn.functional as F # noqa: N812
from torch import Tensor, nn from torch import Tensor, nn
from typing_extensions import Unpack
from lerobot.utils.import_utils import _transformers_available from lerobot.utils.import_utils import _transformers_available
+5 -5
View File
@@ -20,11 +20,12 @@ import logging
import math import math
from collections import deque from collections import deque
from pathlib import Path from pathlib import Path
from typing import TYPE_CHECKING, Literal, TypedDict, Unpack from typing import TYPE_CHECKING, Literal, TypedDict
import torch import torch
import torch.nn.functional as F # noqa: N812 import torch.nn.functional as F # noqa: N812
from torch import Tensor, nn from torch import Tensor, nn
from typing_extensions import Unpack
from lerobot.utils.import_utils import _transformers_available from lerobot.utils.import_utils import _transformers_available
@@ -98,11 +99,10 @@ def create_sinusoidal_pos_embedding( # see openpi `create_sinusoidal_pos_embedd
def sample_beta(alpha, beta, bsize, device): # see openpi `sample_beta` (exact copy) def sample_beta(alpha, beta, bsize, device): # see openpi `sample_beta` (exact copy)
# Beta sampling uses _sample_dirichlet which isn't implemented for MPS, so sample on CPU alpha_t = torch.as_tensor(alpha, dtype=torch.float32, device=device)
alpha_t = torch.tensor(alpha, dtype=torch.float32) beta_t = torch.as_tensor(beta, dtype=torch.float32, device=device)
beta_t = torch.tensor(beta, dtype=torch.float32)
dist = torch.distributions.Beta(alpha_t, beta_t) dist = torch.distributions.Beta(alpha_t, beta_t)
return dist.sample((bsize,)).to(device) return dist.sample((bsize,))
def make_att_2d_masks(pad_masks, att_masks): # see openpi `make_att_2d_masks` (exact copy) def make_att_2d_masks(pad_masks, att_masks): # see openpi `make_att_2d_masks` (exact copy)
@@ -19,12 +19,13 @@ import logging
import math import math
from collections import deque from collections import deque
from pathlib import Path from pathlib import Path
from typing import TYPE_CHECKING, Literal, TypedDict, Unpack from typing import TYPE_CHECKING, Literal, TypedDict
import numpy as np import numpy as np
import torch import torch
import torch.nn.functional as F # noqa: N812 import torch.nn.functional as F # noqa: N812
from torch import Tensor, nn from torch import Tensor, nn
from typing_extensions import Unpack
from lerobot.utils.import_utils import _scipy_available, _transformers_available from lerobot.utils.import_utils import _scipy_available, _transformers_available
+1 -1
View File
@@ -260,7 +260,7 @@ class PiGemmaModel(GemmaModel): # type: ignore[misc]
causal_mask = create_causal_mask( causal_mask = create_causal_mask(
config=self.config, config=self.config,
inputs_embeds=inputs_embeds, input_embeds=inputs_embeds,
attention_mask=attention_mask, attention_mask=attention_mask,
cache_position=cache_position, cache_position=cache_position,
past_key_values=past_key_values, past_key_values=past_key_values,
+2 -1
View File
@@ -19,7 +19,7 @@ import os
from importlib.resources import files from importlib.resources import files
from pathlib import Path from pathlib import Path
from tempfile import TemporaryDirectory from tempfile import TemporaryDirectory
from typing import TypedDict, TypeVar, Unpack from typing import TypedDict, TypeVar
import packaging import packaging
import safetensors import safetensors
@@ -28,6 +28,7 @@ from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
from huggingface_hub.errors import HfHubHTTPError from huggingface_hub.errors import HfHubHTTPError
from safetensors.torch import load_model as load_model_as_safetensor, save_model as save_model_as_safetensor from safetensors.torch import load_model as load_model_as_safetensor, save_model as save_model_as_safetensor
from torch import Tensor, nn from torch import Tensor, nn
from typing_extensions import Unpack
from lerobot.configs.policies import PreTrainedConfig from lerobot.configs.policies import PreTrainedConfig
from lerobot.configs.train import TrainPipelineConfig from lerobot.configs.train import TrainPipelineConfig
+3 -1
View File
@@ -277,7 +277,9 @@ class SARMEncodingProcessorStep(ProcessorStep):
# When language is perturbed, targets are zero so perturbed samples don't contribute to progress loss # When language is perturbed, targets are zero so perturbed samples don't contribute to progress loss
if self.dataset_meta is not None: if self.dataset_meta is not None:
episodes_df = self.dataset_meta.episodes.to_pandas() episodes_df = None
if self.sparse_subtask_names != ["task"]:
episodes_df = self.dataset_meta.episodes.to_pandas()
# Generate sparse targets # Generate sparse targets
if self.sparse_temporal_proportions is not None: if self.sparse_temporal_proportions is not None:
@@ -106,9 +106,6 @@ class SmolVLAConfig(PreTrainedConfig):
# Real-Time Chunking (RTC) configuration # Real-Time Chunking (RTC) configuration
rtc_config: RTCConfig | None = None rtc_config: RTCConfig | None = None
compile_model: bool = False # Whether to use torch.compile for model optimization
compile_mode: str = "max-autotune" # Torch compile mode
def __post_init__(self): def __post_init__(self):
super().__post_init__() super().__post_init__()
@@ -54,11 +54,12 @@ policy = SmolVLAPolicy.from_pretrained("lerobot/smolvla_base")
import math import math
from collections import deque from collections import deque
from typing import TypedDict, Unpack from typing import TypedDict
import torch import torch
import torch.nn.functional as F # noqa: N812 import torch.nn.functional as F # noqa: N812
from torch import Tensor, nn from torch import Tensor, nn
from typing_extensions import Unpack
from lerobot.policies.pretrained import PreTrainedPolicy from lerobot.policies.pretrained import PreTrainedPolicy
from lerobot.policies.rtc.modeling_rtc import RTCProcessor from lerobot.policies.rtc.modeling_rtc import RTCProcessor
@@ -592,12 +593,6 @@ class VLAFlowMatching(nn.Module):
self.prefix_length = self.config.prefix_length self.prefix_length = self.config.prefix_length
self.rtc_processor = rtc_processor self.rtc_processor = rtc_processor
# Compile model if requested
if config.compile_model:
torch.set_float32_matmul_precision("high")
self.sample_actions = torch.compile(self.sample_actions, mode=config.compile_mode)
self.forward = torch.compile(self.forward, mode=config.compile_mode)
def _rtc_enabled(self): def _rtc_enabled(self):
return self.config.rtc_config is not None and self.config.rtc_config.enabled return self.config.rtc_config is not None and self.config.rtc_config.enabled
@@ -77,6 +77,7 @@ class SmolVLMWithExpertModel(nn.Module):
print(f"Loading {model_id} weights ...") print(f"Loading {model_id} weights ...")
self.vlm = AutoModelForImageTextToText.from_pretrained( self.vlm = AutoModelForImageTextToText.from_pretrained(
model_id, model_id,
device_map=device,
torch_dtype="bfloat16", torch_dtype="bfloat16",
low_cpu_mem_usage=True, low_cpu_mem_usage=True,
) )
@@ -55,7 +55,7 @@ class WallXConfig(PreTrainedConfig):
pretrained_name_or_path: str = "x-square-robot/wall-oss-flow" pretrained_name_or_path: str = "x-square-robot/wall-oss-flow"
# Tokenizer settings # Tokenizer settings
action_tokenizer_path: str | None = "lerobot/fast-action-tokenizer" action_tokenizer_path: str | None = "physical-intelligence/fast"
# Action prediction mode: "diffusion" or "fast" # Action prediction mode: "diffusion" or "fast"
prediction_mode: str = "diffusion" prediction_mode: str = "diffusion"
@@ -1596,7 +1596,7 @@ QWEN2_5_VL_INPUTS_DOCSTRING = r"""
class Qwen2_5_VLForConditionalGeneration(Qwen2_5_VLPreTrainedModel, GenerationMixin): class Qwen2_5_VLForConditionalGeneration(Qwen2_5_VLPreTrainedModel, GenerationMixin):
_tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tied_weights_keys = ["lm_head.weight"]
config_class = Qwen2_5_VLConfig config_class = Qwen2_5_VLConfig
_no_split_modules = ["Qwen2_5_VLDecoderLayer", "Qwen2_5_VLVisionBlock"] _no_split_modules = ["Qwen2_5_VLDecoderLayer", "Qwen2_5_VLVisionBlock"]
@@ -13,12 +13,9 @@
import warnings import warnings
from transformers.configuration_utils import PretrainedConfig from transformers.configuration_utils import PretrainedConfig
from transformers.utils import logging
""" Florence-2 configuration""" """ Florence-2 configuration"""
logger = logging.get_logger(__name__)
class Florence2VisionConfig(PretrainedConfig): class Florence2VisionConfig(PretrainedConfig):
r""" r"""
@@ -46,7 +46,6 @@ from transformers.utils import (
add_start_docstrings_to_model_forward, add_start_docstrings_to_model_forward,
is_flash_attn_2_available, is_flash_attn_2_available,
is_flash_attn_greater_or_equal_2_10, is_flash_attn_greater_or_equal_2_10,
logging,
replace_return_docstrings, replace_return_docstrings,
) )
@@ -57,8 +56,6 @@ if is_flash_attn_2_available():
from flash_attn import flash_attn_func, flash_attn_varlen_func from flash_attn import flash_attn_func, flash_attn_varlen_func
from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa
logger = logging.get_logger(__name__)
_CONFIG_FOR_DOC = "Florence2Config" _CONFIG_FOR_DOC = "Florence2Config"
@@ -992,12 +989,6 @@ class Florence2FlashAttention2(Florence2Attention):
else: else:
target_dtype = self.q_proj.weight.dtype target_dtype = self.q_proj.weight.dtype
logger.warning_once(
f"The input hidden states seems to be silently casted in float32, this might be related to"
f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
f" {target_dtype}."
)
query_states = query_states.to(target_dtype) query_states = query_states.to(target_dtype)
key_states = key_states.to(target_dtype) key_states = key_states.to(target_dtype)
value_states = value_states.to(target_dtype) value_states = value_states.to(target_dtype)
@@ -1135,11 +1126,6 @@ class Florence2SdpaAttention(Florence2Attention):
) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]:
"""Input shape: Batch x Time x Channel""" """Input shape: Batch x Time x Channel"""
if output_attentions or layer_head_mask is not None: if output_attentions or layer_head_mask is not None:
# TODO: Improve this warning with e.g. `model.config._attn_implementation = "manual"` once this is implemented.
logger.warning_once(
"Florence2Model is using Florence2SdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True` or `layer_head_mask` not None. Falling back to the manual attention"
' implementation, but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
)
return super().forward( return super().forward(
hidden_states, hidden_states,
key_value_states=key_value_states, key_value_states=key_value_states,
@@ -1860,9 +1846,6 @@ class Florence2Decoder(Florence2LanguagePreTrainedModel):
hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
if self.gradient_checkpointing and self.training and use_cache: if self.gradient_checkpointing and self.training and use_cache:
logger.warning_once(
"`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
)
use_cache = False use_cache = False
# decoder layers # decoder layers
@@ -2160,8 +2143,6 @@ class Florence2LanguageForConditionalGeneration(Florence2LanguagePreTrainedModel
return_dict = return_dict if return_dict is not None else self.config.use_return_dict return_dict = return_dict if return_dict is not None else self.config.use_return_dict
if labels is not None: if labels is not None:
if use_cache:
logger.warning("The `use_cache` argument is changed to `False` since `labels` is provided.")
use_cache = False use_cache = False
if decoder_input_ids is None and decoder_inputs_embeds is None: if decoder_input_ids is None and decoder_inputs_embeds is None:
decoder_input_ids = shift_tokens_right( decoder_input_ids = shift_tokens_right(
+5 -5
View File
@@ -17,7 +17,7 @@
from __future__ import annotations from __future__ import annotations
from enum import Enum from enum import Enum
from typing import Any, TypedDict from typing import Any, TypeAlias, TypedDict
import numpy as np import numpy as np
import torch import torch
@@ -36,10 +36,10 @@ class TransitionKey(str, Enum):
COMPLEMENTARY_DATA = "complementary_data" COMPLEMENTARY_DATA = "complementary_data"
PolicyAction = torch.Tensor PolicyAction: TypeAlias = torch.Tensor
RobotAction = dict[str, Any] RobotAction: TypeAlias = dict[str, Any]
EnvAction = np.ndarray EnvAction: TypeAlias = np.ndarray
RobotObservation = dict[str, Any] RobotObservation: TypeAlias = dict[str, Any]
EnvTransition = TypedDict( EnvTransition = TypedDict(
-38
View File
@@ -153,44 +153,6 @@ class LiberoProcessorStep(ObservationProcessorStep):
return result return result
@dataclass
@ProcessorStepRegistry.register(name="robocasa_processor")
class RoboCasaProcessorStep(ObservationProcessorStep):
"""
Processes RoboCasa observations into LeRobot format.
The RoboCasaEnv wrapper returns:
- ``pixels.<cam_name>``: (B, C, H, W) float32 images (already converted by vectorenv)
- ``observation.robot_state``: (B, 16) float32 proprioception
This step remaps them to:
- ``observation.images.<cam_name>`` (unchanged tensor)
- ``observation.state`` (robot_state renamed)
"""
def _process_observation(self, observation: dict) -> dict:
processed = {}
obs_prefix = OBS_PREFIX # "observation."
for key, value in observation.items():
if key.startswith(f"{OBS_IMAGES}."):
# Already in the right place; pass through
processed[key] = value
elif key == OBS_STATE or key == f"{obs_prefix}robot_state":
# Rename robot_state → observation.state
processed[OBS_STATE] = value.float() if hasattr(value, "float") else value
return processed
def transform_features(
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
return features
def observation(self, observation: dict) -> dict:
return self._process_observation(observation)
@dataclass @dataclass
@ProcessorStepRegistry.register(name="isaaclab_arena_processor") @ProcessorStepRegistry.register(name="isaaclab_arena_processor")
class IsaaclabArenaProcessorStep(ObservationProcessorStep): class IsaaclabArenaProcessorStep(ObservationProcessorStep):
+4 -4
View File
@@ -39,7 +39,7 @@ from collections.abc import Callable, Iterable, Sequence
from copy import deepcopy from copy import deepcopy
from dataclasses import dataclass, field from dataclasses import dataclass, field
from pathlib import Path from pathlib import Path
from typing import Any, TypedDict, TypeVar, cast from typing import Any, Generic, TypeAlias, TypedDict, TypeVar, cast
import torch import torch
from huggingface_hub import hf_hub_download from huggingface_hub import hf_hub_download
@@ -251,7 +251,7 @@ class ProcessorMigrationError(Exception):
@dataclass @dataclass
class DataProcessorPipeline[TInput, TOutput](HubMixin): class DataProcessorPipeline(HubMixin, Generic[TInput, TOutput]):
"""A sequential pipeline for processing data, integrated with the Hugging Face Hub. """A sequential pipeline for processing data, integrated with the Hugging Face Hub.
This class chains together multiple `ProcessorStep` instances to form a complete This class chains together multiple `ProcessorStep` instances to form a complete
@@ -1432,8 +1432,8 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
# Type aliases for semantic clarity. # Type aliases for semantic clarity.
RobotProcessorPipeline = DataProcessorPipeline[TInput, TOutput] RobotProcessorPipeline: TypeAlias = DataProcessorPipeline[TInput, TOutput]
PolicyProcessorPipeline = DataProcessorPipeline[TInput, TOutput] PolicyProcessorPipeline: TypeAlias = DataProcessorPipeline[TInput, TOutput]
class ObservationProcessorStep(ProcessorStep, ABC): class ObservationProcessorStep(ProcessorStep, ABC):
@@ -15,6 +15,7 @@
# limitations under the License. # limitations under the License.
from dataclasses import dataclass, field from dataclasses import dataclass, field
from typing import TypeAlias
from lerobot.cameras import CameraConfig from lerobot.cameras import CameraConfig
@@ -49,5 +50,5 @@ class SOFollowerRobotConfig(RobotConfig, SOFollowerConfig):
pass pass
SO100FollowerConfig = SOFollowerRobotConfig SO100FollowerConfig: TypeAlias = SOFollowerRobotConfig
SO101FollowerConfig = SOFollowerRobotConfig SO101FollowerConfig: TypeAlias = SOFollowerRobotConfig
@@ -17,6 +17,7 @@
import logging import logging
import time import time
from functools import cached_property from functools import cached_property
from typing import TypeAlias
from lerobot.cameras.utils import make_cameras_from_configs from lerobot.cameras.utils import make_cameras_from_configs
from lerobot.motors import Motor, MotorCalibration, MotorNormMode from lerobot.motors import Motor, MotorCalibration, MotorNormMode
@@ -229,5 +230,5 @@ class SOFollower(Robot):
logger.info(f"{self} disconnected.") logger.info(f"{self} disconnected.")
SO100Follower = SOFollower SO100Follower: TypeAlias = SOFollower
SO101Follower = SOFollower SO101Follower: TypeAlias = SOFollower
@@ -16,5 +16,3 @@
from .config_unitree_g1 import UnitreeG1Config from .config_unitree_g1 import UnitreeG1Config
from .unitree_g1 import UnitreeG1 from .unitree_g1 import UnitreeG1
__all__ = ["UnitreeG1", "UnitreeG1Config"]
@@ -27,10 +27,11 @@ _GAINS: dict[str, dict[str, list[float]]] = {
}, # pitch, roll, yaw, knee, ankle_pitch, ankle_roll }, # pitch, roll, yaw, knee, ankle_pitch, ankle_roll
"right_leg": {"kp": [150, 150, 150, 300, 40, 40], "kd": [2, 2, 2, 4, 2, 2]}, "right_leg": {"kp": [150, 150, 150, 300, 40, 40], "kd": [2, 2, 2, 4, 2, 2]},
"waist": {"kp": [250, 250, 250], "kd": [5, 5, 5]}, # yaw, roll, pitch "waist": {"kp": [250, 250, 250], "kd": [5, 5, 5]}, # yaw, roll, pitch
"left_arm": {"kp": [50, 50, 80, 80], "kd": [3, 3, 3, 3]}, # shoulder_pitch/roll/yaw, elbow "left_arm": {"kp": [80, 80, 80, 80], "kd": [3, 3, 3, 3]}, # shoulder_pitch/roll/yaw, elbow
"left_wrist": {"kp": [40, 40, 40], "kd": [1.5, 1.5, 1.5]}, # roll, pitch, yaw "left_wrist": {"kp": [40, 40, 40], "kd": [1.5, 1.5, 1.5]}, # roll, pitch, yaw
"right_arm": {"kp": [50, 50, 80, 80], "kd": [3, 3, 3, 3]}, "right_arm": {"kp": [80, 80, 80, 80], "kd": [3, 3, 3, 3]},
"right_wrist": {"kp": [40, 40, 40], "kd": [1.5, 1.5, 1.5]}, "right_wrist": {"kp": [40, 40, 40], "kd": [1.5, 1.5, 1.5]},
"other": {"kp": [80, 80, 80, 80, 80, 80], "kd": [3, 3, 3, 3, 3, 3]},
} }
@@ -67,7 +68,3 @@ class UnitreeG1Config(RobotConfig):
# Compensates for gravity on the unitree's arms using the arm ik solver # Compensates for gravity on the unitree's arms using the arm ik solver
gravity_compensation: bool = False gravity_compensation: bool = False
# Lower-body controller class name, e.g. "GrootLocomotionController" or
# "HolosomaLocomotionController". None disables it.
controller: str | None = None
+2 -39
View File
@@ -14,34 +14,12 @@
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
import importlib
from enum import IntEnum from enum import IntEnum
import numpy as np
# ruff: noqa: N801, N815 # ruff: noqa: N801, N815
NUM_MOTORS = 29 NUM_MOTORS = 29
REMOTE_AXES = ("remote.lx", "remote.ly", "remote.rx", "remote.ry")
REMOTE_BUTTONS = tuple(f"remote.button.{i}" for i in range(16))
REMOTE_KEYS = REMOTE_AXES + REMOTE_BUTTONS
def default_remote_input() -> dict[str, float]:
"""Return a zeroed-out remote input dict (axes + buttons)."""
return dict.fromkeys(REMOTE_KEYS, 0.0)
def get_gravity_orientation(quaternion: list[float] | np.ndarray) -> np.ndarray:
"""Get gravity orientation from quaternion [w, x, y, z]."""
qw, qx, qy, qz = quaternion
gravity_orientation = np.zeros(3, dtype=np.float32)
gravity_orientation[0] = 2 * (-qz * qx + qw * qy)
gravity_orientation[1] = -2 * (qz * qy + qw * qx)
gravity_orientation[2] = 1 - 2 * (qw * qw + qz * qz)
return gravity_orientation
class G1_29_JointArmIndex(IntEnum): class G1_29_JointArmIndex(IntEnum):
# Left arm # Left arm
@@ -51,7 +29,7 @@ class G1_29_JointArmIndex(IntEnum):
kLeftElbow = 18 kLeftElbow = 18
kLeftWristRoll = 19 kLeftWristRoll = 19
kLeftWristPitch = 20 kLeftWristPitch = 20
kLeftWristYaw = 21 kLeftWristyaw = 21
# Right arm # Right arm
kRightShoulderPitch = 22 kRightShoulderPitch = 22
@@ -63,21 +41,6 @@ class G1_29_JointArmIndex(IntEnum):
kRightWristYaw = 28 kRightWristYaw = 28
def make_locomotion_controller(name: str | None):
"""Instantiate a locomotion controller by class name. Returns None if name is None."""
if name is None:
return None
controllers = {
"GrootLocomotionController": "lerobot.robots.unitree_g1.gr00t_locomotion",
"HolosomaLocomotionController": "lerobot.robots.unitree_g1.holosoma_locomotion",
}
module_path = controllers.get(name)
if module_path is None:
raise ValueError(f"Unknown controller: {name!r}. Available: {list(controllers)}")
module = importlib.import_module(module_path)
return getattr(module, name)()
class G1_29_JointIndex(IntEnum): class G1_29_JointIndex(IntEnum):
# Left leg # Left leg
kLeftHipPitch = 0 kLeftHipPitch = 0
@@ -106,7 +69,7 @@ class G1_29_JointIndex(IntEnum):
kLeftElbow = 18 kLeftElbow = 18
kLeftWristRoll = 19 kLeftWristRoll = 19
kLeftWristPitch = 20 kLeftWristPitch = 20
kLeftWristYaw = 21 kLeftWristyaw = 21
# Right arm # Right arm
kRightShoulderPitch = 22 kRightShoulderPitch = 22
@@ -16,11 +16,13 @@
import logging import logging
import os import os
from collections import deque import sys
import numpy as np import numpy as np
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
parent2_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.append(parent2_dir)
class WeightedMovingFilter: class WeightedMovingFilter:
@@ -29,14 +31,18 @@ class WeightedMovingFilter:
self._weights = np.array(weights) self._weights = np.array(weights)
self._data_size = data_size self._data_size = data_size
self._filtered_data = np.zeros(self._data_size) self._filtered_data = np.zeros(self._data_size)
self._data_queue = deque(maxlen=self._window_size) self._data_queue = []
def _apply_filter(self): def _apply_filter(self):
if len(self._data_queue) < self._window_size: if len(self._data_queue) < self._window_size:
return self._data_queue[-1] return self._data_queue[-1]
data_array = np.array(self._data_queue) data_array = np.array(self._data_queue)
return data_array.T @ self._weights temp_filtered_data = np.zeros(self._data_size)
for i in range(self._data_size):
temp_filtered_data[i] = np.convolve(data_array[:, i], self._weights, mode="valid")[-1]
return temp_filtered_data
def add_data(self, new_data): def add_data(self, new_data):
assert len(new_data) == self._data_size assert len(new_data) == self._data_size
@@ -46,6 +52,9 @@ class WeightedMovingFilter:
): # skip duplicate data ): # skip duplicate data
return return
if len(self._data_queue) >= self._window_size:
self._data_queue.pop(0)
self._data_queue.append(new_data) self._data_queue.append(new_data)
self._filtered_data = self._apply_filter() self._filtered_data = self._apply_filter()
@@ -62,6 +71,8 @@ class G1_29_ArmIK: # noqa: N801
from pinocchio import casadi as cpin from pinocchio import casadi as cpin
self._pin = pin self._pin = pin
np.set_printoptions(precision=5, suppress=True, linewidth=200)
self.unit_test = unit_test self.unit_test = unit_test
self.repo_path = snapshot_download("lerobot/unitree-g1-mujoco") self.repo_path = snapshot_download("lerobot/unitree-g1-mujoco")
@@ -238,35 +249,50 @@ class G1_29_ArmIK: # noqa: N801
self.opti.set_value(self.param_tf_r, right_wrist) self.opti.set_value(self.param_tf_r, right_wrist)
self.opti.set_value(self.var_q_last, self.init_data) # for smooth self.opti.set_value(self.var_q_last, self.init_data) # for smooth
converged = True
try: try:
self.opti.solve() self.opti.solve()
sol_q = self.opti.value(self.var_q) sol_q = self.opti.value(self.var_q)
self.smooth_filter.add_data(sol_q)
sol_q = self.smooth_filter.filtered_data
if current_lr_arm_motor_dq is not None:
v = current_lr_arm_motor_dq * 0.0
else:
v = (sol_q - self.init_data) * 0.0
self.init_data = sol_q
sol_tauff = self._pin.rnea(
self.reduced_robot.model,
self.reduced_robot.data,
sol_q,
v,
np.zeros(self.reduced_robot.model.nv),
)
return sol_q, sol_tauff
except Exception as e: except Exception as e:
converged = False logger.error(f"ERROR in convergence, plotting debug info.{e}")
logger.error(f"IK convergence error: {e}")
sol_q = self.opti.debug.value(self.var_q) sol_q = self.opti.debug.value(self.var_q)
self.smooth_filter.add_data(sol_q)
sol_q = self.smooth_filter.filtered_data
self.smooth_filter.add_data(sol_q) if current_lr_arm_motor_dq is not None:
sol_q = self.smooth_filter.filtered_data v = current_lr_arm_motor_dq * 0.0
self.init_data = sol_q else:
v = (sol_q - self.init_data) * 0.0
self.init_data = sol_q
if not converged:
logger.error( logger.error(
f"sol_q:{sol_q} \nmotorstate: \n{current_lr_arm_motor_q} \nleft_pose: \n{left_wrist} \nright_pose: \n{right_wrist}" f"sol_q:{sol_q} \nmotorstate: \n{current_lr_arm_motor_q} \nleft_pose: \n{left_wrist} \nright_pose: \n{right_wrist}"
) )
return current_lr_arm_motor_q, np.zeros(self.reduced_robot.model.nv) return current_lr_arm_motor_q, np.zeros(self.reduced_robot.model.nv)
sol_tauff = self._pin.rnea(
self.reduced_robot.model,
self.reduced_robot.data,
sol_q,
np.zeros(self.reduced_robot.model.nv),
np.zeros(self.reduced_robot.model.nv),
)
return sol_q, sol_tauff
def solve_tau(self, current_lr_arm_motor_q=None, current_lr_arm_motor_dq=None): def solve_tau(self, current_lr_arm_motor_q=None, current_lr_arm_motor_dq=None):
try: try:
q_g1 = np.array(current_lr_arm_motor_q, dtype=float) q_g1 = np.array(current_lr_arm_motor_q, dtype=float)
@@ -24,7 +24,6 @@ This server runs on the robot and forwards:
Uses JSON for secure serialization instead of pickle. Uses JSON for secure serialization instead of pickle.
""" """
import argparse
import base64 import base64
import contextlib import contextlib
import json import json
@@ -39,8 +38,6 @@ from unitree_sdk2py.idl.default import unitree_hg_msg_dds__LowCmd_
from unitree_sdk2py.idl.unitree_hg.msg.dds_ import LowCmd_ as hg_LowCmd, LowState_ as hg_LowState from unitree_sdk2py.idl.unitree_hg.msg.dds_ import LowCmd_ as hg_LowCmd, LowState_ as hg_LowState
from unitree_sdk2py.utils.crc import CRC from unitree_sdk2py.utils.crc import CRC
from lerobot.cameras.zmq.image_server import ImageServer
# DDS topic names follow Unitree SDK naming conventions # DDS topic names follow Unitree SDK naming conventions
# ruff: noqa: N816 # ruff: noqa: N816
kTopicLowCommand_Debug = "rt/lowcmd" # action to robot kTopicLowCommand_Debug = "rt/lowcmd" # action to robot
@@ -153,32 +150,6 @@ def cmd_forward_loop(
def main() -> None: def main() -> None:
"""Main entry point for the robot server bridge.""" """Main entry point for the robot server bridge."""
parser = argparse.ArgumentParser(description="DDS-to-ZMQ bridge server for Unitree G1")
parser.add_argument("--camera", action="store_true", help="Also launch camera server")
parser.add_argument("--camera-device", type=int, default=4, help="Camera device ID (default: 4)")
parser.add_argument("--camera-fps", type=int, default=30, help="Camera FPS (default: 30)")
parser.add_argument("--camera-width", type=int, default=640, help="Camera width (default: 640)")
parser.add_argument("--camera-height", type=int, default=480, help="Camera height (default: 480)")
parser.add_argument("--camera-port", type=int, default=5555, help="Camera ZMQ port (default: 5555)")
args = parser.parse_args()
# Optionally start camera server in background thread
camera_thread = None
if args.camera:
camera_config = {
"fps": args.camera_fps,
"cameras": {
"head_camera": {
"device_id": args.camera_device,
"shape": [args.camera_height, args.camera_width],
}
},
}
camera_server = ImageServer(camera_config, port=args.camera_port)
camera_thread = threading.Thread(target=camera_server.run, daemon=True)
camera_thread.start()
print(f"Camera server started on port {args.camera_port} (device {args.camera_device})")
# initialize DDS # initialize DDS
ChannelFactoryInitialize(0) ChannelFactoryInitialize(0)
@@ -235,8 +206,6 @@ def main() -> None:
shutdown_event.set() shutdown_event.set()
ctx.term() # terminates blocking zmq.recv() calls ctx.term() # terminates blocking zmq.recv() calls
t_state.join(timeout=2.0) t_state.join(timeout=2.0)
if camera_thread is not None:
camera_thread.join(timeout=2.0)
if __name__ == "__main__": if __name__ == "__main__":
+132 -252
View File
@@ -14,67 +14,27 @@
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
from __future__ import annotations
import logging import logging
import struct
import threading import threading
import time import time
from dataclasses import dataclass, field from dataclasses import dataclass, field
from functools import cached_property from functools import cached_property
from typing import TYPE_CHECKING, Protocol, runtime_checkable from typing import Any
import numpy as np import numpy as np
from lerobot.cameras.utils import make_cameras_from_configs from lerobot.cameras.utils import make_cameras_from_configs
from lerobot.envs.factory import make_env from lerobot.envs.factory import make_env
from lerobot.processor import RobotAction, RobotObservation from lerobot.processor import RobotAction, RobotObservation
from lerobot.robots.unitree_g1.g1_kinematics import G1_29_ArmIK from lerobot.robots.unitree_g1.g1_utils import G1_29_JointArmIndex, G1_29_JointIndex
from lerobot.robots.unitree_g1.g1_utils import ( from lerobot.robots.unitree_g1.robot_kinematic_processor import G1_29_ArmIK
REMOTE_AXES,
REMOTE_KEYS,
G1_29_JointArmIndex,
G1_29_JointIndex,
default_remote_input,
make_locomotion_controller,
)
from lerobot.utils.import_utils import _unitree_sdk_available
from ..robot import Robot from ..robot import Robot
from .config_unitree_g1 import UnitreeG1Config from .config_unitree_g1 import UnitreeG1Config
if TYPE_CHECKING or _unitree_sdk_available:
from unitree_sdk2py.core.channel import (
ChannelFactoryInitialize as _SDKChannelFactoryInitialize,
ChannelPublisher as _SDKChannelPublisher,
ChannelSubscriber as _SDKChannelSubscriber,
)
from unitree_sdk2py.idl.default import unitree_hg_msg_dds__LowCmd_
from unitree_sdk2py.idl.unitree_hg.msg.dds_ import (
LowCmd_ as hg_LowCmd,
LowState_ as hg_LowState,
)
from unitree_sdk2py.utils.crc import CRC
else:
_SDKChannelFactoryInitialize = None
_SDKChannelPublisher = None
_SDKChannelSubscriber = None
unitree_hg_msg_dds__LowCmd_ = None
hg_LowCmd = None
hg_LowState = None
CRC = None
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@runtime_checkable
class LocomotionController(Protocol):
control_dt: float
def run_step(self, action: dict, lowstate) -> dict: ...
def reset(self) -> None: ...
# DDS topic names follow Unitree SDK naming conventions # DDS topic names follow Unitree SDK naming conventions
# ruff: noqa: N816 # ruff: noqa: N816
kTopicLowCommand_Debug = "rt/lowcmd" kTopicLowCommand_Debug = "rt/lowcmd"
@@ -103,7 +63,7 @@ class IMUState:
class G1_29_LowState: # noqa: N801 class G1_29_LowState: # noqa: N801
motor_state: list[MotorState] = field(default_factory=lambda: [MotorState() for _ in G1_29_JointIndex]) motor_state: list[MotorState] = field(default_factory=lambda: [MotorState() for _ in G1_29_JointIndex])
imu_state: IMUState = field(default_factory=IMUState) imu_state: IMUState = field(default_factory=IMUState)
wireless_remote: bytes | None = None # Raw wireless remote data wireless_remote: Any = None # Raw wireless remote data
mode_machine: int = 0 # Robot mode mode_machine: int = 0 # Robot mode
@@ -111,6 +71,25 @@ class UnitreeG1(Robot):
config_class = UnitreeG1Config config_class = UnitreeG1Config
name = "unitree_g1" name = "unitree_g1"
# unitree remote controller
class RemoteController:
def __init__(self):
self.lx = 0
self.ly = 0
self.rx = 0
self.ry = 0
self.button = [0] * 16
def set(self, data):
# wireless_remote
keys = struct.unpack("H", data[2:4])[0]
for i in range(16):
self.button[i] = (keys & (1 << i)) >> i
self.lx = struct.unpack("f", data[4:8])[0]
self.rx = struct.unpack("f", data[8:12])[0]
self.ry = struct.unpack("f", data[12:16])[0]
self.ly = struct.unpack("f", data[20:24])[0]
def __init__(self, config: UnitreeG1Config): def __init__(self, config: UnitreeG1Config):
super().__init__(config) super().__init__(config)
@@ -124,9 +103,11 @@ class UnitreeG1(Robot):
# Import channel classes based on mode # Import channel classes based on mode
if config.is_simulation: if config.is_simulation:
self._ChannelFactoryInitialize = _SDKChannelFactoryInitialize from unitree_sdk2py.core.channel import (
self._ChannelPublisher = _SDKChannelPublisher ChannelFactoryInitialize,
self._ChannelSubscriber = _SDKChannelSubscriber ChannelPublisher,
ChannelSubscriber,
)
else: else:
from lerobot.robots.unitree_g1.unitree_sdk2_socket import ( from lerobot.robots.unitree_g1.unitree_sdk2_socket import (
ChannelFactoryInitialize, ChannelFactoryInitialize,
@@ -134,30 +115,22 @@ class UnitreeG1(Robot):
ChannelSubscriber, ChannelSubscriber,
) )
self._ChannelFactoryInitialize = ChannelFactoryInitialize # Store for use in connect()
self._ChannelPublisher = ChannelPublisher self._ChannelFactoryInitialize = ChannelFactoryInitialize
self._ChannelSubscriber = ChannelSubscriber self._ChannelPublisher = ChannelPublisher
self._ChannelSubscriber = ChannelSubscriber
# Initialize state variables # Initialize state variables
self.sim_env = None self.sim_env = None
self._env_wrapper = None self._env_wrapper = None
self._lowstate = None self._lowstate = None
self._lowstate_lock = threading.Lock()
self._shutdown_event = threading.Event() self._shutdown_event = threading.Event()
self.subscribe_thread = None self.subscribe_thread = None
self.remote_controller = self.RemoteController()
self.arm_ik = G1_29_ArmIK() if config.gravity_compensation else None self.arm_ik = G1_29_ArmIK()
# Lower-body controller loaded dynamically def _subscribe_motor_state(self): # polls robot state @ 250Hz
self.controller: LocomotionController | None = make_locomotion_controller(config.controller)
# Controller thread state
self._controller_thread = None
self._controller_action_lock = threading.Lock()
self.controller_input = default_remote_input()
self.controller_output = {}
def _subscribe_lowstate(self): # polls robot state @ 250Hz
while not self._shutdown_event.is_set(): while not self._shutdown_event.is_set():
start_time = time.time() start_time = time.time()
@@ -170,11 +143,11 @@ class UnitreeG1(Robot):
lowstate = G1_29_LowState() lowstate = G1_29_LowState()
# Capture motor states using jointindex # Capture motor states using jointindex
for joint in G1_29_JointIndex: for id in G1_29_JointIndex:
lowstate.motor_state[joint].q = msg.motor_state[joint].q lowstate.motor_state[id].q = msg.motor_state[id].q
lowstate.motor_state[joint].dq = msg.motor_state[joint].dq lowstate.motor_state[id].dq = msg.motor_state[id].dq
lowstate.motor_state[joint].tau_est = msg.motor_state[joint].tau_est lowstate.motor_state[id].tau_est = msg.motor_state[id].tau_est
lowstate.motor_state[joint].temperature = msg.motor_state[joint].temperature lowstate.motor_state[id].temperature = msg.motor_state[id].temperature
# Capture IMU state # Capture IMU state
lowstate.imu_state.quaternion = list(msg.imu_state.quaternion) lowstate.imu_state.quaternion = list(msg.imu_state.quaternion)
@@ -189,106 +162,31 @@ class UnitreeG1(Robot):
# Capture mode_machine # Capture mode_machine
lowstate.mode_machine = msg.mode_machine lowstate.mode_machine = msg.mode_machine
with self._lowstate_lock: self._lowstate = lowstate
self._lowstate = lowstate
current_time = time.time() current_time = time.time()
all_t_elapsed = current_time - start_time all_t_elapsed = current_time - start_time
sleep_time = max(0, (self.control_dt - all_t_elapsed)) # maintain constant control dt sleep_time = max(0, (self.control_dt - all_t_elapsed)) # maintain constant control dt
time.sleep(sleep_time) time.sleep(sleep_time)
def publish_lowcmd(
self,
action: RobotAction,
kp: np.ndarray | list[float] | None = None,
kd: np.ndarray | list[float] | None = None,
tau: np.ndarray | list[float] | None = None,
) -> None: # writes robot command whenever requested
for motor in G1_29_JointIndex:
key = f"{motor.name}.q"
if key in action:
self.msg.motor_cmd[motor.value].q = action[key]
self.msg.motor_cmd[motor.value].qd = 0
self.msg.motor_cmd[motor.value].kp = (
kp[motor.value] if kp is not None else self.kp[motor.value]
)
self.msg.motor_cmd[motor.value].kd = (
kd[motor.value] if kd is not None else self.kd[motor.value]
)
self.msg.motor_cmd[motor.value].tau = tau[motor.value] if tau is not None else 0.0
self.msg.crc = self.crc.Crc(self.msg)
self.lowcmd_publisher.Write(self.msg)
@property
def _cameras_ft(self) -> dict[str, tuple]:
return {
cam: (self.config.cameras[cam].height, self.config.cameras[cam].width, 3) for cam in self.cameras
}
@cached_property
def observation_features(self) -> dict[str, type | tuple]:
return {**self._motors_ft, **self._cameras_ft}
@cached_property @cached_property
def action_features(self) -> dict[str, type]: def action_features(self) -> dict[str, type]:
if self.controller is None: return {f"{G1_29_JointIndex(motor).name}.q": float for motor in G1_29_JointIndex}
return {f"{G1_29_JointIndex(motor).name}.q": float for motor in G1_29_JointIndex}
arm_features = {f"{G1_29_JointArmIndex(motor).name}.q": float for motor in G1_29_JointArmIndex} def calibrate(self) -> None: # robot is already calibrated
remote_features = dict.fromkeys(REMOTE_AXES, float)
return {**arm_features, **remote_features}
def _controller_loop(self):
"""Background thread that runs controller at policy's control_dt."""
control_dt = self.controller.control_dt
logger.info(f"Controller loop starting with control_dt={control_dt} ({1.0 / control_dt:.1f}Hz)")
loop_count = 0
last_log_time = time.time()
while not self._shutdown_event.is_set():
start_time = time.time()
with self._lowstate_lock:
lowstate = self._lowstate
if lowstate is not None and self.controller is not None:
loop_count += 1
if time.time() - last_log_time >= 5.0: # Log every 5 seconds
actual_hz = loop_count / (time.time() - last_log_time)
logger.info(
f"Controller actual rate: {actual_hz:.1f}Hz (target: {1.0 / control_dt:.1f}Hz)"
)
loop_count = 0
last_log_time = time.time()
# Read controller input snapshot
with self._controller_action_lock:
controller_input = dict(self.controller_input)
# Run controller step
controller_action = self.controller.run_step(controller_input, lowstate)
# Write controller output snapshot
with self._controller_action_lock:
self.controller_output = dict(controller_action)
ctrl_kp = self.controller.kp if hasattr(self.controller, "kp") else None
ctrl_kd = self.controller.kd if hasattr(self.controller, "kd") else None
self.publish_lowcmd(controller_action, kp=ctrl_kp, kd=ctrl_kd)
elapsed = time.time() - start_time
sleep_time = max(0, control_dt - elapsed)
time.sleep(sleep_time)
def calibrate(self) -> None:
# TODO: implement g1_29 calibration
pass pass
def configure(self) -> None: def configure(self) -> None:
pass pass
def connect(self, calibrate: bool = True) -> None: # connect to DDS def connect(self, calibrate: bool = True) -> None: # connect to DDS
from unitree_sdk2py.idl.default import unitree_hg_msg_dds__LowCmd_
from unitree_sdk2py.idl.unitree_hg.msg.dds_ import (
LowCmd_ as hg_LowCmd,
LowState_ as hg_LowState,
)
from unitree_sdk2py.utils.crc import CRC
# Initialize DDS channel and simulation environment # Initialize DDS channel and simulation environment
if self.config.is_simulation: if self.config.is_simulation:
self._ChannelFactoryInitialize(0, "lo") self._ChannelFactoryInitialize(0, "lo")
@@ -296,7 +194,7 @@ class UnitreeG1(Robot):
# Extract the actual gym env from the dict structure # Extract the actual gym env from the dict structure
self.sim_env = self._env_wrapper["hub_env"][0].envs[0] self.sim_env = self._env_wrapper["hub_env"][0].envs[0]
else: else:
self._ChannelFactoryInitialize(0, config=self.config) self._ChannelFactoryInitialize(0)
# Initialize direct motor control interface # Initialize direct motor control interface
self.lowcmd_publisher = self._ChannelPublisher(kTopicLowCommand_Debug, hg_LowCmd) self.lowcmd_publisher = self._ChannelPublisher(kTopicLowCommand_Debug, hg_LowCmd)
@@ -305,7 +203,7 @@ class UnitreeG1(Robot):
self.lowstate_subscriber.Init() self.lowstate_subscriber.Init()
# Start subscribe thread to read robot state # Start subscribe thread to read robot state
self.subscribe_thread = threading.Thread(target=self._subscribe_lowstate) self.subscribe_thread = threading.Thread(target=self._subscribe_motor_state)
self.subscribe_thread.start() self.subscribe_thread.start()
# Connect cameras # Connect cameras
@@ -322,53 +220,25 @@ class UnitreeG1(Robot):
# Wait for first state message to arrive # Wait for first state message to arrive
lowstate = None lowstate = None
deadline = time.time() + 10.0
while lowstate is None: while lowstate is None:
with self._lowstate_lock: lowstate = self._lowstate
lowstate = self._lowstate
if lowstate is None: if lowstate is None:
if time.time() > deadline:
raise TimeoutError("Timed out waiting for robot state (10s)")
logger.warning("[UnitreeG1] Waiting for robot state...")
time.sleep(0.01) time.sleep(0.01)
logger.info("[UnitreeG1] Connected to robot.") logger.warning("[UnitreeG1] Waiting for robot state...")
logger.warning("[UnitreeG1] Connected to robot.")
self.msg.mode_machine = lowstate.mode_machine self.msg.mode_machine = lowstate.mode_machine
# Initialize all motors with unified kp/kd from config
self.kp = np.array(self.config.kp, dtype=np.float32) self.kp = np.array(self.config.kp, dtype=np.float32)
self.kd = np.array(self.config.kd, dtype=np.float32) self.kd = np.array(self.config.kd, dtype=np.float32)
for joint in G1_29_JointIndex: for id in G1_29_JointIndex:
self.msg.motor_cmd[joint].mode = 1 self.msg.motor_cmd[id].mode = 1
self.msg.motor_cmd[joint].kp = self.kp[joint.value] self.msg.motor_cmd[id].kp = self.kp[id.value]
self.msg.motor_cmd[joint].kd = self.kd[joint.value] self.msg.motor_cmd[id].kd = self.kd[id.value]
self.msg.motor_cmd[joint].q = lowstate.motor_state[joint.value].q self.msg.motor_cmd[id].q = lowstate.motor_state[id.value].q
# Start controller thread if enabled
if self.controller is not None:
self._controller_thread = threading.Thread(target=self._controller_loop, daemon=True)
self._controller_thread.start()
fps = int(1.0 / self.controller.control_dt)
logger.info(f"Controller thread started ({fps}Hz)")
def _send_zero_torque(self) -> None:
"""Send a zero-gain command to make joints passive before shutting down."""
try:
with self._lowstate_lock:
lowstate = self._lowstate
if lowstate is None:
return
action = {f"{motor.name}.q": lowstate.motor_state[motor.value].q for motor in G1_29_JointIndex}
zero_gains = np.zeros(29, dtype=np.float32)
self.publish_lowcmd(action, kp=zero_gains, kd=zero_gains, tau=zero_gains)
logger.info("Sent zero-torque command for safe shutdown")
except Exception as e:
logger.warning(f"Failed to send zero-torque on disconnect: {e}")
def disconnect(self): def disconnect(self):
# Put robot in passive mode before stopping threads
if not self.config.is_simulation:
self._send_zero_torque()
# Signal thread to stop and unblock any waits # Signal thread to stop and unblock any waits
self._shutdown_event.set() self._shutdown_event.set()
@@ -378,12 +248,6 @@ class UnitreeG1(Robot):
if self.subscribe_thread.is_alive(): if self.subscribe_thread.is_alive():
logger.warning("Subscribe thread did not stop cleanly") logger.warning("Subscribe thread did not stop cleanly")
# Wait for controller thread to finish
if self._controller_thread is not None:
self._controller_thread.join(timeout=2.0)
if self._controller_thread.is_alive():
logger.warning("Controller thread did not stop cleanly")
# Close simulation environment # Close simulation environment
if self.config.is_simulation and self.sim_env is not None: if self.config.is_simulation and self.sim_env is not None:
try: try:
@@ -410,8 +274,7 @@ class UnitreeG1(Robot):
cam.disconnect() cam.disconnect()
def get_observation(self) -> RobotObservation: def get_observation(self) -> RobotObservation:
with self._lowstate_lock: lowstate = self._lowstate
lowstate = self._lowstate
if lowstate is None: if lowstate is None:
return {} return {}
@@ -450,9 +313,14 @@ class UnitreeG1(Robot):
obs["imu.rpy.pitch"] = lowstate.imu_state.rpy[1] obs["imu.rpy.pitch"] = lowstate.imu_state.rpy[1]
obs["imu.rpy.yaw"] = lowstate.imu_state.rpy[2] obs["imu.rpy.yaw"] = lowstate.imu_state.rpy[2]
# Wireless remote (raw bytes for teleoperator) # Controller - parse wireless_remote and add to obs
if lowstate.wireless_remote: if lowstate.wireless_remote and len(lowstate.wireless_remote) >= 24:
obs["wireless_remote"] = lowstate.wireless_remote self.remote_controller.set(lowstate.wireless_remote)
obs["remote.buttons"] = self.remote_controller.button.copy()
obs["remote.lx"] = self.remote_controller.lx
obs["remote.ly"] = self.remote_controller.ly
obs["remote.rx"] = self.remote_controller.rx
obs["remote.ry"] = self.remote_controller.ry
# Cameras - read images from ZMQ cameras # Cameras - read images from ZMQ cameras
for cam_name, cam in self._cameras.items(): for cam_name, cam in self._cameras.items():
@@ -460,63 +328,73 @@ class UnitreeG1(Robot):
return obs return obs
def send_action(self, action: RobotAction) -> RobotAction:
action_to_publish = action
if self.controller is not None:
# Controller thread owns legs/waist. Here we only update joystick inputs
# and publish arm targets from the teleoperator.
self._update_controller_action(action)
arm_prefixes = tuple(j.name for j in G1_29_JointArmIndex)
action_to_publish = {
key: value
for key, value in action.items()
if key.endswith(".q") and key.startswith(arm_prefixes)
}
tau = None
if self.config.gravity_compensation and self.arm_ik is not None:
tau = np.zeros(29, dtype=np.float32)
action_np = np.array(
[
action_to_publish.get(f"{joint.name}.q", self.msg.motor_cmd[joint.value].q)
for joint in G1_29_JointArmIndex
],
dtype=np.float32,
)
arm_tau = self.arm_ik.solve_tau(action_np)
arm_start_idx = G1_29_JointArmIndex.kLeftShoulderPitch.value
for joint in G1_29_JointArmIndex:
local_idx = joint.value - arm_start_idx
tau[joint.value] = arm_tau[local_idx]
self.publish_lowcmd(action_to_publish, tau=tau)
return action
def _update_controller_action(self, action: RobotAction) -> None:
"""Update controller input state from incoming teleop action."""
with self._controller_action_lock:
for key in REMOTE_KEYS:
if key in action:
self.controller_input[key] = action[key]
@property @property
def is_calibrated(self) -> bool: def is_calibrated(self) -> bool:
return True return True
@property @property
def is_connected(self) -> bool: def is_connected(self) -> bool:
with self._lowstate_lock: return self._lowstate is not None
return self._lowstate is not None
@property @property
def _motors_ft(self) -> dict[str, type]: def _motors_ft(self) -> dict[str, type]:
"""Joint positions for all 29 joints."""
return {f"{G1_29_JointIndex(motor).name}.q": float for motor in G1_29_JointIndex} return {f"{G1_29_JointIndex(motor).name}.q": float for motor in G1_29_JointIndex}
@property @property
def cameras(self) -> dict: def cameras(self) -> dict:
return self._cameras return self._cameras
@property
def _cameras_ft(self) -> dict[str, tuple]:
return {
cam: (self.config.cameras[cam].height, self.config.cameras[cam].width, 3) for cam in self.cameras
}
@cached_property
def observation_features(self) -> dict[str, type | tuple]:
return {**self._motors_ft, **self._cameras_ft}
def send_action(self, action: RobotAction) -> RobotAction:
for motor in G1_29_JointIndex:
key = f"{motor.name}.q"
if key in action:
self.msg.motor_cmd[motor.value].q = action[key]
self.msg.motor_cmd[motor.value].qd = 0
self.msg.motor_cmd[motor.value].kp = self.kp[motor.value]
self.msg.motor_cmd[motor.value].kd = self.kd[motor.value]
self.msg.motor_cmd[motor.value].tau = 0
if self.config.gravity_compensation:
# Build action_np from motor commands (arm joints are indices 15-28, local indices 0-13)
action_np = np.zeros(14)
arm_start_idx = G1_29_JointArmIndex.kLeftShoulderPitch.value # 15
for joint in G1_29_JointArmIndex:
local_idx = joint.value - arm_start_idx
action_np[local_idx] = self.msg.motor_cmd[joint.value].q
tau = self.arm_ik.solve_tau(action_np)
# Apply tau back to motor commands
for joint in G1_29_JointArmIndex:
local_idx = joint.value - arm_start_idx
self.msg.motor_cmd[joint.value].tau = tau[local_idx]
self.msg.crc = self.crc.Crc(self.msg)
self.lowcmd_publisher.Write(self.msg)
return action
def get_gravity_orientation(self, quaternion): # get gravity orientation from quaternion
"""Get gravity orientation from quaternion."""
qw = quaternion[0]
qx = quaternion[1]
qy = quaternion[2]
qz = quaternion[3]
gravity_orientation = np.zeros(3)
gravity_orientation[0] = 2 * (-qz * qx + qw * qy)
gravity_orientation[1] = -2 * (qz * qy + qw * qx)
gravity_orientation[2] = 1 - 2 * (qw * qw + qz * qz)
return gravity_orientation
def reset( def reset(
self, self,
control_dt: float | None = None, control_dt: float | None = None,
@@ -529,9 +407,15 @@ class UnitreeG1(Robot):
if self.config.is_simulation and self.sim_env is not None: if self.config.is_simulation and self.sim_env is not None:
self.sim_env.reset() self.sim_env.reset()
self.publish_lowcmd(
{f"{motor.name}.q": float(default_positions[motor.value]) for motor in G1_29_JointIndex} for motor in G1_29_JointIndex:
) self.msg.motor_cmd[motor.value].q = default_positions[motor.value]
self.msg.motor_cmd[motor.value].qd = 0
self.msg.motor_cmd[motor.value].kp = self.kp[motor.value]
self.msg.motor_cmd[motor.value].kd = self.kd[motor.value]
self.msg.motor_cmd[motor.value].tau = 0
self.msg.crc = self.crc.Crc(self.msg)
self.lowcmd_publisher.Write(self.msg)
else: else:
total_time = 3.0 total_time = 3.0
num_steps = int(total_time / control_dt) num_steps = int(total_time / control_dt)
@@ -562,8 +446,4 @@ class UnitreeG1(Robot):
sleep_time = max(0, control_dt - elapsed) sleep_time = max(0, control_dt - elapsed)
time.sleep(sleep_time) time.sleep(sleep_time)
# Reset controller internal state (gait phase, obs history, etc.)
if self.controller is not None and hasattr(self.controller, "reset"):
self.controller.reset()
logger.info("Reached default position") logger.info("Reached default position")
@@ -22,8 +22,6 @@ import zmq
from lerobot.robots.unitree_g1.config_unitree_g1 import UnitreeG1Config from lerobot.robots.unitree_g1.config_unitree_g1 import UnitreeG1Config
# Module-level ZMQ state mirrors the Unitree SDK's global ChannelFactory Singleton.
# Only one robot connection per process is supported.
_ctx: zmq.Context | None = None _ctx: zmq.Context | None = None
_lowcmd_sock: zmq.Socket | None = None _lowcmd_sock: zmq.Socket | None = None
_lowstate_sock: zmq.Socket | None = None _lowstate_sock: zmq.Socket | None = None
@@ -99,22 +97,17 @@ def lowcmd_to_dict(topic: str, msg: Any) -> dict[str, Any]:
} }
def ChannelFactoryInitialize(domain_id: int = 0, config: Any = None) -> None: # noqa: N802 def ChannelFactoryInitialize(*args: Any, **kwargs: Any) -> None: # noqa: N802
""" """
Initialize ZMQ sockets for robot communication. Initialize ZMQ sockets for robot communication.
This function mimics the Unitree SDK's ChannelFactoryInitialize but uses This function mimics the Unitree SDK's ChannelFactoryInitialize but uses
ZMQ sockets to connect to the robot server bridge instead of DDS. ZMQ sockets to connect to the robot server bridge instead of DDS.
Args:
domain_id: Ignored (for API compatibility with Unitree SDK)
config: UnitreeG1Config instance with robot_ip
""" """
global _ctx, _lowcmd_sock, _lowstate_sock global _ctx, _lowcmd_sock, _lowstate_sock
# read socket config # read socket config
if config is None: config = UnitreeG1Config()
config = UnitreeG1Config()
robot_ip = config.robot_ip robot_ip = config.robot_ip
ctx = zmq.Context.instance() ctx = zmq.Context.instance()
-1
View File
@@ -56,7 +56,6 @@ from lerobot.teleoperators import ( # noqa: F401
make_teleoperator_from_config, make_teleoperator_from_config,
omx_leader, omx_leader,
openarm_leader, openarm_leader,
openarm_mini,
so_leader, so_leader,
unitree_g1, unitree_g1,
) )
+1 -4
View File
@@ -132,13 +132,10 @@ def visualize_dataset(
logging.info("Logging to Rerun") logging.info("Logging to Rerun")
first_index = None
for batch in tqdm.tqdm(dataloader, total=len(dataloader)): for batch in tqdm.tqdm(dataloader, total=len(dataloader)):
if first_index is None:
first_index = batch["index"][0].item()
# iterate over the batch # iterate over the batch
for i in range(len(batch["index"])): for i in range(len(batch["index"])):
rr.set_time("frame_index", sequence=batch["index"][i].item() - first_index) rr.set_time("frame_index", sequence=batch["frame_index"][i].item())
rr.set_time("timestamp", timestamp=batch["timestamp"][i].item()) rr.set_time("timestamp", timestamp=batch["timestamp"][i].item())
# display each camera image # display each camera image
+59 -140
View File
@@ -21,9 +21,6 @@ This script allows you to delete episodes, split datasets, merge datasets,
remove features, modify tasks, and convert image datasets to video format. remove features, modify tasks, and convert image datasets to video format.
When new_repo_id is specified, creates a new dataset. When new_repo_id is specified, creates a new dataset.
Path semantics (v2): --root and --new_root are exact dataset folders containing
meta/, data/, videos/. When omitted, defaults to $HF_LEROBOT_HOME/{repo_id}.
Usage Examples: Usage Examples:
Delete episodes 0, 2, and 5 from a dataset: Delete episodes 0, 2, and 5 from a dataset:
@@ -32,34 +29,19 @@ Delete episodes 0, 2, and 5 from a dataset:
--operation.type delete_episodes \ --operation.type delete_episodes \
--operation.episode_indices "[0, 2, 5]" --operation.episode_indices "[0, 2, 5]"
Delete episodes from a local dataset at a specific path: Delete episodes and save to a new dataset:
lerobot-edit-dataset \
--repo_id lerobot/pusht \
--root /path/to/pusht \
--operation.type delete_episodes \
--operation.episode_indices "[0, 2, 5]"
Delete episodes and save to a new dataset at a specific path and with a new repo_id:
lerobot-edit-dataset \ lerobot-edit-dataset \
--repo_id lerobot/pusht \ --repo_id lerobot/pusht \
--new_repo_id lerobot/pusht_filtered \ --new_repo_id lerobot/pusht_filtered \
--new_root /path/to/pusht_filtered \
--operation.type delete_episodes \ --operation.type delete_episodes \
--operation.episode_indices "[0, 2, 5]" --operation.episode_indices "[0, 2, 5]"
Split dataset by fractions (pusht_train, pusht_val): Split dataset by fractions:
lerobot-edit-dataset \ lerobot-edit-dataset \
--repo_id lerobot/pusht \ --repo_id lerobot/pusht \
--operation.type split \ --operation.type split \
--operation.splits '{"train": 0.8, "val": 0.2}' --operation.splits '{"train": 0.8, "val": 0.2}'
Split dataset by fractions and save split datasets to a specific folder (base_folder/train, base_folder/val):
lerobot-edit-dataset \
--repo_id lerobot/pusht \
--new_root /path/to/base_folder \
--operation.type split \
--operation.splits '{"train": 0.8, "val": 0.2}'
Split dataset by episode indices: Split dataset by episode indices:
lerobot-edit-dataset \ lerobot-edit-dataset \
--repo_id lerobot/pusht \ --repo_id lerobot/pusht \
@@ -74,29 +56,15 @@ Split into more than two splits:
Merge multiple datasets: Merge multiple datasets:
lerobot-edit-dataset \ lerobot-edit-dataset \
--new_repo_id lerobot/pusht_merged \ --repo_id lerobot/pusht_merged \
--operation.type merge \ --operation.type merge \
--operation.repo_ids "['lerobot/pusht_train', 'lerobot/pusht_val']" --operation.repo_ids "['lerobot/pusht_train', 'lerobot/pusht_val']"
Merge multiple datasets to a specific output path:
lerobot-edit-dataset \
--new_repo_id lerobot/pusht_merged \
--new_root /path/to/pusht_merged \
--operation.type merge \
--operation.repo_ids "['lerobot/pusht_train', 'lerobot/pusht_val']"
Merge multiple datasets from a list of local dataset paths:
lerobot-edit-dataset \
--new_repo_id lerobot/pusht_merged \
--operation.type merge \
--operation.repo_ids "['pusht_train', 'pusht_val']" \
--operation.roots "['/path/to/pusht_train', '/path/to/pusht_val']"
Remove camera feature: Remove camera feature:
lerobot-edit-dataset \ lerobot-edit-dataset \
--repo_id lerobot/pusht \ --repo_id lerobot/pusht \
--operation.type remove_feature \ --operation.type remove_feature \
--operation.feature_names "['observation.image']" --operation.feature_names "['observation.images.top']"
Modify tasks - set a single task for all episodes (WARNING: modifies in-place): Modify tasks - set a single task for all episodes (WARNING: modifies in-place):
lerobot-edit-dataset \ lerobot-edit-dataset \
@@ -120,8 +88,8 @@ Modify tasks - set default task with overrides for specific episodes (WARNING: m
Convert image dataset to video format and save locally: Convert image dataset to video format and save locally:
lerobot-edit-dataset \ lerobot-edit-dataset \
--repo_id lerobot/pusht_image \ --repo_id lerobot/pusht_image \
--new_root /path/to/output/pusht_video \ --operation.type convert_image_to_video \
--operation.type convert_image_to_video --operation.output_dir /path/to/output/pusht_video
Convert image dataset to video format and save with new repo_id: Convert image dataset to video format and save with new repo_id:
lerobot-edit-dataset \ lerobot-edit-dataset \
@@ -199,7 +167,6 @@ class SplitConfig(OperationConfig):
@dataclass @dataclass
class MergeConfig(OperationConfig): class MergeConfig(OperationConfig):
repo_ids: list[str] | None = None repo_ids: list[str] | None = None
roots: list[str] | None = None
@OperationConfig.register_subclass("remove_feature") @OperationConfig.register_subclass("remove_feature")
@@ -233,46 +200,36 @@ class ConvertImageToVideoConfig(OperationConfig):
@OperationConfig.register_subclass("info") @OperationConfig.register_subclass("info")
@dataclass @dataclass
class InfoConfig(OperationConfig): class InfoConfig(OperationConfig):
type: str = "info"
show_features: bool = False show_features: bool = False
@dataclass @dataclass
class EditDatasetConfig: class EditDatasetConfig:
# Operation configuration. repo_id: str
operation: OperationConfig operation: OperationConfig
# Input dataset identifier. Always required unless for Merge operation.
repo_id: str | None = None
# Root directory where the input dataset is stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id.
root: str | None = None root: str | None = None
# Edited dataset identifier. When both new_repo_id (resp. new_root) and repo_id (resp. root) are identical, modifications are applied in-place and a backup of the original dataset is created. Required for Merge operation.
new_repo_id: str | None = None new_repo_id: str | None = None
# Root directory where the edited dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/new_repo_id. For Split operation, this is the base directory for the split datasets.
new_root: str | None = None
# Upload dataset to Hugging Face hub.
push_to_hub: bool = False push_to_hub: bool = False
def get_output_path( def get_output_path(repo_id: str, new_repo_id: str | None, root: Path | None) -> tuple[str, Path]:
repo_id: str, if new_repo_id:
new_repo_id: str | None, output_repo_id = new_repo_id
root: Path | str | None, output_dir = root / new_repo_id if root else HF_LEROBOT_HOME / new_repo_id
new_root: Path | str | None, else:
) -> tuple[str, Path]: output_repo_id = repo_id
input_path = Path(root) if root else HF_LEROBOT_HOME / repo_id dataset_path = root / repo_id if root else HF_LEROBOT_HOME / repo_id
old_path = Path(str(dataset_path) + "_old")
output_repo_id = new_repo_id if new_repo_id else repo_id if dataset_path.exists():
output_path = Path(new_root) if new_root else HF_LEROBOT_HOME / output_repo_id if old_path.exists():
shutil.rmtree(old_path)
shutil.move(str(dataset_path), str(old_path))
# In case of in-place modification, create a backup of the original dataset (if it exists) output_dir = dataset_path
if output_path == input_path:
backup_path = input_path.with_name(input_path.name + "_old")
if input_path.exists(): return output_repo_id, output_dir
if backup_path.exists():
shutil.rmtree(backup_path)
shutil.move(input_path, backup_path)
return output_repo_id, output_path
def handle_delete_episodes(cfg: EditDatasetConfig) -> None: def handle_delete_episodes(cfg: EditDatasetConfig) -> None:
@@ -284,15 +241,11 @@ def handle_delete_episodes(cfg: EditDatasetConfig) -> None:
dataset = LeRobotDataset(cfg.repo_id, root=cfg.root) dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
output_repo_id, output_dir = get_output_path( output_repo_id, output_dir = get_output_path(
cfg.repo_id, cfg.repo_id, cfg.new_repo_id, Path(cfg.root) if cfg.root else None
new_repo_id=cfg.new_repo_id,
root=cfg.root,
new_root=cfg.new_root,
) )
# In case of in-place modification, make the dataset point to the backup directory if cfg.new_repo_id is None:
if output_dir == dataset.root: dataset.root = Path(str(dataset.root) + "_old")
dataset.root = dataset.root.with_name(dataset.root.name + "_old")
logging.info(f"Deleting episodes {cfg.operation.episode_indices} from {cfg.repo_id}") logging.info(f"Deleting episodes {cfg.operation.episode_indices} from {cfg.repo_id}")
new_dataset = delete_episodes( new_dataset = delete_episodes(
@@ -319,27 +272,19 @@ def handle_split(cfg: EditDatasetConfig) -> None:
"splits dict must be specified with split names as keys and fractions/episode lists as values" "splits dict must be specified with split names as keys and fractions/episode lists as values"
) )
if cfg.new_repo_id is not None:
logging.warning(
"split uses the original dataset identifier --repo_id to generate split names. The --new_repo_id parameter is ignored."
)
dataset = LeRobotDataset(cfg.repo_id, root=cfg.root) dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
logging.info(f"Splitting dataset {cfg.repo_id} with splits: {cfg.operation.splits}") logging.info(f"Splitting dataset {cfg.repo_id} with splits: {cfg.operation.splits}")
split_datasets = split_dataset( split_datasets = split_dataset(dataset, splits=cfg.operation.splits)
dataset,
splits=cfg.operation.splits,
output_dir=cfg.new_root,
)
for split_name, split_ds in split_datasets.items(): for split_name, split_ds in split_datasets.items():
split_repo_id = f"{cfg.repo_id}_{split_name}"
logging.info( logging.info(
f"{split_name}: {split_ds.meta.total_episodes} episodes, {split_ds.meta.total_frames} frames" f"{split_name}: {split_ds.meta.total_episodes} episodes, {split_ds.meta.total_frames} frames"
) )
if cfg.push_to_hub: if cfg.push_to_hub:
logging.info(f"Pushing {split_name} split to hub as {split_ds.repo_id}") logging.info(f"Pushing {split_name} split to hub as {split_repo_id}")
LeRobotDataset(split_ds.repo_id, root=split_ds.root).push_to_hub() LeRobotDataset(split_ds.repo_id, root=split_ds.root).push_to_hub()
@@ -350,29 +295,18 @@ def handle_merge(cfg: EditDatasetConfig) -> None:
if not cfg.operation.repo_ids: if not cfg.operation.repo_ids:
raise ValueError("repo_ids must be specified for merge operation") raise ValueError("repo_ids must be specified for merge operation")
if cfg.repo_id is not None or cfg.root is not None: if not cfg.repo_id:
logging.warning( raise ValueError("repo_id must be specified as the output repository for merged dataset")
"merge uses --new_repo_id and --new_root for the merged dataset. The --repo_id and --root parameters are ignored."
)
if cfg.operation.roots: logging.info(f"Loading {len(cfg.operation.repo_ids)} datasets to merge")
if len(cfg.operation.roots) != len(cfg.operation.repo_ids): datasets = [LeRobotDataset(repo_id, root=cfg.root) for repo_id in cfg.operation.repo_ids]
raise ValueError("repo_ids and roots must have the same length for merge operation")
logging.info(f"Loading {len(cfg.operation.roots)} datasets to merge")
datasets = [
LeRobotDataset(repo_id=repo_id, root=root)
for repo_id, root in zip(cfg.operation.repo_ids, cfg.operation.roots, strict=True)
]
else:
logging.info(f"Loading {len(cfg.operation.repo_ids)} datasets to merge")
datasets = [LeRobotDataset(repo_id) for repo_id in cfg.operation.repo_ids]
output_dir = Path(cfg.new_root) if cfg.new_root else HF_LEROBOT_HOME / cfg.new_repo_id output_dir = Path(cfg.root) / cfg.repo_id if cfg.root else HF_LEROBOT_HOME / cfg.repo_id
logging.info(f"Merging datasets into {cfg.new_repo_id}") logging.info(f"Merging datasets into {cfg.repo_id}")
merged_dataset = merge_datasets( merged_dataset = merge_datasets(
datasets, datasets,
output_repo_id=cfg.new_repo_id, output_repo_id=cfg.repo_id,
output_dir=output_dir, output_dir=output_dir,
) )
@@ -382,7 +316,7 @@ def handle_merge(cfg: EditDatasetConfig) -> None:
) )
if cfg.push_to_hub: if cfg.push_to_hub:
logging.info(f"Pushing to hub as {cfg.new_repo_id}") logging.info(f"Pushing to hub as {cfg.repo_id}")
LeRobotDataset(merged_dataset.repo_id, root=output_dir).push_to_hub() LeRobotDataset(merged_dataset.repo_id, root=output_dir).push_to_hub()
@@ -395,15 +329,11 @@ def handle_remove_feature(cfg: EditDatasetConfig) -> None:
dataset = LeRobotDataset(cfg.repo_id, root=cfg.root) dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
output_repo_id, output_dir = get_output_path( output_repo_id, output_dir = get_output_path(
cfg.repo_id, cfg.repo_id, cfg.new_repo_id, Path(cfg.root) if cfg.root else None
new_repo_id=cfg.new_repo_id,
root=cfg.root,
new_root=cfg.new_root,
) )
# In case of in-place modification, make the dataset point to the backup directory if cfg.new_repo_id is None:
if output_dir == dataset.root: dataset.root = Path(str(dataset.root) + "_old")
dataset.root = dataset.root.with_name(dataset.root.name + "_old")
logging.info(f"Removing features {cfg.operation.feature_names} from {cfg.repo_id}") logging.info(f"Removing features {cfg.operation.feature_names} from {cfg.repo_id}")
new_dataset = remove_feature( new_dataset = remove_feature(
@@ -431,10 +361,9 @@ def handle_modify_tasks(cfg: EditDatasetConfig) -> None:
if new_task is None and episode_tasks_raw is None: if new_task is None and episode_tasks_raw is None:
raise ValueError("Must specify at least one of new_task or episode_tasks for modify_tasks operation") raise ValueError("Must specify at least one of new_task or episode_tasks for modify_tasks operation")
if cfg.new_repo_id is not None or cfg.new_root is not None: # Warn about in-place modification behavior
logging.warning( if cfg.new_repo_id is not None:
"modify_tasks modifies datasets in-place. The --new_repo_id and --new_root parameters are ignored." logging.warning("modify_tasks modifies datasets in-place. The --new_repo_id parameter is ignored.")
)
dataset = LeRobotDataset(cfg.repo_id, root=cfg.root) dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
logging.warning(f"Modifying dataset in-place at {dataset.root}. Original data will be overwritten.") logging.warning(f"Modifying dataset in-place at {dataset.root}. Original data will be overwritten.")
@@ -470,30 +399,32 @@ def handle_convert_image_to_video(cfg: EditDatasetConfig) -> None:
dataset = LeRobotDataset(cfg.repo_id, root=cfg.root) dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
# Determine output directory and repo_id # Determine output directory and repo_id
# Priority: 1) new_root, 2) new_repo_id, 3) operation.output_dir, 4) auto-generated name # Priority: 1) new_repo_id, 2) operation.output_dir, 3) auto-generated name
output_dir_config = getattr(cfg.operation, "output_dir", None) output_dir_config = getattr(cfg.operation, "output_dir", None)
if output_dir_config:
logging.warning(
"--operation.output_dir is deprecated and will be removed in future versions. "
"Please use --new_root instead."
)
if cfg.new_root: if cfg.new_repo_id:
output_dir = Path(cfg.new_root) # Use new_repo_id for both local storage and hub push
output_repo_id = cfg.new_repo_id or f"{cfg.repo_id}_video"
logging.info(f"Saving to new_root: {output_dir} as {output_repo_id}")
elif cfg.new_repo_id:
output_repo_id = cfg.new_repo_id output_repo_id = cfg.new_repo_id
output_dir = HF_LEROBOT_HOME / cfg.new_repo_id # Place new dataset as a sibling to the original dataset
# Get the parent of the actual dataset root (not cfg.root which might be the lerobot cache dir)
# Extract just the dataset name (after last slash) for the local directory
local_dir_name = cfg.new_repo_id.split("/")[-1]
output_dir = dataset.root.parent / local_dir_name
logging.info(f"Saving to new dataset: {cfg.new_repo_id} at {output_dir}") logging.info(f"Saving to new dataset: {cfg.new_repo_id} at {output_dir}")
elif output_dir_config: elif output_dir_config:
# Use custom output directory for local-only storage
output_dir = Path(output_dir_config) output_dir = Path(output_dir_config)
# Extract repo name from output_dir for the dataset
output_repo_id = output_dir.name output_repo_id = output_dir.name
logging.info(f"Saving to local directory: {output_dir} as {output_repo_id}") logging.info(f"Saving to local directory: {output_dir}")
else: else:
# Auto-generate name: append "_video" to original repo_id
output_repo_id = f"{cfg.repo_id}_video" output_repo_id = f"{cfg.repo_id}_video"
output_dir = HF_LEROBOT_HOME / output_repo_id # Place new dataset as a sibling to the original dataset
logging.info(f"Saving to auto-generated location: {output_dir} as {output_repo_id}") # Extract just the dataset name (after last slash) for the local directory
local_dir_name = output_repo_id.split("/")[-1]
output_dir = dataset.root.parent / local_dir_name
logging.info(f"Saving to auto-generated location: {output_dir}")
logging.info(f"Converting dataset {cfg.repo_id} to video format") logging.info(f"Converting dataset {cfg.repo_id} to video format")
@@ -568,20 +499,8 @@ def handle_info(cfg: EditDatasetConfig):
sys.stdout.write(f"{feature_dump_str}\n") sys.stdout.write(f"{feature_dump_str}\n")
def _validate_config(cfg: EditDatasetConfig) -> None:
if isinstance(cfg.operation, MergeConfig):
if not cfg.new_repo_id:
raise ValueError("--new_repo_id is required for merge operation (the merged dataset identifier)")
else:
if not cfg.repo_id:
raise ValueError(
f"--repo_id is required for {cfg.operation.type} operation (the input dataset identifier)"
)
@parser.wrap() @parser.wrap()
def edit_dataset(cfg: EditDatasetConfig) -> None: def edit_dataset(cfg: EditDatasetConfig) -> None:
_validate_config(cfg)
operation_type = cfg.operation.type operation_type = cfg.operation.type
if operation_type == "delete_episodes": if operation_type == "delete_episodes":
@@ -61,7 +61,6 @@ from lerobot.teleoperators import ( # noqa: F401
make_teleoperator_from_config, make_teleoperator_from_config,
omx_leader, omx_leader,
openarm_leader, openarm_leader,
openarm_mini,
so_leader, so_leader,
) )
from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.robot_utils import precise_sleep
+10 -12
View File
@@ -125,7 +125,6 @@ from lerobot.teleoperators import ( # noqa: F401
make_teleoperator_from_config, make_teleoperator_from_config,
omx_leader, omx_leader,
openarm_leader, openarm_leader,
openarm_mini,
reachy2_teleoperator, reachy2_teleoperator,
so_leader, so_leader,
unitree_g1, unitree_g1,
@@ -155,7 +154,7 @@ class DatasetRecordConfig:
repo_id: str repo_id: str
# A short but accurate description of the task performed during the recording (e.g. "Pick the Lego block and drop it in the box on the right.") # A short but accurate description of the task performed during the recording (e.g. "Pick the Lego block and drop it in the box on the right.")
single_task: str single_task: str
# Root directory where the dataset will be stored (e.g. 'dataset/path'). If None, defaults to $HF_LEROBOT_HOME/repo_id. # Root directory where the dataset will be stored (e.g. 'dataset/path').
root: str | Path | None = None root: str | Path | None = None
# Limit the frames per second. # Limit the frames per second.
fps: int = 30 fps: int = 30
@@ -334,7 +333,6 @@ def record_loop(
preprocessor.reset() preprocessor.reset()
postprocessor.reset() postprocessor.reset()
no_action_count = 0
timestamp = 0 timestamp = 0
start_episode_t = time.perf_counter() start_episode_t = time.perf_counter()
while timestamp < control_time_s: while timestamp < control_time_s:
@@ -369,8 +367,6 @@ def record_loop(
act_processed_policy: RobotAction = make_robot_action(action_values, dataset.features) act_processed_policy: RobotAction = make_robot_action(action_values, dataset.features)
elif policy is None and isinstance(teleop, Teleoperator): elif policy is None and isinstance(teleop, Teleoperator):
if robot.name == "unitree_g1":
teleop.send_feedback(obs)
act = teleop.get_action() act = teleop.get_action()
# Applies a pipeline to the raw teleop action, default is IdentityProcessor # Applies a pipeline to the raw teleop action, default is IdentityProcessor
@@ -384,13 +380,11 @@ def record_loop(
act = {**arm_action, **base_action} if len(base_action) > 0 else arm_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, obs))
else: else:
no_action_count += 1 logging.info(
if no_action_count == 1 or no_action_count % 10 == 0: "No policy or teleoperator provided, skipping action generation."
logging.warning( "This is likely to happen when resetting the environment without a teleop device."
"No policy or teleoperator provided, skipping action generation. " "The robot won't be at its rest position at the start of the next episode."
"This is likely to happen when resetting the environment without a teleop device. " )
"The robot won't be at its rest position at the start of the next episode."
)
continue continue
# Applies a pipeline to the action, default is IdentityProcessor # Applies a pipeline to the action, default is IdentityProcessor
@@ -558,6 +552,10 @@ def record(cfg: RecordConfig) -> LeRobotDataset:
): ):
log_say("Reset the environment", cfg.play_sounds) log_say("Reset the environment", cfg.play_sounds)
# reset g1 robot
if robot.name == "unitree_g1":
robot.reset()
record_loop( record_loop(
robot=robot, robot=robot,
events=events, events=events,
+1 -1
View File
@@ -80,7 +80,7 @@ class DatasetReplayConfig:
repo_id: str repo_id: str
# Episode to replay. # Episode to replay.
episode: int episode: int
# Root directory where the dataset will be stored (e.g. 'dataset/path'). If None, defaults to $HF_LEROBOT_HOME/repo_id. # Root directory where the dataset will be stored (e.g. 'dataset/path').
root: str | Path | None = None root: str | Path | None = None
# Limit the frames per second. By default, uses the policy fps. # Limit the frames per second. By default, uses the policy fps.
fps: int = 30 fps: int = 30
@@ -43,7 +43,6 @@ from lerobot.teleoperators import ( # noqa: F401
koch_leader, koch_leader,
make_teleoperator_from_config, make_teleoperator_from_config,
omx_leader, omx_leader,
openarm_mini,
so_leader, so_leader,
) )
@@ -52,7 +51,6 @@ COMPATIBLE_DEVICES = [
"koch_leader", "koch_leader",
"omx_follower", "omx_follower",
"omx_leader", "omx_leader",
"openarm_mini",
"so100_follower", "so100_follower",
"so100_leader", "so100_leader",
"so101_follower", "so101_follower",
+1 -5
View File
@@ -60,7 +60,6 @@ import rerun as rr
from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig # noqa: F401 from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig # noqa: F401
from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig # noqa: F401 from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig # noqa: F401
from lerobot.cameras.zmq.configuration_zmq import ZMQCameraConfig # noqa: F401
from lerobot.configs import parser from lerobot.configs import parser
from lerobot.processor import ( from lerobot.processor import (
RobotAction, RobotAction,
@@ -95,7 +94,6 @@ from lerobot.teleoperators import ( # noqa: F401
make_teleoperator_from_config, make_teleoperator_from_config,
omx_leader, omx_leader,
openarm_leader, openarm_leader,
openarm_mini,
reachy2_teleoperator, reachy2_teleoperator,
so_leader, so_leader,
unitree_g1, unitree_g1,
@@ -154,6 +152,7 @@ def teleop_loop(
display_len = max(len(key) for key in robot.action_features) display_len = max(len(key) for key in robot.action_features)
start = time.perf_counter() start = time.perf_counter()
while True: while True:
loop_start = time.perf_counter() loop_start = time.perf_counter()
@@ -163,9 +162,6 @@ def teleop_loop(
# given that it is the identity processor as default # given that it is the identity processor as default
obs = robot.get_observation() obs = robot.get_observation()
if robot.name == "unitree_g1":
teleop.send_feedback(obs)
# Get teleop action # Get teleop action
raw_action = teleop.get_action() raw_action = teleop.get_action()
+7 -39
View File
@@ -24,13 +24,11 @@ import torch
from accelerate import Accelerator from accelerate import Accelerator
from termcolor import colored from termcolor import colored
from torch.optim import Optimizer from torch.optim import Optimizer
from tqdm import tqdm
from lerobot.configs import parser from lerobot.configs import parser
from lerobot.configs.train import TrainPipelineConfig from lerobot.configs.train import TrainPipelineConfig
from lerobot.datasets.factory import make_dataset from lerobot.datasets.factory import make_dataset
from lerobot.datasets.multi_dataset import NewMultiLeRobotDataset from lerobot.datasets.sampler import EpisodeAwareSampler
from lerobot.datasets.sampler import EpisodeAwareSampler, WeightedEpisodeAwareSampler
from lerobot.datasets.utils import cycle from lerobot.datasets.utils import cycle
from lerobot.envs.factory import make_env, make_env_pre_post_processors from lerobot.envs.factory import make_env, make_env_pre_post_processors
from lerobot.envs.utils import close_envs from lerobot.envs.utils import close_envs
@@ -53,7 +51,6 @@ from lerobot.utils.utils import (
format_big_number, format_big_number,
has_method, has_method,
init_logging, init_logging,
inside_slurm,
) )
@@ -210,11 +207,7 @@ def train(cfg: TrainPipelineConfig, accelerator: Accelerator | None = None):
# Use accelerator's device # Use accelerator's device
device = accelerator.device device = accelerator.device
if cfg.cudnn_deterministic: torch.backends.cudnn.benchmark = True
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
else:
torch.backends.cudnn.benchmark = True
torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cuda.matmul.allow_tf32 = True
# Dataset loading synchronization: main process downloads first to avoid race conditions # Dataset loading synchronization: main process downloads first to avoid race conditions
@@ -344,25 +337,13 @@ def train(cfg: TrainPipelineConfig, accelerator: Accelerator | None = None):
logging.info(f"{num_total_params=} ({format_big_number(num_total_params)})") logging.info(f"{num_total_params=} ({format_big_number(num_total_params)})")
# create dataloader for offline training # create dataloader for offline training
drop_n_last = getattr(cfg.policy, "drop_n_last_frames", 0) if hasattr(cfg.policy, "drop_n_last_frames"):
if isinstance(dataset, NewMultiLeRobotDataset):
shuffle = False
sampler = WeightedEpisodeAwareSampler(
dataset.meta.episodes["dataset_from_index"],
dataset.meta.episodes["dataset_to_index"],
dataset_membership=dataset.meta.episodes["dataset_source"],
dataset_weights=dataset.dataset_weights,
drop_n_last_frames=drop_n_last,
shuffle=True,
)
elif drop_n_last > 0:
shuffle = False shuffle = False
sampler = EpisodeAwareSampler( sampler = EpisodeAwareSampler(
dataset.meta.episodes["dataset_from_index"], dataset.meta.episodes["dataset_from_index"],
dataset.meta.episodes["dataset_to_index"], dataset.meta.episodes["dataset_to_index"],
episode_indices_to_use=dataset.episodes, episode_indices_to_use=dataset.episodes,
drop_n_last_frames=drop_n_last, drop_n_last_frames=cfg.policy.drop_n_last_frames,
shuffle=True, shuffle=True,
) )
else: else:
@@ -373,7 +354,7 @@ def train(cfg: TrainPipelineConfig, accelerator: Accelerator | None = None):
dataset, dataset,
num_workers=cfg.num_workers, num_workers=cfg.num_workers,
batch_size=cfg.batch_size, batch_size=cfg.batch_size,
shuffle=shuffle and not getattr(cfg.dataset, "streaming", False), shuffle=shuffle and not cfg.dataset.streaming,
sampler=sampler, sampler=sampler,
pin_memory=device.type == "cuda", pin_memory=device.type == "cuda",
drop_last=False, drop_last=False,
@@ -397,10 +378,10 @@ def train(cfg: TrainPipelineConfig, accelerator: Accelerator | None = None):
"dataloading_s": AverageMeter("data_s", ":.3f"), "dataloading_s": AverageMeter("data_s", ":.3f"),
} }
# Keep global batch size for logging; MetricsTracker handles world size internally. # Use effective batch size for proper epoch calculation in distributed training
effective_batch_size = cfg.batch_size * accelerator.num_processes effective_batch_size = cfg.batch_size * accelerator.num_processes
train_tracker = MetricsTracker( train_tracker = MetricsTracker(
cfg.batch_size, effective_batch_size,
dataset.num_frames, dataset.num_frames,
dataset.num_episodes, dataset.num_episodes,
train_metrics, train_metrics,
@@ -409,14 +390,6 @@ def train(cfg: TrainPipelineConfig, accelerator: Accelerator | None = None):
) )
if is_main_process: if is_main_process:
progbar = tqdm(
total=cfg.steps - step,
desc="Training",
unit="step",
disable=inside_slurm(),
position=0,
leave=True,
)
logging.info( logging.info(
f"Start offline training on a fixed dataset, with effective batch size: {effective_batch_size}" f"Start offline training on a fixed dataset, with effective batch size: {effective_batch_size}"
) )
@@ -441,8 +414,6 @@ def train(cfg: TrainPipelineConfig, accelerator: Accelerator | None = None):
# Note: eval and checkpoint happens *after* the `step`th training update has completed, so we # Note: eval and checkpoint happens *after* the `step`th training update has completed, so we
# increment `step` here. # increment `step` here.
step += 1 step += 1
if is_main_process:
progbar.update(1)
train_tracker.step() train_tracker.step()
is_log_step = cfg.log_freq > 0 and step % cfg.log_freq == 0 and is_main_process is_log_step = cfg.log_freq > 0 and step % cfg.log_freq == 0 and is_main_process
is_saving_step = step % cfg.save_freq == 0 or step == cfg.steps is_saving_step = step % cfg.save_freq == 0 or step == cfg.steps
@@ -536,9 +507,6 @@ def train(cfg: TrainPipelineConfig, accelerator: Accelerator | None = None):
accelerator.wait_for_everyone() accelerator.wait_for_everyone()
if is_main_process:
progbar.close()
if eval_env: if eval_env:
close_envs(eval_env) close_envs(eval_env)
@@ -1,20 +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.
from .config_openarm_mini import OpenArmMiniConfig
from .openarm_mini import OpenArmMini
__all__ = ["OpenArmMini", "OpenArmMiniConfig"]
@@ -1,30 +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.
from dataclasses import dataclass
from ..config import TeleoperatorConfig
@TeleoperatorConfig.register_subclass("openarm_mini")
@dataclass
class OpenArmMiniConfig(TeleoperatorConfig):
"""Configuration for OpenArm Mini teleoperator with Feetech motors (dual arms)."""
port_right: str = "/dev/ttyUSB0"
port_left: str = "/dev/ttyUSB1"
use_degrees: bool = True
@@ -1,296 +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 logging
import time
from typing import Any
from lerobot.motors import Motor, MotorCalibration, MotorNormMode
from lerobot.motors.feetech import (
FeetechMotorsBus,
OperatingMode,
)
from lerobot.processor import RobotAction
from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
from ..teleoperator import Teleoperator
from .config_openarm_mini import OpenArmMiniConfig
logger = logging.getLogger(__name__)
# Motors whose direction is inverted during readout
RIGHT_MOTORS_TO_FLIP = ["joint_1", "joint_2", "joint_3", "joint_4", "joint_5"]
LEFT_MOTORS_TO_FLIP = ["joint_1", "joint_3", "joint_4", "joint_5", "joint_6", "joint_7"]
class OpenArmMini(Teleoperator):
"""
OpenArm Mini Teleoperator with dual Feetech-based arms (8 motors per arm).
Each arm has 7 joints plus a gripper, using Feetech STS3215 servos.
"""
config_class = OpenArmMiniConfig
name = "openarm_mini"
def __init__(self, config: OpenArmMiniConfig):
super().__init__(config)
self.config = config
norm_mode_body = MotorNormMode.DEGREES
motors_right = {
"joint_1": Motor(1, "sts3215", norm_mode_body),
"joint_2": Motor(2, "sts3215", norm_mode_body),
"joint_3": Motor(3, "sts3215", norm_mode_body),
"joint_4": Motor(4, "sts3215", norm_mode_body),
"joint_5": Motor(5, "sts3215", norm_mode_body),
"joint_6": Motor(6, "sts3215", norm_mode_body),
"joint_7": Motor(7, "sts3215", norm_mode_body),
"gripper": Motor(8, "sts3215", MotorNormMode.RANGE_0_100),
}
motors_left = {
"joint_1": Motor(1, "sts3215", norm_mode_body),
"joint_2": Motor(2, "sts3215", norm_mode_body),
"joint_3": Motor(3, "sts3215", norm_mode_body),
"joint_4": Motor(4, "sts3215", norm_mode_body),
"joint_5": Motor(5, "sts3215", norm_mode_body),
"joint_6": Motor(6, "sts3215", norm_mode_body),
"joint_7": Motor(7, "sts3215", norm_mode_body),
"gripper": Motor(8, "sts3215", MotorNormMode.RANGE_0_100),
}
cal_right = {
k.replace("right_", ""): v for k, v in (self.calibration or {}).items() if k.startswith("right_")
}
cal_left = {
k.replace("left_", ""): v for k, v in (self.calibration or {}).items() if k.startswith("left_")
}
self.bus_right = FeetechMotorsBus(
port=self.config.port_right,
motors=motors_right,
calibration=cal_right,
)
self.bus_left = FeetechMotorsBus(
port=self.config.port_left,
motors=motors_left,
calibration=cal_left,
)
@property
def action_features(self) -> dict[str, type]:
features: dict[str, type] = {}
for motor in self.bus_right.motors:
features[f"right_{motor}.pos"] = float
for motor in self.bus_left.motors:
features[f"left_{motor}.pos"] = float
return features
@property
def feedback_features(self) -> dict[str, type]:
return {}
@property
def is_connected(self) -> bool:
return self.bus_right.is_connected and self.bus_left.is_connected
@check_if_already_connected
def connect(self, calibrate: bool = True) -> None:
logger.info(f"Connecting right arm on {self.config.port_right}...")
self.bus_right.connect()
logger.info(f"Connecting left arm on {self.config.port_left}...")
self.bus_left.connect()
if calibrate:
self.calibrate()
self.configure()
logger.info(f"{self} connected.")
@property
def is_calibrated(self) -> bool:
return self.bus_right.is_calibrated and self.bus_left.is_calibrated
def calibrate(self) -> None:
"""
Run calibration procedure for OpenArm Mini.
1. Disable torque
2. Ask user to position arms in hanging position with grippers closed
3. Set this as zero position via half-turn homing
4. Interactive gripper calibration (open/close positions)
5. Save calibration
"""
if self.calibration:
user_input = input(
f"Press ENTER to use existing calibration for {self.id}, "
f"or type 'c' and press ENTER to run new calibration: "
)
if user_input.strip().lower() != "c":
logger.info(f"Using existing calibration for {self.id}")
cal_right = {
k.replace("right_", ""): v for k, v in self.calibration.items() if k.startswith("right_")
}
cal_left = {
k.replace("left_", ""): v for k, v in self.calibration.items() if k.startswith("left_")
}
self.bus_right.write_calibration(cal_right)
self.bus_left.write_calibration(cal_left)
return
logger.info(f"\nRunning calibration for {self}")
self._calibrate_arm("right", self.bus_right)
self._calibrate_arm("left", self.bus_left)
self._save_calibration()
print(f"\nCalibration complete and saved to {self.calibration_fpath}")
def _calibrate_arm(self, arm_name: str, bus: FeetechMotorsBus) -> None:
"""Calibrate a single arm with Feetech motors."""
logger.info(f"\n=== Calibrating {arm_name.upper()} arm ===")
bus.disable_torque()
logger.info(f"Setting Phase to 12 for all motors in {arm_name.upper()} arm...")
for motor in bus.motors:
bus.write("Phase", motor, 12)
for motor in bus.motors:
bus.write("Operating_Mode", motor, OperatingMode.POSITION.value)
input(
f"\nCalibration: Zero Position ({arm_name.upper()} arm)\n"
"Position the arm in the following configuration:\n"
" - Arm hanging straight down\n"
" - Gripper closed\n"
"Press ENTER when ready..."
)
homing_offsets = bus.set_half_turn_homings()
logger.info(f"{arm_name.capitalize()} arm zero position set.")
print(f"\nSetting motor ranges for {arm_name.upper()} arm\n")
if self.calibration is None:
self.calibration = {}
motor_resolution = bus.model_resolution_table[list(bus.motors.values())[0].model]
max_res = motor_resolution - 1
for motor_name, motor in bus.motors.items():
prefixed_name = f"{arm_name}_{motor_name}"
if motor_name == "gripper":
input(
f"\nGripper Calibration ({arm_name.upper()} arm)\n"
f"Step 1: CLOSE the gripper fully\n"
f"Press ENTER when gripper is closed..."
)
closed_pos = bus.read("Present_Position", motor_name, normalize=False)
logger.info(f" Gripper closed position recorded: {closed_pos}")
input("\nStep 2: OPEN the gripper fully\nPress ENTER when gripper is fully open...")
open_pos = bus.read("Present_Position", motor_name, normalize=False)
logger.info(f" Gripper open position recorded: {open_pos}")
if closed_pos < open_pos:
range_min = int(closed_pos)
range_max = int(open_pos)
drive_mode = 0
else:
range_min = int(open_pos)
range_max = int(closed_pos)
drive_mode = 1
logger.info(
f" {prefixed_name}: range set to [{range_min}, {range_max}] "
f"(0=closed, 100=open, drive_mode={drive_mode})"
)
else:
range_min = 0
range_max = max_res
drive_mode = 0
logger.info(f" {prefixed_name}: range set to [0, {max_res}] (full motor range)")
self.calibration[prefixed_name] = MotorCalibration(
id=motor.id,
drive_mode=drive_mode,
homing_offset=homing_offsets[motor_name],
range_min=range_min,
range_max=range_max,
)
cal_for_bus = {
k.replace(f"{arm_name}_", ""): v
for k, v in self.calibration.items()
if k.startswith(f"{arm_name}_")
}
bus.write_calibration(cal_for_bus)
def configure(self) -> None:
self.bus_right.disable_torque()
self.bus_right.configure_motors()
for motor in self.bus_right.motors:
self.bus_right.write("Operating_Mode", motor, OperatingMode.POSITION.value)
self.bus_left.disable_torque()
self.bus_left.configure_motors()
for motor in self.bus_left.motors:
self.bus_left.write("Operating_Mode", motor, OperatingMode.POSITION.value)
def setup_motors(self) -> None:
print("\nSetting up RIGHT arm motors...")
for motor in reversed(self.bus_right.motors):
input(f"Connect the controller board to the RIGHT '{motor}' motor only and press enter.")
self.bus_right.setup_motor(motor)
print(f"RIGHT '{motor}' motor id set to {self.bus_right.motors[motor].id}")
print("\nSetting up LEFT arm motors...")
for motor in reversed(self.bus_left.motors):
input(f"Connect the controller board to the LEFT '{motor}' motor only and press enter.")
self.bus_left.setup_motor(motor)
print(f"LEFT '{motor}' motor id set to {self.bus_left.motors[motor].id}")
@check_if_not_connected
def get_action(self) -> RobotAction:
"""Get current action from both arms (read positions from all motors)."""
start = time.perf_counter()
right_positions = self.bus_right.sync_read("Present_Position")
left_positions = self.bus_left.sync_read("Present_Position")
action: dict[str, Any] = {}
for motor, val in right_positions.items():
action[f"right_{motor}.pos"] = -val if motor in RIGHT_MOTORS_TO_FLIP else val
for motor, val in left_positions.items():
action[f"left_{motor}.pos"] = -val if motor in LEFT_MOTORS_TO_FLIP else val
dt_ms = (time.perf_counter() - start) * 1e3
logger.debug(f"{self} read action: {dt_ms:.1f}ms")
return action
def send_feedback(self, feedback: dict[str, float]) -> None:
raise NotImplementedError("Feedback is not yet implemented for OpenArm Mini.")
@check_if_not_connected
def disconnect(self) -> None:
self.bus_right.disconnect()
self.bus_left.disconnect()
logger.info(f"{self} disconnected.")
@@ -15,6 +15,7 @@
# limitations under the License. # limitations under the License.
from dataclasses import dataclass from dataclasses import dataclass
from typing import TypeAlias
from ..config import TeleoperatorConfig from ..config import TeleoperatorConfig
@@ -37,5 +38,5 @@ class SOLeaderTeleopConfig(TeleoperatorConfig, SOLeaderConfig):
pass pass
SO100LeaderConfig = SOLeaderTeleopConfig SO100LeaderConfig: TypeAlias = SOLeaderTeleopConfig
SO101LeaderConfig = SOLeaderTeleopConfig SO101LeaderConfig: TypeAlias = SOLeaderTeleopConfig
@@ -16,6 +16,7 @@
import logging import logging
import time import time
from typing import TypeAlias
from lerobot.motors import Motor, MotorCalibration, MotorNormMode from lerobot.motors import Motor, MotorCalibration, MotorNormMode
from lerobot.motors.feetech import ( from lerobot.motors.feetech import (
@@ -155,5 +156,5 @@ class SOLeader(Teleoperator):
logger.info(f"{self} disconnected.") logger.info(f"{self} disconnected.")
SO100Leader = SOLeader SO100Leader: TypeAlias = SOLeader
SO101Leader = SOLeader SO101Leader: TypeAlias = SOLeader

Some files were not shown because too many files have changed in this diff Show More