test(cameras): add opencv camera dependency injection tests suite

for more information, see https://pre-commit.ci

.dockerignore

@@ -1,3 +1,17 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 # Misc
 .git
 tmp
@@ -59,7 +73,7 @@ pip-log.txt
 pip-delete-this-directory.txt
 
 # Unit test / coverage reports
-!tests/data
+!tests/artifacts
 htmlcov/
 .tox/
 .nox/

.gitattributes

@@ -1,3 +1,17 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 *.memmap filter=lfs diff=lfs merge=lfs -text
 *.stl filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text

.github/ISSUE_TEMPLATE/bug-report.yml

@@ -1,3 +1,17 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 name: "\U0001F41B Bug Report"
 description: Submit a bug report to help us improve LeRobot
 body:

.github/workflows/build-docker-images.yml

@@ -1,3 +1,17 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 # Inspired by
 # https://github.com/huggingface/peft/blob/main/.github/workflows/build_docker_images.yml
 name: Builds

.github/workflows/nightly-tests.yml

@@ -1,3 +1,17 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 # Inspired by
 # https://github.com/huggingface/peft/blob/main/.github/workflows/nightly.yml
 name: Nightly

.github/workflows/pr_style_bot.yml

@@ -1,125 +0,0 @@
-# Adapted from https://github.com/huggingface/diffusers/blob/main/.github/workflows/pr_style_bot.yml
-name: PR Style Bot
-
-on:
-  issue_comment:
-    types: [created]
-
-permissions:
-  contents: write
-  pull-requests: write
-
-jobs:
-  run-style-bot:
-    if: >
-      contains(github.event.comment.body, '@bot /style') &&
-      github.event.issue.pull_request != null
-    runs-on: ubuntu-latest
-
-    steps:
-      - name: Extract PR details
-        id: pr_info
-        uses: actions/github-script@v6
-        with:
-          script: |
-            const prNumber = context.payload.issue.number;
-            const { data: pr } = await github.rest.pulls.get({
-              owner: context.repo.owner,
-              repo: context.repo.repo,
-              pull_number: prNumber
-            });
-
-            // We capture both the branch ref and the "full_name" of the head repo
-            // so that we can check out the correct repository & branch (including forks).
-            core.setOutput("prNumber", prNumber);
-            core.setOutput("headRef", pr.head.ref);
-            core.setOutput("headRepoFullName", pr.head.repo.full_name);
-
-      - name: Check out PR branch
-        uses: actions/checkout@v4
-        env:
-          HEADREPOFULLNAME: ${{ steps.pr_info.outputs.headRepoFullName }}
-          HEADREF: ${{ steps.pr_info.outputs.headRef }}
-        with:
-          persist-credentials: true
-          # Instead of checking out the base repo, use the contributor's repo name
-          repository: ${{ env.HEADREPOFULLNAME }}
-          ref: ${{ env.HEADREF }}
-          # You may need fetch-depth: 0 for being able to push
-          fetch-depth: 0
-          token: ${{ secrets.GITHUB_TOKEN }}
-
-      - name: Debug
-        env:
-          HEADREPOFULLNAME: ${{ steps.pr_info.outputs.headRepoFullName }}
-          HEADREF: ${{ steps.pr_info.outputs.headRef }}
-          PRNUMBER: ${{ steps.pr_info.outputs.prNumber }}
-        run: |
-          echo "PR number: ${PRNUMBER}"
-          echo "Head Ref: ${HEADREF}"
-          echo "Head Repo Full Name: ${HEADREPOFULLNAME}"
-
-      - name: Set up Python
-        uses: actions/setup-python@v4
-
-      - name: Get Ruff Version from pre-commit-config.yaml
-        id: get-ruff-version
-        run: |
-          RUFF_VERSION=$(awk '/repo: https:\/\/github.com\/astral-sh\/ruff-pre-commit/{flag=1;next}/rev:/{if(flag){print $2;exit}}' .pre-commit-config.yaml)
-          echo "ruff_version=${RUFF_VERSION}" >> $GITHUB_OUTPUT
-
-      - name: Install Ruff
-        env:
-          RUFF_VERSION: ${{ steps.get-ruff-version.outputs.ruff_version }}
-        run: python -m pip install "ruff==${RUFF_VERSION}"
-
-      - name: Ruff check
-        run: ruff check --fix
-
-      - name: Ruff format
-        run: ruff format
-
-      - name: Commit and push changes
-        id: commit_and_push
-        env:
-          HEADREPOFULLNAME: ${{ steps.pr_info.outputs.headRepoFullName }}
-          HEADREF: ${{ steps.pr_info.outputs.headRef }}
-          PRNUMBER: ${{ steps.pr_info.outputs.prNumber }}
-          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-        run: |
-          echo "HEADREPOFULLNAME: ${HEADREPOFULLNAME}, HEADREF: ${HEADREF}"
-          # Configure git with the Actions bot user
-          git config user.name "github-actions[bot]"
-          git config user.email "github-actions[bot]@users.noreply.github.com"
-
-          # Make sure your 'origin' remote is set to the contributor's fork
-          git remote set-url origin "https://x-access-token:${GITHUB_TOKEN}@github.com/${HEADREPOFULLNAME}.git"
-
-          # If there are changes after running style/quality, commit them
-          if [ -n "$(git status --porcelain)" ]; then
-            git add .
-            git commit -m "Apply style fixes"
-            # Push to the original contributor's forked branch
-            git push origin HEAD:${HEADREF}
-            echo "changes_pushed=true" >> $GITHUB_OUTPUT
-          else
-            echo "No changes to commit."
-            echo "changes_pushed=false" >> $GITHUB_OUTPUT
-          fi
-
-      - name: Comment on PR with workflow run link
-        if: steps.commit_and_push.outputs.changes_pushed == 'true'
-        uses: actions/github-script@v6
-        with:
-          script: |
-            const prNumber = parseInt(process.env.prNumber, 10);
-            const runUrl = `${process.env.GITHUB_SERVER_URL}/${process.env.GITHUB_REPOSITORY}/actions/runs/${process.env.GITHUB_RUN_ID}`
-
-            await github.rest.issues.createComment({
-              owner: context.repo.owner,
-              repo: context.repo.repo,
-              issue_number: prNumber,
-              body: `Style fixes have been applied. [View the workflow run here](${runUrl}).`
-            });
-        env:
-          prNumber: ${{ steps.pr_info.outputs.prNumber }}

.github/workflows/quality.yml

@@ -1,3 +1,17 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 name: Quality
 
 on:

.github/workflows/test-docker-build.yml

@@ -1,3 +1,17 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 # Inspired by
 # https://github.com/huggingface/peft/blob/main/.github/workflows/test-docker-build.yml
 name: Test Dockerfiles
@@ -27,7 +41,7 @@ jobs:
 
       - name: Get changed files
         id: changed-files
-        uses: tj-actions/changed-files@v44
+        uses: tj-actions/changed-files@3f54ebb830831fc121d3263c1857cfbdc310cdb9 #v42
         with:
           files: docker/**
           json: "true"

.github/workflows/test.yml

@@ -1,3 +1,17 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 name: Tests
 
 on:
@@ -112,7 +126,7 @@ jobs:
       # portaudio19-dev is needed to install pyaudio
         run: |
           sudo apt-get update && \
-          sudo apt-get install -y libegl1-mesa-dev portaudio19-dev
+          sudo apt-get install -y libegl1-mesa-dev ffmpeg portaudio19-dev
 
       - name: Install uv and python
         uses: astral-sh/setup-uv@v5

.github/workflows/trufflehog.yml

@@ -1,3 +1,17 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 on:
   push:
 

.gitignore

@@ -1,3 +1,17 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+.dev
 # Logging
 logs
 tmp
@@ -64,7 +78,7 @@ pip-log.txt
 pip-delete-this-directory.txt
 
 # Unit test / coverage reports
-!tests/data
+!tests/artifacts
 htmlcov/
 .tox/
 .nox/

.pre-commit-config.yaml

@@ -1,7 +1,29 @@
-exclude: ^(tests/data)
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+exclude: "tests/artifacts/.*\\.safetensors$"
 default_language_version:
     python: python3.10
 repos:
+  ##### Meta #####
+  - repo: meta
+    hooks:
+      - id: check-useless-excludes
+      - id: check-hooks-apply
+
+
+  ##### Style / Misc. #####
   - repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v5.0.0
     hooks:
@@ -13,26 +35,40 @@ repos:
       - id: check-toml
       - id: end-of-file-fixer
       - id: trailing-whitespace
-  - repo: https://github.com/crate-ci/typos
-    rev: v1.29.10
+
+  - repo: https://github.com/adhtruong/mirrors-typos
+    rev: v1.31.1
     hooks:
       - id: typos
         args: [--force-exclude]
+
   - repo: https://github.com/asottile/pyupgrade
     rev: v3.19.1
     hooks:
     -   id: pyupgrade
+
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.9.6
+    rev: v0.11.4
     hooks:
       - id: ruff
         args: [--fix]
       - id: ruff-format
+
+
+  ##### Security #####
   - repo: https://github.com/gitleaks/gitleaks
-    rev: v8.23.3
+    rev: v8.24.2
     hooks:
       - id: gitleaks
+
   - repo: https://github.com/woodruffw/zizmor-pre-commit
-    rev: v1.3.1
+    rev: v1.5.2
     hooks:
       - id: zizmor
+
+  - repo: https://github.com/PyCQA/bandit
+    rev: 1.8.3
+    hooks:
+    - id: bandit
+      args: ["-c", "pyproject.toml"]
+      additional_dependencies: ["bandit[toml]"]

CONTRIBUTING.md

@@ -291,7 +291,7 @@ sudo apt-get install git-lfs
 git lfs install
 ```
 
-Pull artifacts if they're not in [tests/data](tests/data)
+Pull artifacts if they're not in [tests/artifacts](tests/artifacts)
 ```bash
 git lfs pull
 ```

Makefile

@@ -1,3 +1,17 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 .PHONY: tests
 
 PYTHON_PATH := $(shell which python)
@@ -33,6 +47,7 @@ test-act-ete-train:
 		--policy.dim_model=64 \
 		--policy.n_action_steps=20 \
 		--policy.chunk_size=20 \
+		--policy.device=$(DEVICE) \
 		--env.type=aloha \
 		--env.episode_length=5 \
 		--dataset.repo_id=lerobot/aloha_sim_transfer_cube_human \
@@ -47,7 +62,6 @@ test-act-ete-train:
 		--save_checkpoint=true \
 		--log_freq=1 \
 		--wandb.enable=false \
-		--device=$(DEVICE) \
 		--output_dir=tests/outputs/act/
 
 test-act-ete-train-resume:
@@ -58,11 +72,11 @@ test-act-ete-train-resume:
 test-act-ete-eval:
 	python lerobot/scripts/eval.py \
 		--policy.path=tests/outputs/act/checkpoints/000004/pretrained_model \
+		--policy.device=$(DEVICE) \
 		--env.type=aloha \
 		--env.episode_length=5 \
 		--eval.n_episodes=1 \
-		--eval.batch_size=1 \
-		--device=$(DEVICE)
+		--eval.batch_size=1
 
 test-diffusion-ete-train:
 	python lerobot/scripts/train.py \
@@ -70,6 +84,7 @@ test-diffusion-ete-train:
 		--policy.down_dims='[64,128,256]' \
 		--policy.diffusion_step_embed_dim=32 \
 		--policy.num_inference_steps=10 \
+		--policy.device=$(DEVICE) \
 		--env.type=pusht \
 		--env.episode_length=5 \
 		--dataset.repo_id=lerobot/pusht \
@@ -84,21 +99,21 @@ test-diffusion-ete-train:
 		--save_freq=2 \
 		--log_freq=1 \
 		--wandb.enable=false \
-		--device=$(DEVICE) \
 		--output_dir=tests/outputs/diffusion/
 
 test-diffusion-ete-eval:
 	python lerobot/scripts/eval.py \
 		--policy.path=tests/outputs/diffusion/checkpoints/000002/pretrained_model \
+		--policy.device=$(DEVICE) \
 		--env.type=pusht \
 		--env.episode_length=5 \
 		--eval.n_episodes=1 \
-		--eval.batch_size=1 \
-		--device=$(DEVICE)
+		--eval.batch_size=1
 
 test-tdmpc-ete-train:
 	python lerobot/scripts/train.py \
 		--policy.type=tdmpc \
+		--policy.device=$(DEVICE) \
 		--env.type=xarm \
 		--env.task=XarmLift-v0 \
 		--env.episode_length=5 \
@@ -114,15 +129,14 @@ test-tdmpc-ete-train:
 		--save_freq=2 \
 		--log_freq=1 \
 		--wandb.enable=false \
-		--device=$(DEVICE) \
 		--output_dir=tests/outputs/tdmpc/
 
 test-tdmpc-ete-eval:
 	python lerobot/scripts/eval.py \
 		--policy.path=tests/outputs/tdmpc/checkpoints/000002/pretrained_model \
+		--policy.device=$(DEVICE) \
 		--env.type=xarm \
 		--env.episode_length=5 \
 		--env.task=XarmLift-v0 \
 		--eval.n_episodes=1 \
-		--eval.batch_size=1 \
-		--device=$(DEVICE)
+		--eval.batch_size=1

README.md

@@ -23,15 +23,24 @@
 </div>
 
 <h2 align="center">
-    <p><a href="https://github.com/huggingface/lerobot/blob/main/examples/10_use_so100.md">New robot in town: SO-100</a></p>
+    <p><a href="https://github.com/huggingface/lerobot/blob/main/examples/10_use_so100.md">
+        Build Your Own SO-100 Robot!</a></p>
 </h2>
 
 <div align="center">
-    <img src="media/so100/leader_follower.webp?raw=true" alt="SO-100 leader and follower arms" title="SO-100 leader and follower arms" width="50%">
-    <p>We just added a new tutorial on how to build a more affordable robot, at the price of $110 per arm!</p>
-    <p>Teach it new skills by showing it a few moves with just a laptop.</p>
-    <p>Then watch your homemade robot act autonomously 🤯</p>
-    <p>Follow the link to the <a href="https://github.com/huggingface/lerobot/blob/main/examples/10_use_so100.md">full tutorial for SO-100</a>.</p>
+  <img src="media/so100/leader_follower.webp?raw=true" alt="SO-100 leader and follower arms" title="SO-100 leader and follower arms" width="50%">
+
+  <p><strong>Meet the SO-100 – Just $110 per arm!</strong></p>
+  <p>Train it in minutes with a few simple moves on your laptop.</p>
+  <p>Then sit back and watch your creation act autonomously! 🤯</p>
+
+  <p><a href="https://github.com/huggingface/lerobot/blob/main/examples/10_use_so100.md">
+      Get the full SO-100 tutorial here.</a></p>
+
+  <p>Want to take it to the next level? Make your SO-100 mobile by building LeKiwi!</p>
+  <p>Check out the <a href="https://github.com/huggingface/lerobot/blob/main/examples/11_use_lekiwi.md">LeKiwi tutorial</a> and bring your robot to life on wheels.</p>
+
+  <img src="media/lekiwi/kiwi.webp?raw=true" alt="LeKiwi mobile robot" title="LeKiwi mobile robot" width="50%">
 </div>
 
 <br/>
@@ -89,14 +98,18 @@ conda create -y -n lerobot python=3.10
 conda activate lerobot
 ```
 
+When using `miniconda`, install `ffmpeg` in your environment:
+```bash
+conda install ffmpeg -c conda-forge
+```
+
 Install 🤗 LeRobot:
 ```bash
 pip install -e .
 ```
 
-> **NOTE:** Depending on your platform, If you encounter any build errors during this step
-you may need to install `cmake` and `build-essential` for building some of our dependencies.
-On linux: `sudo apt-get install cmake build-essential`
+> **NOTE:** If you encounter build errors, you may need to install additional dependencies (`cmake`, `build-essential`, and `ffmpeg libs`). On Linux, run:
+`sudo apt-get install cmake build-essential python-dev pkg-config libavformat-dev libavcodec-dev libavdevice-dev libavutil-dev libswscale-dev libswresample-dev libavfilter-dev pkg-config`. For other systems, see: [Compiling PyAV](https://pyav.org/docs/develop/overview/installation.html#bring-your-own-ffmpeg)
 
 For simulations, 🤗 LeRobot comes with gymnasium environments that can be installed as extras:
 - [aloha](https://github.com/huggingface/gym-aloha)
@@ -223,8 +236,8 @@ python lerobot/scripts/eval.py \
     --env.type=pusht \
     --eval.batch_size=10 \
     --eval.n_episodes=10 \
-    --use_amp=false \
-    --device=cuda
+    --policy.use_amp=false \
+    --policy.device=cuda
 ```
 
 Note: After training your own policy, you can re-evaluate the checkpoints with:
@@ -375,3 +388,6 @@ Additionally, if you are using any of the particular policy architecture, pretra
   year={2024}
 }
 ```
+## Star History
+
+[![Star History Chart](https://api.star-history.com/svg?repos=huggingface/lerobot&type=Timeline)](https://star-history.com/#huggingface/lerobot&Timeline)

benchmarks/video/README.md

@@ -51,7 +51,7 @@ For a comprehensive list and documentation of these parameters, see the ffmpeg d
 ### Decoding parameters
 **Decoder**
 We tested two video decoding backends from torchvision:
-- `pyav` (default)
+- `pyav`
 - `video_reader` (requires to build torchvision from source)
 
 **Requested timestamps**

benchmarks/video/capture_camera_feed.py

@@ -17,12 +17,21 @@
 
 import argparse
 import datetime as dt
+import os
+import time
 from pathlib import Path
 
 import cv2
+import rerun as rr
+
+# see https://rerun.io/docs/howto/visualization/limit-ram
+RERUN_MEMORY_LIMIT = os.getenv("LEROBOT_RERUN_MEMORY_LIMIT", "5%")
 
 
-def display_and_save_video_stream(output_dir: Path, fps: int, width: int, height: int):
+def display_and_save_video_stream(output_dir: Path, fps: int, width: int, height: int, duration: int):
+    rr.init("lerobot_capture_camera_feed")
+    rr.spawn(memory_limit=RERUN_MEMORY_LIMIT)
+
     now = dt.datetime.now()
     capture_dir = output_dir / f"{now:%Y-%m-%d}" / f"{now:%H-%M-%S}"
     if not capture_dir.exists():
@@ -39,24 +48,21 @@ def display_and_save_video_stream(output_dir: Path, fps: int, width: int, height
     cap.set(cv2.CAP_PROP_FRAME_HEIGHT, height)
 
     frame_index = 0
-    while True:
+    start_time = time.time()
+    while time.time() - start_time < duration:
         ret, frame = cap.read()
 
         if not ret:
             print("Error: Could not read frame.")
             break
-
-        cv2.imshow("Video Stream", frame)
+        rr.log("video/stream", rr.Image(frame.numpy()), static=True)
         cv2.imwrite(str(capture_dir / f"frame_{frame_index:06d}.png"), frame)
         frame_index += 1
 
-        # Break the loop on 'q' key press
-        if cv2.waitKey(1) & 0xFF == ord("q"):
-            break
-
-    # Release the capture and destroy all windows
+    # Release the capture
     cap.release()
-    cv2.destroyAllWindows()
+
+    # TODO(Steven): Add a graceful shutdown via a close() method for the Viewer context, though not currently supported in the Rerun API.
 
 
 if __name__ == "__main__":
@@ -86,5 +92,11 @@ if __name__ == "__main__":
         default=720,
         help="Height of the captured images.",
     )
+    parser.add_argument(
+        "--duration",
+        type=int,
+        default=20,
+        help="Duration in seconds for which the video stream should be captured.",
+    )
     args = parser.parse_args()
     display_and_save_video_stream(**vars(args))

benchmarks/video/run_video_benchmark.py

@@ -67,7 +67,7 @@ def parse_int_or_none(value) -> int | None:
 def check_datasets_formats(repo_ids: list) -> None:
     for repo_id in repo_ids:
         dataset = LeRobotDataset(repo_id)
-        if dataset.video:
+        if len(dataset.meta.video_keys) > 0:
             raise ValueError(
                 f"Use only image dataset for running this benchmark. Video dataset provided: {repo_id}"
             )

examples/1_load_lerobot_dataset.py

@@ -1,3 +1,17 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 """
 This script demonstrates the use of `LeRobotDataset` class for handling and processing robotic datasets from Hugging Face.
 It illustrates how to load datasets, manipulate them, and apply transformations suitable for machine learning tasks in PyTorch.
@@ -105,7 +119,7 @@ print(dataset.features[camera_key]["shape"])
 delta_timestamps = {
     # loads 4 images: 1 second before current frame, 500 ms before, 200 ms before, and current frame
     camera_key: [-1, -0.5, -0.20, 0],
-    # loads 8 state vectors: 1.5 seconds before, 1 second before, ... 200 ms, 100 ms, and current frame
+    # loads 6 state vectors: 1.5 seconds before, 1 second before, ... 200 ms, 100 ms, and current frame
     "observation.state": [-1.5, -1, -0.5, -0.20, -0.10, 0],
     # loads 64 action vectors: current frame, 1 frame in the future, 2 frames, ... 63 frames in the future
     "action": [t / dataset.fps for t in range(64)],
@@ -129,6 +143,6 @@ dataloader = torch.utils.data.DataLoader(
 
 for batch in dataloader:
     print(f"{batch[camera_key].shape=}")  # (32, 4, c, h, w)
-    print(f"{batch['observation.state'].shape=}")  # (32, 5, c)
+    print(f"{batch['observation.state'].shape=}")  # (32, 6, c)
     print(f"{batch['action'].shape=}")  # (32, 64, c)
     break

examples/2_evaluate_pretrained_policy.py

@@ -1,10 +1,24 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 """
 This scripts demonstrates how to evaluate a pretrained policy from the HuggingFace Hub or from your local
 training outputs directory. In the latter case, you might want to run examples/3_train_policy.py first.
 
 It requires the installation of the 'gym_pusht' simulation environment. Install it by running:
 ```bash
-pip install -e ".[pusht]"`
+pip install -e ".[pusht]"
 ```
 """
 
@@ -30,7 +44,7 @@ pretrained_policy_path = "lerobot/diffusion_pusht"
 # OR a path to a local outputs/train folder.
 # pretrained_policy_path = Path("outputs/train/example_pusht_diffusion")
 
-policy = DiffusionPolicy.from_pretrained(pretrained_policy_path, map_location=device)
+policy = DiffusionPolicy.from_pretrained(pretrained_policy_path)
 
 # Initialize evaluation environment to render two observation types:
 # an image of the scene and state/position of the agent. The environment

examples/3_train_policy.py

@@ -1,3 +1,17 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 """This scripts demonstrates how to train Diffusion Policy on the PushT environment.
 
 Once you have trained a model with this script, you can try to evaluate it on
@@ -85,7 +99,7 @@ def main():
     done = False
     while not done:
         for batch in dataloader:
-            batch = {k: v.to(device, non_blocking=True) for k, v in batch.items()}
+            batch = {k: (v.to(device) if isinstance(v, torch.Tensor) else v) for k, v in batch.items()}
             loss, _ = policy.forward(batch)
             loss.backward()
             optimizer.step()

examples/4_train_policy_with_script.md

@@ -1,5 +1,5 @@
 This tutorial will explain the training script, how to use it, and particularly how to configure everything needed for the training run.
-> **Note:** The following assume you're running these commands on a machine equipped with a cuda GPU. If you don't have one (or if you're using a Mac), you can add `--device=cpu` (`--device=mps` respectively). However, be advised that the code executes much slower on cpu.
+> **Note:** The following assume you're running these commands on a machine equipped with a cuda GPU. If you don't have one (or if you're using a Mac), you can add `--policy.device=cpu` (`--policy.device=mps` respectively). However, be advised that the code executes much slower on cpu.
 
 
 ## The training script

examples/7_get_started_with_real_robot.md

@@ -46,13 +46,6 @@ Using `uv`:
 uv sync --extra "dynamixel"
 ```
 
-/!\ For Linux only, ffmpeg and opencv requires conda install for now. Run this exact sequence of commands:
-```bash
-conda install -c conda-forge ffmpeg
-pip uninstall opencv-python
-conda install -c conda-forge "opencv>=4.10.0"
-```
-
 You are now ready to plug the 5V power supply to the motor bus of the leader arm (the smaller one) since all its motors only require 5V.
 
 Then plug the 12V power supply to the motor bus of the follower arm. It has two motors that need 12V, and the rest will be powered with 5V through the voltage convertor.
@@ -62,6 +55,9 @@ Finally, connect both arms to your computer via USB. Note that the USB doesn't p
 Now you are ready to configure your motors for the first time, as detailed in the sections below. In the upcoming sections, you'll learn about our classes and functions by running some python code in an interactive session, or by copy-pasting it in a python file.
 
 If you have already configured your motors the first time, you can streamline the process by directly running the teleoperate script (which is detailed further in the tutorial):
+
+> **NOTE:** To visualize the data, enable `--control.display_data=true`. This streams the data using `rerun`.
+
 ```bash
 python lerobot/scripts/control_robot.py \
   --robot.type=koch \
@@ -87,7 +83,7 @@ python lerobot/scripts/configure_motor.py \
   --brand dynamixel \
   --model xl330-m288 \
   --baudrate 1000000 \
-  --ID 1
+  --id 1
 ```
 
 Then unplug your first motor and plug the second motor and set its ID to 2.
@@ -97,7 +93,7 @@ python lerobot/scripts/configure_motor.py \
   --brand dynamixel \
   --model xl330-m288 \
   --baudrate 1000000 \
-  --ID 2
+  --id 2
 ```
 
 Redo the process for all your motors until ID 6.
@@ -292,6 +288,11 @@ Steps:
    - Scan for devices. All 12 motors should appear.
    - Select the motors one by one and move the arm. Check that the graphical indicator near the top right shows the movement.
 
+** There is a common issue with the Dynamixel XL430-W250 motors where the motors become undiscoverable after upgrading their firmware from Mac and Windows Dynamixel Wizard2 applications.  When this occurs, it is required to do a firmware recovery (Select `DYNAMIXEL Firmware Recovery` and follow the prompts).   There are two known workarounds to conduct this firmware reset:
+  1) Install the Dynamixel Wizard on a linux machine and complete the firmware recovery
+  2) Use the Dynamixel U2D2 in order to perform the reset with Windows or Mac.  This U2D2 can be purchased [here](https://www.robotis.us/u2d2/).
+  For either solution, open DYNAMIXEL Wizard 2.0 and select the appropriate port. You will likely be unable to see the motor in the GUI at this time. Select `Firmware Recovery`, carefully choose the correct model, and wait for the process to complete. Finally, re-scan to confirm the firmware recovery was successful.
+
 **Read and Write with DynamixelMotorsBus**
 
 To get familiar with how `DynamixelMotorsBus` communicates with the motors, you can start by reading data from them. Copy past this code in the same interactive python session:
@@ -386,14 +387,14 @@ When you connect your robot for the first time, the [`ManipulatorRobot`](../lero
 
 Here are the positions you'll move the follower arm to:
 
-| 1. Zero position | 2. Rotated position | 3. Rest position |
-|---|---|---|
+| 1. Zero position                                                                                                                                                  | 2. Rotated position                                                                                                                                                        | 3. Rest position                                                                                                                                                  |
+| ----------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------- |
 | <img src="../media/koch/follower_zero.webp?raw=true" alt="Koch v1.1 follower arm zero position" title="Koch v1.1 follower arm zero position" style="width:100%;"> | <img src="../media/koch/follower_rotated.webp?raw=true" alt="Koch v1.1 follower arm rotated position" title="Koch v1.1 follower arm rotated position" style="width:100%;"> | <img src="../media/koch/follower_rest.webp?raw=true" alt="Koch v1.1 follower arm rest position" title="Koch v1.1 follower arm rest position" style="width:100%;"> |
 
 And here are the corresponding positions for the leader arm:
 
-| 1. Zero position | 2. Rotated position | 3. Rest position |
-|---|---|---|
+| 1. Zero position                                                                                                                                            | 2. Rotated position                                                                                                                                                  | 3. Rest position                                                                                                                                            |
+| ----------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------- |
 | <img src="../media/koch/leader_zero.webp?raw=true" alt="Koch v1.1 leader arm zero position" title="Koch v1.1 leader arm zero position" style="width:100%;"> | <img src="../media/koch/leader_rotated.webp?raw=true" alt="Koch v1.1 leader arm rotated position" title="Koch v1.1 leader arm rotated position" style="width:100%;"> | <img src="../media/koch/leader_rest.webp?raw=true" alt="Koch v1.1 leader arm rest position" title="Koch v1.1 leader arm rest position" style="width:100%;"> |
 
 You can watch a [video tutorial of the calibration procedure](https://youtu.be/8drnU9uRY24) for more details.
@@ -829,16 +830,11 @@ It contains:
 - `dtRphone:33.84 (29.5hz)` which is the delta time of capturing an image from the phone camera in the thread running asynchronously.
 
 Troubleshooting:
-- On Linux, if you encounter a hanging issue when using cameras, uninstall opencv and re-install it with conda:
-```bash
-pip uninstall opencv-python
-conda install -c conda-forge opencv=4.10.0
-```
 - On Linux, if you encounter any issue during video encoding with `ffmpeg: unknown encoder libsvtav1`, you can:
-  - install with conda-forge by running `conda install -c conda-forge ffmpeg` (it should be compiled with `libsvtav1`),
-  - or, install [Homebrew](https://brew.sh) and run `brew install ffmpeg` (it should be compiled with `libsvtav1`),
-  - or, install [ffmpeg build dependencies](https://trac.ffmpeg.org/wiki/CompilationGuide/Ubuntu#GettheDependencies) and [compile ffmpeg from source with libsvtav1](https://trac.ffmpeg.org/wiki/CompilationGuide/Ubuntu#libsvtav1),
-  - and, make sure you use the corresponding ffmpeg binary to your install with `which ffmpeg`.
+   - install with conda-forge by running `conda install -c conda-forge ffmpeg` (it should be compiled with `libsvtav1`),
+> **NOTE:** This usually installs `ffmpeg 7.X` for your platform (check the version installed with `ffmpeg -encoders | grep libsvtav1`). If it isn't `ffmpeg 7.X` or lacks `libsvtav1` support, you can explicitly install `ffmpeg 7.X` using: `conda install ffmpeg=7.1.1 -c conda-forge`
+   - or, install [ffmpeg build dependencies](https://trac.ffmpeg.org/wiki/CompilationGuide/Ubuntu#GettheDependencies) and [compile ffmpeg from source with libsvtav1](https://trac.ffmpeg.org/wiki/CompilationGuide/Ubuntu#libsvtav1),
+   - and, make sure you use the corresponding ffmpeg binary to your install with `which ffmpeg`.
 - On Linux, if the left and right arrow keys and escape key don't have any effect during data recording, make sure you've set the `$DISPLAY` environment variable. See [pynput limitations](https://pynput.readthedocs.io/en/latest/limitations.html#linux).
 
 At the end of data recording, your dataset will be uploaded on your Hugging Face page (e.g. https://huggingface.co/datasets/cadene/koch_test) that you can obtain by running:
@@ -898,14 +894,14 @@ python lerobot/scripts/train.py \
   --policy.type=act \
   --output_dir=outputs/train/act_koch_test \
   --job_name=act_koch_test \
-  --device=cuda \
+  --policy.device=cuda \
   --wandb.enable=true
 ```
 
 Let's explain it:
 1. We provided the dataset as argument with `--dataset.repo_id=${HF_USER}/koch_test`.
 2. We provided the policy with `policy.type=act`. This loads configurations from [`configuration_act.py`](../lerobot/common/policies/act/configuration_act.py). Importantly, this policy will automatically adapt to the number of motor sates, motor actions and cameras of your robot (e.g. `laptop` and `phone`) which have been saved in your dataset.
-4. We provided `device=cuda` since we are training on a Nvidia GPU, but you could use `device=mps` to train on Apple silicon.
+4. We provided `policy.device=cuda` since we are training on a Nvidia GPU, but you could use `policy.device=mps` to train on Apple silicon.
 5. We provided `wandb.enable=true` to use [Weights and Biases](https://docs.wandb.ai/quickstart) for visualizing training plots. This is optional but if you use it, make sure you are logged in by running `wandb login`.
 
 For more information on the `train` script see the previous tutorial: [`examples/4_train_policy_with_script.md`](../examples/4_train_policy_with_script.md)

examples/advanced/1_add_image_transforms.py

@@ -1,3 +1,17 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 """
 This script demonstrates how to use torchvision's image transformation with LeRobotDataset for data
 augmentation purposes. The transformations are passed to the dataset as an argument upon creation, and

examples/advanced/2_calculate_validation_loss.py

@@ -1,3 +1,17 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 """This script demonstrates how to slice a dataset and calculate the loss on a subset of the data.
 
 This technique can be useful for debugging and testing purposes, as well as identifying whether a policy

examples/port_datasets/pusht_zarr.py

@@ -1,229 +0,0 @@
-import shutil
-from pathlib import Path
-
-import numpy as np
-from huggingface_hub import HfApi
-
-from lerobot.common.constants import HF_LEROBOT_HOME
-from lerobot.common.datasets.lerobot_dataset import CODEBASE_VERSION, LeRobotDataset
-from lerobot.common.datasets.push_dataset_to_hub._download_raw import download_raw
-
-PUSHT_TASK = "Push the T-shaped blue block onto the T-shaped green target surface."
-PUSHT_FEATURES = {
-    "observation.state": {
-        "dtype": "float32",
-        "shape": (2,),
-        "names": {
-            "axes": ["x", "y"],
-        },
-    },
-    "action": {
-        "dtype": "float32",
-        "shape": (2,),
-        "names": {
-            "axes": ["x", "y"],
-        },
-    },
-    "next.reward": {
-        "dtype": "float32",
-        "shape": (1,),
-        "names": None,
-    },
-    "next.success": {
-        "dtype": "bool",
-        "shape": (1,),
-        "names": None,
-    },
-    "observation.environment_state": {
-        "dtype": "float32",
-        "shape": (16,),
-        "names": [
-            "keypoints",
-        ],
-    },
-    "observation.image": {
-        "dtype": None,
-        "shape": (3, 96, 96),
-        "names": [
-            "channels",
-            "height",
-            "width",
-        ],
-    },
-}
-
-
-def build_features(mode: str) -> dict:
-    features = PUSHT_FEATURES
-    if mode == "keypoints":
-        features.pop("observation.image")
-    else:
-        features.pop("observation.environment_state")
-        features["observation.image"]["dtype"] = mode
-
-    return features
-
-
-def load_raw_dataset(zarr_path: Path):
-    try:
-        from lerobot.common.datasets.push_dataset_to_hub._diffusion_policy_replay_buffer import (
-            ReplayBuffer as DiffusionPolicyReplayBuffer,
-        )
-    except ModuleNotFoundError as e:
-        print("`gym_pusht` is not installed. Please install it with `pip install 'lerobot[gym_pusht]'`")
-        raise e
-
-    zarr_data = DiffusionPolicyReplayBuffer.copy_from_path(zarr_path)
-    return zarr_data
-
-
-def calculate_coverage(zarr_data):
-    try:
-        import pymunk
-        from gym_pusht.envs.pusht import PushTEnv, pymunk_to_shapely
-    except ModuleNotFoundError as e:
-        print("`gym_pusht` is not installed. Please install it with `pip install 'lerobot[gym_pusht]'`")
-        raise e
-
-    block_pos = zarr_data["state"][:, 2:4]
-    block_angle = zarr_data["state"][:, 4]
-
-    num_frames = len(block_pos)
-
-    coverage = np.zeros((num_frames,), dtype=np.float32)
-    # 8 keypoints with 2 coords each
-    keypoints = np.zeros((num_frames, 16), dtype=np.float32)
-
-    # Set x, y, theta (in radians)
-    goal_pos_angle = np.array([256, 256, np.pi / 4])
-    goal_body = PushTEnv.get_goal_pose_body(goal_pos_angle)
-
-    for i in range(num_frames):
-        space = pymunk.Space()
-        space.gravity = 0, 0
-        space.damping = 0
-
-        # Add walls.
-        walls = [
-            PushTEnv.add_segment(space, (5, 506), (5, 5), 2),
-            PushTEnv.add_segment(space, (5, 5), (506, 5), 2),
-            PushTEnv.add_segment(space, (506, 5), (506, 506), 2),
-            PushTEnv.add_segment(space, (5, 506), (506, 506), 2),
-        ]
-        space.add(*walls)
-
-        block_body, block_shapes = PushTEnv.add_tee(space, block_pos[i].tolist(), block_angle[i].item())
-        goal_geom = pymunk_to_shapely(goal_body, block_body.shapes)
-        block_geom = pymunk_to_shapely(block_body, block_body.shapes)
-        intersection_area = goal_geom.intersection(block_geom).area
-        goal_area = goal_geom.area
-        coverage[i] = intersection_area / goal_area
-        keypoints[i] = PushTEnv.get_keypoints(block_shapes).flatten()
-
-    return coverage, keypoints
-
-
-def calculate_success(coverage: float, success_threshold: float):
-    return coverage > success_threshold
-
-
-def calculate_reward(coverage: float, success_threshold: float):
-    return np.clip(coverage / success_threshold, 0, 1)
-
-
-def main(raw_dir: Path, repo_id: str, mode: str = "video", push_to_hub: bool = True):
-    if mode not in ["video", "image", "keypoints"]:
-        raise ValueError(mode)
-
-    if (HF_LEROBOT_HOME / repo_id).exists():
-        shutil.rmtree(HF_LEROBOT_HOME / repo_id)
-
-    if not raw_dir.exists():
-        download_raw(raw_dir, repo_id="lerobot-raw/pusht_raw")
-
-    zarr_data = load_raw_dataset(zarr_path=raw_dir / "pusht_cchi_v7_replay.zarr")
-
-    env_state = zarr_data["state"][:]
-    agent_pos = env_state[:, :2]
-
-    action = zarr_data["action"][:]
-    image = zarr_data["img"]  # (b, h, w, c)
-
-    if image.dtype == np.float32 and image.max() == np.float32(255):
-        # HACK: images are loaded as float32 but they actually encode uint8 data
-        image = image.astype(np.uint8)
-
-    episode_data_index = {
-        "from": np.concatenate(([0], zarr_data.meta["episode_ends"][:-1])),
-        "to": zarr_data.meta["episode_ends"],
-    }
-
-    # Calculate success and reward based on the overlapping area
-    # of the T-object and the T-area.
-    coverage, keypoints = calculate_coverage(zarr_data)
-    success = calculate_success(coverage, success_threshold=0.95)
-    reward = calculate_reward(coverage, success_threshold=0.95)
-
-    features = build_features(mode)
-    dataset = LeRobotDataset.create(
-        repo_id=repo_id,
-        fps=10,
-        robot_type="2d pointer",
-        features=features,
-        image_writer_threads=4,
-    )
-    episodes = range(len(episode_data_index["from"]))
-    for ep_idx in episodes:
-        from_idx = episode_data_index["from"][ep_idx]
-        to_idx = episode_data_index["to"][ep_idx]
-        num_frames = to_idx - from_idx
-
-        for frame_idx in range(num_frames):
-            i = from_idx + frame_idx
-            idx = i + (frame_idx < num_frames - 1)
-            frame = {
-                "action": action[i],
-                # Shift reward and success by +1 until the last item of the episode
-                "next.reward": reward[idx : idx + 1],
-                "next.success": success[idx : idx + 1],
-                "task": PUSHT_TASK,
-            }
-
-            frame["observation.state"] = agent_pos[i]
-
-            if mode == "keypoints":
-                frame["observation.environment_state"] = keypoints[i]
-            else:
-                frame["observation.image"] = image[i]
-
-            dataset.add_frame(frame)
-
-        dataset.save_episode()
-
-    if push_to_hub:
-        dataset.push_to_hub()
-        hub_api = HfApi()
-        hub_api.create_tag(repo_id, tag=CODEBASE_VERSION, repo_type="dataset")
-
-
-if __name__ == "__main__":
-    # To try this script, modify the repo id with your own HuggingFace user (e.g cadene/pusht)
-    repo_id = "lerobot/pusht"
-
-    modes = ["video", "image", "keypoints"]
-    # Uncomment if you want to try with a specific mode
-    # modes = ["video"]
-    # modes = ["image"]
-    # modes = ["keypoints"]
-
-    raw_dir = Path("data/lerobot-raw/pusht_raw")
-    for mode in modes:
-        if mode in ["image", "keypoints"]:
-            repo_id += f"_{mode}"
-
-        # download and load raw dataset, create LeRobotDataset, populate it, push to hub
-        main(raw_dir, repo_id=repo_id, mode=mode)
-
-        # Uncomment if you want to load the local dataset and explore it
-        # dataset = LeRobotDataset(repo_id=repo_id)
-        # breakpoint()

lerobot/common/cameras/__init__.py

@@ -0,0 +1,4 @@
+from .camera import Camera
+from .configs import CameraConfig
+
+__all__ = ["Camera", "CameraConfig"]

lerobot/common/cameras/camera.py

@@ -0,0 +1,25 @@
+import abc
+
+import numpy as np
+
+
+class Camera(abc.ABC):
+    @abc.abstractmethod
+    def connect(self):
+        pass
+
+    @abc.abstractmethod
+    def read(self, temporary_color_mode: str | None = None) -> np.ndarray:
+        pass
+
+    @abc.abstractmethod
+    def async_read(self) -> np.ndarray:
+        pass
+
+    @abc.abstractmethod
+    def disconnect(self):
+        pass
+
+    def __del__(self):
+        if getattr(self, "is_connected", False):
+            self.disconnect()

lerobot/common/cameras/configs.py

@@ -0,0 +1,11 @@
+import abc
+from dataclasses import dataclass
+
+import draccus
+
+
+@dataclass
+class CameraConfig(draccus.ChoiceRegistry, abc.ABC):
+    @property
+    def type(self) -> str:
+        return self.get_choice_name(self.__class__)

lerobot/common/cameras/intel/__init__.py

@@ -0,0 +1,4 @@
+from .camera_realsense import RealSenseCamera
+from .configuration_realsense import RealSenseCameraConfig
+
+__all__ = ["RealSenseCamera", "RealSenseCameraConfig"]

lerobot/common/cameras/intel/camera_realsense.py

@@ -1,3 +1,17 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 """
 This file contains utilities for recording frames from Intel Realsense cameras.
 """
@@ -17,14 +31,15 @@ from threading import Thread
 import numpy as np
 from PIL import Image
 
-from lerobot.common.robot_devices.cameras.configs import IntelRealSenseCameraConfig
-from lerobot.common.robot_devices.utils import (
-    RobotDeviceAlreadyConnectedError,
-    RobotDeviceNotConnectedError,
+from lerobot.common.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
+from lerobot.common.utils.robot_utils import (
     busy_wait,
 )
 from lerobot.common.utils.utils import capture_timestamp_utc
 
+from ..camera import Camera
+from .configuration_realsense import RealSenseCameraConfig
+
 SERIAL_NUMBER_INDEX = 1
 
 
@@ -34,7 +49,7 @@ def find_cameras(raise_when_empty=True, mock=False) -> list[dict]:
     connected to the computer.
     """
     if mock:
-        import tests.mock_pyrealsense2 as rs
+        import tests.cameras.mock_pyrealsense2 as rs
     else:
         import pyrealsense2 as rs
 
@@ -86,7 +101,7 @@ def save_images_from_cameras(
         serial_numbers = [cam["serial_number"] for cam in camera_infos]
 
     if mock:
-        import tests.mock_cv2 as cv2
+        import tests.cameras.mock_cv2 as cv2
     else:
         import cv2
 
@@ -94,13 +109,11 @@ def save_images_from_cameras(
     cameras = []
     for cam_sn in serial_numbers:
         print(f"{cam_sn=}")
-        config = IntelRealSenseCameraConfig(
-            serial_number=cam_sn, fps=fps, width=width, height=height, mock=mock
-        )
-        camera = IntelRealSenseCamera(config)
+        config = RealSenseCameraConfig(serial_number=cam_sn, fps=fps, width=width, height=height, mock=mock)
+        camera = RealSenseCamera(config)
         camera.connect()
         print(
-            f"IntelRealSenseCamera({camera.serial_number}, fps={camera.fps}, width={camera.width}, height={camera.height}, color_mode={camera.color_mode})"
+            f"RealSenseCamera({camera.serial_number}, fps={camera.fps}, width={camera.capture_width}, height={camera.capture_height}, color_mode={camera.color_mode})"
         )
         cameras.append(camera)
 
@@ -152,11 +165,11 @@ def save_images_from_cameras(
             camera.disconnect()
 
 
-class IntelRealSenseCamera:
+class RealSenseCamera(Camera):
     """
-    The IntelRealSenseCamera class is similar to OpenCVCamera class but adds additional features for Intel Real Sense cameras:
+    The RealSenseCamera class is similar to OpenCVCamera class but adds additional features for Intel Real Sense cameras:
     - is instantiated with the serial number of the camera - won't randomly change as it can be the case of OpenCVCamera for Linux,
-    - can also be instantiated with the camera's name — if it's unique — using IntelRealSenseCamera.init_from_name(),
+    - can also be instantiated with the camera's name — if it's unique — using RealSenseCamera.init_from_name(),
     - depth map can be returned.
 
     To find the camera indices of your cameras, you can run our utility script that will save a few frames for each camera:
@@ -164,15 +177,15 @@ class IntelRealSenseCamera:
     python lerobot/common/robot_devices/cameras/intelrealsense.py --images-dir outputs/images_from_intelrealsense_cameras
     ```
 
-    When an IntelRealSenseCamera is instantiated, if no specific config is provided, the default fps, width, height and color_mode
+    When an RealSenseCamera is instantiated, if no specific config is provided, the default fps, width, height and color_mode
     of the given camera will be used.
 
     Example of instantiating with a serial number:
     ```python
-    from lerobot.common.robot_devices.cameras.configs import IntelRealSenseCameraConfig
+    from lerobot.common.robot_devices.cameras.configs import RealSenseCameraConfig
 
-    config = IntelRealSenseCameraConfig(serial_number=128422271347)
-    camera = IntelRealSenseCamera(config)
+    config = RealSenseCameraConfig(serial_number=128422271347)
+    camera = RealSenseCamera(config)
     camera.connect()
     color_image = camera.read()
     # when done using the camera, consider disconnecting
@@ -181,21 +194,21 @@ class IntelRealSenseCamera:
 
     Example of instantiating with a name if it's unique:
     ```
-    config = IntelRealSenseCameraConfig(name="Intel RealSense D405")
+    config = RealSenseCameraConfig(name="Intel RealSense D405")
     ```
 
     Example of changing default fps, width, height and color_mode:
     ```python
-    config = IntelRealSenseCameraConfig(serial_number=128422271347, fps=30, width=1280, height=720)
-    config = IntelRealSenseCameraConfig(serial_number=128422271347, fps=90, width=640, height=480)
-    config = IntelRealSenseCameraConfig(serial_number=128422271347, fps=90, width=640, height=480, color_mode="bgr")
+    config = RealSenseCameraConfig(serial_number=128422271347, fps=30, width=1280, height=720)
+    config = RealSenseCameraConfig(serial_number=128422271347, fps=90, width=640, height=480)
+    config = RealSenseCameraConfig(serial_number=128422271347, fps=90, width=640, height=480, color_mode="bgr")
     # Note: might error out upon `camera.connect()` if these settings are not compatible with the camera
     ```
 
     Example of returning depth:
     ```python
-    config = IntelRealSenseCameraConfig(serial_number=128422271347, use_depth=True)
-    camera = IntelRealSenseCamera(config)
+    config = RealSenseCameraConfig(serial_number=128422271347, use_depth=True)
+    camera = RealSenseCamera(config)
     camera.connect()
     color_image, depth_map = camera.read()
     ```
@@ -203,16 +216,27 @@ class IntelRealSenseCamera:
 
     def __init__(
         self,
-        config: IntelRealSenseCameraConfig,
+        config: RealSenseCameraConfig,
     ):
         self.config = config
         if config.name is not None:
             self.serial_number = self.find_serial_number_from_name(config.name)
         else:
             self.serial_number = config.serial_number
+
+        # Store the raw (capture) resolution from the config.
+        self.capture_width = config.width
+        self.capture_height = config.height
+
+        # If rotated by ±90, swap width and height.
+        if config.rotation in [-90, 90]:
+            self.width = config.height
+            self.height = config.width
+        else:
+            self.width = config.width
+            self.height = config.height
+
         self.fps = config.fps
-        self.width = config.width
-        self.height = config.height
         self.channels = config.channels
         self.color_mode = config.color_mode
         self.use_depth = config.use_depth
@@ -228,11 +252,10 @@ class IntelRealSenseCamera:
         self.logs = {}
 
         if self.mock:
-            import tests.mock_cv2 as cv2
+            import tests.cameras.mock_cv2 as cv2
         else:
             import cv2
 
-        # TODO(alibets): Do we keep original width/height or do we define them after rotation?
         self.rotation = None
         if config.rotation == -90:
             self.rotation = cv2.ROTATE_90_COUNTERCLOCKWISE
@@ -258,27 +281,29 @@ class IntelRealSenseCamera:
 
     def connect(self):
         if self.is_connected:
-            raise RobotDeviceAlreadyConnectedError(
-                f"IntelRealSenseCamera({self.serial_number}) is already connected."
-            )
+            raise DeviceAlreadyConnectedError(f"RealSenseCamera({self.serial_number}) is already connected.")
 
         if self.mock:
-            import tests.mock_pyrealsense2 as rs
+            import tests.cameras.mock_pyrealsense2 as rs
         else:
             import pyrealsense2 as rs
 
         config = rs.config()
         config.enable_device(str(self.serial_number))
 
-        if self.fps and self.width and self.height:
+        if self.fps and self.capture_width and self.capture_height:
             # TODO(rcadene): can we set rgb8 directly?
-            config.enable_stream(rs.stream.color, self.width, self.height, rs.format.rgb8, self.fps)
+            config.enable_stream(
+                rs.stream.color, self.capture_width, self.capture_height, rs.format.rgb8, self.fps
+            )
         else:
             config.enable_stream(rs.stream.color)
 
         if self.use_depth:
-            if self.fps and self.width and self.height:
-                config.enable_stream(rs.stream.depth, self.width, self.height, rs.format.z16, self.fps)
+            if self.fps and self.capture_width and self.capture_height:
+                config.enable_stream(
+                    rs.stream.depth, self.capture_width, self.capture_height, rs.format.z16, self.fps
+                )
             else:
                 config.enable_stream(rs.stream.depth)
 
@@ -302,7 +327,7 @@ class IntelRealSenseCamera:
                     "To find the serial number you should use, run `python lerobot/common/robot_devices/cameras/intelrealsense.py`."
                 )
 
-            raise OSError(f"Can't access IntelRealSenseCamera({self.serial_number}).")
+            raise OSError(f"Can't access RealSenseCamera({self.serial_number}).")
 
         color_stream = profile.get_stream(rs.stream.color)
         color_profile = color_stream.as_video_stream_profile()
@@ -314,20 +339,20 @@ class IntelRealSenseCamera:
         if self.fps is not None and not math.isclose(self.fps, actual_fps, rel_tol=1e-3):
             # Using `OSError` since it's a broad that encompasses issues related to device communication
             raise OSError(
-                f"Can't set {self.fps=} for IntelRealSenseCamera({self.serial_number}). Actual value is {actual_fps}."
+                f"Can't set {self.fps=} for RealSenseCamera({self.serial_number}). Actual value is {actual_fps}."
             )
-        if self.width is not None and self.width != actual_width:
+        if self.capture_width is not None and self.capture_width != actual_width:
             raise OSError(
-                f"Can't set {self.width=} for IntelRealSenseCamera({self.serial_number}). Actual value is {actual_width}."
+                f"Can't set {self.capture_width=} for RealSenseCamera({self.serial_number}). Actual value is {actual_width}."
             )
-        if self.height is not None and self.height != actual_height:
+        if self.capture_height is not None and self.capture_height != actual_height:
             raise OSError(
-                f"Can't set {self.height=} for IntelRealSenseCamera({self.serial_number}). Actual value is {actual_height}."
+                f"Can't set {self.capture_height=} for RealSenseCamera({self.serial_number}). Actual value is {actual_height}."
             )
 
         self.fps = round(actual_fps)
-        self.width = round(actual_width)
-        self.height = round(actual_height)
+        self.capture_width = round(actual_width)
+        self.capture_height = round(actual_height)
 
         self.is_connected = True
 
@@ -342,12 +367,12 @@ class IntelRealSenseCamera:
         If you are reading data from other sensors, we advise to use `camera.async_read()` which is non blocking version of `camera.read()`.
         """
         if not self.is_connected:
-            raise RobotDeviceNotConnectedError(
-                f"IntelRealSenseCamera({self.serial_number}) is not connected. Try running `camera.connect()` first."
+            raise DeviceNotConnectedError(
+                f"RealSenseCamera({self.serial_number}) is not connected. Try running `camera.connect()` first."
             )
 
         if self.mock:
-            import tests.mock_cv2 as cv2
+            import tests.cameras.mock_cv2 as cv2
         else:
             import cv2
 
@@ -358,7 +383,7 @@ class IntelRealSenseCamera:
         color_frame = frame.get_color_frame()
 
         if not color_frame:
-            raise OSError(f"Can't capture color image from IntelRealSenseCamera({self.serial_number}).")
+            raise OSError(f"Can't capture color image from RealSenseCamera({self.serial_number}).")
 
         color_image = np.asanyarray(color_frame.get_data())
 
@@ -373,7 +398,7 @@ class IntelRealSenseCamera:
             color_image = cv2.cvtColor(color_image, cv2.COLOR_RGB2BGR)
 
         h, w, _ = color_image.shape
-        if h != self.height or w != self.width:
+        if h != self.capture_height or w != self.capture_width:
             raise OSError(
                 f"Can't capture color image with expected height and width ({self.height} x {self.width}). ({h} x {w}) returned instead."
             )
@@ -390,12 +415,12 @@ class IntelRealSenseCamera:
         if self.use_depth:
             depth_frame = frame.get_depth_frame()
             if not depth_frame:
-                raise OSError(f"Can't capture depth image from IntelRealSenseCamera({self.serial_number}).")
+                raise OSError(f"Can't capture depth image from RealSenseCamera({self.serial_number}).")
 
             depth_map = np.asanyarray(depth_frame.get_data())
 
             h, w = depth_map.shape
-            if h != self.height or w != self.width:
+            if h != self.capture_height or w != self.capture_width:
                 raise OSError(
                     f"Can't capture depth map with expected height and width ({self.height} x {self.width}). ({h} x {w}) returned instead."
                 )
@@ -417,8 +442,8 @@ class IntelRealSenseCamera:
     def async_read(self):
         """Access the latest color image"""
         if not self.is_connected:
-            raise RobotDeviceNotConnectedError(
-                f"IntelRealSenseCamera({self.serial_number}) is not connected. Try running `camera.connect()` first."
+            raise DeviceNotConnectedError(
+                f"RealSenseCamera({self.serial_number}) is not connected. Try running `camera.connect()` first."
             )
 
         if self.thread is None:
@@ -444,8 +469,8 @@ class IntelRealSenseCamera:
 
     def disconnect(self):
         if not self.is_connected:
-            raise RobotDeviceNotConnectedError(
-                f"IntelRealSenseCamera({self.serial_number}) is not connected. Try running `camera.connect()` first."
+            raise DeviceNotConnectedError(
+                f"RealSenseCamera({self.serial_number}) is not connected. Try running `camera.connect()` first."
             )
 
         if self.thread is not None and self.thread.is_alive():
@@ -467,14 +492,14 @@ class IntelRealSenseCamera:
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(
-        description="Save a few frames using `IntelRealSenseCamera` for all cameras connected to the computer, or a selected subset."
+        description="Save a few frames using `RealSenseCamera` for all cameras connected to the computer, or a selected subset."
     )
     parser.add_argument(
         "--serial-numbers",
         type=int,
         nargs="*",
         default=None,
-        help="List of serial numbers used to instantiate the `IntelRealSenseCamera`. If not provided, find and use all available camera indices.",
+        help="List of serial numbers used to instantiate the `RealSenseCamera`. If not provided, find and use all available camera indices.",
     )
     parser.add_argument(
         "--fps",

lerobot/common/cameras/intel/configuration_realsense.py

@@ -1,64 +1,35 @@
-import abc
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from dataclasses import dataclass
 
-import draccus
-
-
-@dataclass
-class CameraConfig(draccus.ChoiceRegistry, abc.ABC):
-    @property
-    def type(self) -> str:
-        return self.get_choice_name(self.__class__)
-
-
-@CameraConfig.register_subclass("opencv")
-@dataclass
-class OpenCVCameraConfig(CameraConfig):
-    """
-    Example of tested options for Intel Real Sense D405:
-
-    ```python
-    OpenCVCameraConfig(0, 30, 640, 480)
-    OpenCVCameraConfig(0, 60, 640, 480)
-    OpenCVCameraConfig(0, 90, 640, 480)
-    OpenCVCameraConfig(0, 30, 1280, 720)
-    ```
-    """
-
-    camera_index: int
-    fps: int | None = None
-    width: int | None = None
-    height: int | None = None
-    color_mode: str = "rgb"
-    channels: int | None = None
-    rotation: int | None = None
-    mock: bool = False
-
-    def __post_init__(self):
-        if self.color_mode not in ["rgb", "bgr"]:
-            raise ValueError(
-                f"`color_mode` is expected to be 'rgb' or 'bgr', but {self.color_mode} is provided."
-            )
-
-        self.channels = 3
-
-        if self.rotation not in [-90, None, 90, 180]:
-            raise ValueError(f"`rotation` must be in [-90, None, 90, 180] (got {self.rotation})")
+from ..configs import CameraConfig
 
 
 @CameraConfig.register_subclass("intelrealsense")
 @dataclass
-class IntelRealSenseCameraConfig(CameraConfig):
+class RealSenseCameraConfig(CameraConfig):
     """
     Example of tested options for Intel Real Sense D405:
 
     ```python
-    IntelRealSenseCameraConfig(128422271347, 30, 640, 480)
-    IntelRealSenseCameraConfig(128422271347, 60, 640, 480)
-    IntelRealSenseCameraConfig(128422271347, 90, 640, 480)
-    IntelRealSenseCameraConfig(128422271347, 30, 1280, 720)
-    IntelRealSenseCameraConfig(128422271347, 30, 640, 480, use_depth=True)
-    IntelRealSenseCameraConfig(128422271347, 30, 640, 480, rotation=90)
+    RealSenseCameraConfig(128422271347, 30, 640, 480)
+    RealSenseCameraConfig(128422271347, 60, 640, 480)
+    RealSenseCameraConfig(128422271347, 90, 640, 480)
+    RealSenseCameraConfig(128422271347, 30, 1280, 720)
+    RealSenseCameraConfig(128422271347, 30, 640, 480, use_depth=True)
+    RealSenseCameraConfig(128422271347, 30, 640, 480, rotation=90)
     ```
     """
 

lerobot/common/cameras/interface_camera_sdk.py

@@ -0,0 +1,305 @@
+# ruff: noqa: N802,N803
+import abc
+from typing import Optional, Tuple
+
+import numpy as np
+
+
+# --- Interface Definition ---
+class IVideoCapture(abc.ABC):
+    """Interface for the cv2.VideoCapture class."""
+
+    @abc.abstractmethod
+    def __init__(self, index: int | str, backend: Optional[int] = None):
+        pass
+
+    @abc.abstractmethod
+    def isOpened(self) -> bool:
+        pass
+
+    @abc.abstractmethod
+    def release(self) -> None:
+        pass
+
+    @abc.abstractmethod
+    def set(self, propId: int, value: float) -> bool:
+        pass
+
+    @abc.abstractmethod
+    def get(self, propId: int) -> float:
+        pass
+
+    @abc.abstractmethod
+    def read(self) -> Tuple[bool, Optional[np.ndarray]]:
+        pass
+
+
+class IOpenCVSDK(abc.ABC):
+    """Interface defining the contract for OpenCV SDK interactions."""
+
+    # --- Constants ---
+    CAP_PROP_FPS: int
+    CAP_PROP_FRAME_WIDTH: int
+    CAP_PROP_FRAME_HEIGHT: int
+    COLOR_BGR2RGB: int
+    ROTATE_90_COUNTERCLOCKWISE: int
+    ROTATE_90_CLOCKWISE: int
+    ROTATE_180: int
+    CAP_V4L2: int
+    CAP_DSHOW: int
+    CAP_AVFOUNDATION: int
+    CAP_ANY: int
+
+    # --- Inner Class Type Hint ---
+    VideoCapture: type[IVideoCapture]
+
+    # --- Methods ---
+    @abc.abstractmethod
+    def setNumThreads(self, nthreads: int) -> None:
+        pass
+
+    @abc.abstractmethod
+    def cvtColor(self, src: np.ndarray, code: int) -> np.ndarray:
+        pass
+
+    @abc.abstractmethod
+    def rotate(self, src: np.ndarray, rotateCode: int) -> np.ndarray:
+        pass
+
+
+# --- Real SDK Adapter ---
+class OpenCVSDKAdapter(IOpenCVSDK):
+    """Adapts the real cv2 library to the IOpenCVSDK interface."""
+
+    _cv2 = None
+
+    def __init__(self):
+        try:
+            import cv2
+
+            OpenCVSDKAdapter._cv2 = cv2
+        except ImportError as e:
+            raise ImportError(
+                "OpenCV (cv2) is not installed. Please install it to use the real camera."
+            ) from e
+
+        # --- Constants ---
+        self.CAP_PROP_FPS = self._cv2.CAP_PROP_FPS
+        self.CAP_PROP_FRAME_WIDTH = self._cv2.CAP_PROP_FRAME_WIDTH
+        self.CAP_PROP_FRAME_HEIGHT = self._cv2.CAP_PROP_FRAME_HEIGHT
+        self.COLOR_BGR2RGB = self._cv2.COLOR_BGR2RGB
+        self.ROTATE_90_COUNTERCLOCKWISE = self._cv2.ROTATE_90_COUNTERCLOCKWISE
+        self.ROTATE_90_CLOCKWISE = self._cv2.ROTATE_90_CLOCKWISE
+        self.ROTATE_180 = self._cv2.ROTATE_180
+        self.CAP_V4L2 = self._cv2.CAP_V4L2
+        self.CAP_DSHOW = self._cv2.CAP_DSHOW
+        self.CAP_AVFOUNDATION = self._cv2.CAP_AVFOUNDATION
+        self.CAP_ANY = self._cv2.CAP_ANY
+
+        # --- Inner Class Implementation ---
+        class RealVideoCapture(IVideoCapture):
+            def __init__(self, index: int | str, backend: Optional[int] = None):
+                self._cap = OpenCVSDKAdapter._cv2.VideoCapture(index, backend)
+
+            def isOpened(self) -> bool:
+                return self._cap.isOpened()
+
+            def release(self) -> None:
+                self._cap.release()
+
+            def set(self, propId: int, value: float) -> bool:
+                return self._cap.set(propId, value)
+
+            def get(self, propId: int) -> float:
+                return self._cap.get(propId)
+
+            def read(self) -> Tuple[bool, Optional[np.ndarray]]:
+                return self._cap.read()
+
+            def __del__(self):
+                if hasattr(self, "_cap") and self._cap and self._cap.isOpened():
+                    self._cap.release()
+
+        self.VideoCapture = RealVideoCapture
+
+    # --- Methods ---
+    def setNumThreads(self, nthreads: int) -> None:
+        self._cv2.setNumThreads(nthreads)
+
+    def cvtColor(self, src: np.ndarray, code: int) -> np.ndarray:
+        return self._cv2.cvtColor(src, code)
+
+    def rotate(self, src: np.ndarray, rotateCode: int) -> np.ndarray:
+        return self._cv2.rotate(src, rotateCode)
+
+
+# Emulates the cheap USB camera
+VALID_INDICES = {0, 1, 2, "/dev/video0", "/dev/video1", "/dev/video2"}
+DEFAULT_FPS = 30.0
+DEFAULT_WIDTH = 1280
+DEFAULT_HEIGHT = 720
+
+
+# --- Fake SDK Adapter ---
+class FakeOpenCVSDKAdapter(IOpenCVSDK):
+    """Implements the IOpenCVSDK interface with fake behavior for testing."""
+
+    # --- Constants ---
+    CAP_PROP_FPS = DEFAULT_FPS
+    CAP_PROP_FRAME_WIDTH = DEFAULT_WIDTH
+    CAP_PROP_FRAME_HEIGHT = DEFAULT_HEIGHT
+    COLOR_BGR2RGB = 99
+    ROTATE_90_COUNTERCLOCKWISE = -90
+    ROTATE_90_CLOCKWISE = 90
+    ROTATE_180 = 180
+    CAP_V4L2 = 91
+    CAP_DSHOW = 92
+    CAP_AVFOUNDATION = 93
+    CAP_ANY = 90
+
+    _cameras_opened: dict[int | str, bool] = {}
+    _camera_properties: dict[tuple[int | str, int], float] = {}
+    _simulated_image: np.ndarray = np.random.randint(
+        0, 256, (DEFAULT_HEIGHT, DEFAULT_WIDTH, 3), dtype=np.uint8
+    )
+    _simulated_fps: float = DEFAULT_FPS
+    _image_read_count: int = 0
+    _fail_read_after: Optional[int] = None  # Simulate read failure
+
+    @classmethod
+    def init_configure_fake(
+        cls,
+        simulated_image: Optional[np.ndarray] = None,
+        simulated_fps: Optional[float] = None,
+        fail_read_after: Optional[int] = None,
+    ):
+        if simulated_image is not None:
+            cls._simulated_image = simulated_image
+        if simulated_fps is not None:
+            cls._simulated_fps = simulated_fps
+        cls._fail_read_after = fail_read_after
+        cls._image_read_count = 0
+        cls._cameras_opened = {}
+        cls._camera_properties = {}
+
+    @classmethod
+    def configure_fake_simulated_image(cls, simulated_image: Optional[np.ndarray] = None):
+        if simulated_image is not None:
+            cls._simulated_image = simulated_image
+
+    @classmethod
+    def configure_fail_read_after(cls, fail_read_after: Optional[int] = None):
+        cls._fail_read_after = fail_read_after
+
+    @classmethod
+    def configure_fake_simulated_fps(cls, simulated_fps: Optional[float] = None):
+        if simulated_fps is not None:
+            cls._simulated_fps = simulated_fps
+
+    # --- Inner Class Implementation ---
+    class FakeVideoCapture(IVideoCapture):
+        def __init__(self, index: int | str, backend: Optional[int] = None):
+            self.index = index
+            self.backend = backend
+            valid_indices = VALID_INDICES
+            if self.index in valid_indices:
+                FakeOpenCVSDKAdapter._cameras_opened[self.index] = True
+                print(f"[FAKE SDK] Opened camera {self.index}")
+                # Set some default fake properties
+                FakeOpenCVSDKAdapter._camera_properties[(self.index, FakeOpenCVSDKAdapter.CAP_PROP_FPS)] = (
+                    DEFAULT_FPS
+                )
+                FakeOpenCVSDKAdapter._camera_properties[
+                    (self.index, FakeOpenCVSDKAdapter.CAP_PROP_FRAME_WIDTH)
+                ] = float(FakeOpenCVSDKAdapter._simulated_image.shape[1])
+                FakeOpenCVSDKAdapter._camera_properties[
+                    (self.index, FakeOpenCVSDKAdapter.CAP_PROP_FRAME_HEIGHT)
+                ] = float(FakeOpenCVSDKAdapter._simulated_image.shape[0])
+            else:
+                FakeOpenCVSDKAdapter._cameras_opened[self.index] = False
+                print(f"[FAKE SDK] Failed to open camera {self.index}")
+
+        def isOpened(self) -> bool:
+            return FakeOpenCVSDKAdapter._cameras_opened.get(self.index, False)
+
+        def release(self) -> None:
+            if self.index in FakeOpenCVSDKAdapter._cameras_opened:
+                FakeOpenCVSDKAdapter._cameras_opened[self.index] = False
+                print(f"[FAKE SDK] Released camera {self.index}")
+                # Clear properties on release
+                props_to_remove = [k for k in FakeOpenCVSDKAdapter._camera_properties if k[0] == self.index]
+                for k in props_to_remove:
+                    del FakeOpenCVSDKAdapter._camera_properties[k]
+
+        def set(self, propId: int, value: float) -> bool:
+            if not self.isOpened():
+                return False
+            print(
+                f"[FAKE SDK] Ignoring set property {propId} = {value} for camera {self.index} to preserve state."
+            )
+            # FakeOpenCVSDKAdapter._camera_properties[(self.index, propId)] = value
+            # Simulate failure for specific unrealistic settings if needed
+            return True
+
+        def get(self, propId: int) -> float:
+            if not self.isOpened():
+                return 0.0  # Or raise error? Mimic cv2 behavior
+            val = FakeOpenCVSDKAdapter._camera_properties.get((self.index, propId))
+            print(f"[FAKE SDK] Get property {propId} for camera {self.index} -> {val}")
+            return val
+
+        def read(self) -> Tuple[bool, Optional[np.ndarray]]:
+            if not self.isOpened():
+                print(f"[FAKE SDK] Read failed: Camera {self.index} not open.")
+                return False, None
+
+            FakeOpenCVSDKAdapter._image_read_count += 1
+            if (
+                FakeOpenCVSDKAdapter._fail_read_after is not None
+                and FakeOpenCVSDKAdapter._image_read_count > FakeOpenCVSDKAdapter._fail_read_after
+            ):
+                print(
+                    f"[FAKE SDK] Simulated read failure for camera {self.index} after {FakeOpenCVSDKAdapter._fail_read_after} reads."
+                )
+                return False, None
+
+            print(
+                f"[FAKE SDK] Read image from camera {self.index} (read #{FakeOpenCVSDKAdapter._image_read_count})"
+            )
+            # Return a copy to prevent modification issues if the caller changes it
+            return True, FakeOpenCVSDKAdapter._simulated_image.copy()
+
+        def __del__(self):
+            # Ensure cleanup if garbage collected
+            self.release()
+
+    VideoCapture = FakeVideoCapture  # Assign inner class
+
+    # --- Methods ---
+    def setNumThreads(self, nthreads: int) -> None:
+        print(f"[FAKE SDK] setNumThreads({nthreads}) called.")
+        # No actual behavior needed in fake
+
+    def cvtColor(self, src: np.ndarray, code: int) -> np.ndarray:
+        print(f"[FAKE SDK] cvtColor called with code {code}.")
+        # Just return the source image, or simulate channel swap if needed
+        if code == self.COLOR_BGR2RGB and src.shape[2] == 3:
+            print("[FAKE SDK] Simulating BGR -> RGB conversion.")
+            return src[..., ::-1]
+        return src.copy()
+
+    def rotate(self, src: np.ndarray, rotateCode: int) -> np.ndarray:
+        print(f"[FAKE SDK] rotate called with code {rotateCode}.")
+        if rotateCode == self.ROTATE_90_COUNTERCLOCKWISE:
+            print("[FAKE SDK] Simulating 90 degree counter-clockwise rotation.")
+            rotated_img = np.rot90(np.rot90(np.rot90(src.copy())))
+            return rotated_img
+        elif rotateCode == self.ROTATE_90_CLOCKWISE:
+            print("[FAKE SDK] Simulating 90 degree clockwise rotation.")
+            rotated_img = np.rot90(src.copy())
+            return rotated_img
+        elif rotateCode == self.ROTATE_180:
+            print("[FAKE SDK] Simulating 180 degree rotation.")
+            rotated_img = np.rot90(np.rot90(src.copy()))
+            return rotated_img
+        return src.copy()

lerobot/common/cameras/opencv/__init__.py

@@ -0,0 +1,4 @@
+from .camera_opencv import OpenCVCamera
+from .configuration_opencv import OpenCVCameraConfig
+
+__all__ = ["OpenCVCamera", "OpenCVCameraConfig"]

lerobot/common/cameras/opencv/camera_opencv.py

@@ -1,3 +1,17 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 """
 This file contains utilities for recording frames from cameras. For more info look at `OpenCVCamera` docstring.
 """
@@ -10,19 +24,20 @@ import shutil
 import threading
 import time
 from pathlib import Path
-from threading import Thread
 
 import numpy as np
 from PIL import Image
 
-from lerobot.common.robot_devices.cameras.configs import OpenCVCameraConfig
-from lerobot.common.robot_devices.utils import (
-    RobotDeviceAlreadyConnectedError,
-    RobotDeviceNotConnectedError,
+from lerobot.common.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
+from lerobot.common.utils.robot_utils import (
     busy_wait,
 )
 from lerobot.common.utils.utils import capture_timestamp_utc
 
+from ..camera import Camera
+from ..interface_camera_sdk import IOpenCVSDK, OpenCVSDKAdapter
+from .configuration_opencv import OpenCVCameraConfig
+
 # The maximum opencv device index depends on your operating system. For instance,
 # if you have 3 cameras, they should be associated to index 0, 1, and 2. This is the case
 # on MacOS. However, on Ubuntu, the indices are different like 6, 16, 23.
@@ -31,12 +46,17 @@ from lerobot.common.utils.utils import capture_timestamp_utc
 MAX_OPENCV_INDEX = 60
 
 
-def find_cameras(raise_when_empty=False, max_index_search_range=MAX_OPENCV_INDEX, mock=False) -> list[dict]:
+def find_cameras(
+    raise_when_empty=False, max_index_search_range=MAX_OPENCV_INDEX, cv2_sdk: IOpenCVSDK = None
+) -> list[dict]:
+    if cv2_sdk is None:
+        cv2_sdk = OpenCVSDKAdapter()
+
     cameras = []
     if platform.system() == "Linux":
         print("Linux detected. Finding available camera indices through scanning '/dev/video*' ports")
         possible_ports = [str(port) for port in Path("/dev").glob("video*")]
-        ports = _find_cameras(possible_ports, mock=mock)
+        ports = _find_cameras(possible_ports, cv2_sdk=cv2_sdk)
         for port in ports:
             cameras.append(
                 {
@@ -50,7 +70,7 @@ def find_cameras(raise_when_empty=False, max_index_search_range=MAX_OPENCV_INDEX
             f"scanning all indices from 0 to {MAX_OPENCV_INDEX}"
         )
         possible_indices = range(max_index_search_range)
-        indices = _find_cameras(possible_indices, mock=mock)
+        indices = _find_cameras(possible_indices, cv2_sdk=cv2_sdk)
         for index in indices:
             cameras.append(
                 {
@@ -63,16 +83,14 @@ def find_cameras(raise_when_empty=False, max_index_search_range=MAX_OPENCV_INDEX
 
 
 def _find_cameras(
-    possible_camera_ids: list[int | str], raise_when_empty=False, mock=False
+    possible_camera_ids: list[int | str], raise_when_empty=False, cv2_sdk: IOpenCVSDK = None
 ) -> list[int | str]:
-    if mock:
-        import tests.mock_cv2 as cv2
-    else:
-        import cv2
+    if cv2_sdk is None:
+        cv2_sdk = OpenCVSDKAdapter()
 
     camera_ids = []
     for camera_idx in possible_camera_ids:
-        camera = cv2.VideoCapture(camera_idx)
+        camera = cv2_sdk.VideoCapture(camera_idx)
         is_open = camera.isOpened()
         camera.release()
 
@@ -113,25 +131,28 @@ def save_images_from_cameras(
     width=None,
     height=None,
     record_time_s=2,
-    mock=False,
+    cv2_sdk: IOpenCVSDK = None,
 ):
     """
     Initializes all the cameras and saves images to the directory. Useful to visually identify the camera
     associated to a given camera index.
     """
+    if cv2_sdk is None:
+        cv2_sdk = OpenCVSDKAdapter()
+
     if camera_ids is None or len(camera_ids) == 0:
-        camera_infos = find_cameras(mock=mock)
+        camera_infos = find_cameras(cv2_sdk=cv2_sdk)
         camera_ids = [cam["index"] for cam in camera_infos]
 
     print("Connecting cameras")
     cameras = []
     for cam_idx in camera_ids:
-        config = OpenCVCameraConfig(camera_index=cam_idx, fps=fps, width=width, height=height, mock=mock)
-        camera = OpenCVCamera(config)
+        config = OpenCVCameraConfig(camera_index=cam_idx, fps=fps, width=width, height=height)
+        camera = OpenCVCamera(config, cv2_sdk=cv2_sdk)
         camera.connect()
         print(
-            f"OpenCVCamera({camera.camera_index}, fps={camera.fps}, width={camera.width}, "
-            f"height={camera.height}, color_mode={camera.color_mode})"
+            f"OpenCVCamera({camera.camera_index}, fps={camera.fps}, width={camera.capture_width}, "
+            f"height={camera.capture_height}, color_mode={camera.color_mode})"
         )
         cameras.append(camera)
 
@@ -176,7 +197,7 @@ def save_images_from_cameras(
     print(f"Images have been saved to {images_dir}")
 
 
-class OpenCVCamera:
+class OpenCVCamera(Camera):
     """
     The OpenCVCamera class allows to efficiently record images from cameras. It relies on opencv2 to communicate
     with the cameras. Most cameras are compatible. For more info, see the [Video I/O with OpenCV Overview](https://docs.opencv.org/4.x/d0/da7/videoio_overview.html).
@@ -214,11 +235,16 @@ class OpenCVCamera:
     ```
     """
 
-    def __init__(self, config: OpenCVCameraConfig):
+    def __init__(self, config: OpenCVCameraConfig, cv2_sdk: IOpenCVSDK = None):
         self.config = config
         self.camera_index = config.camera_index
         self.port = None
 
+        if cv2_sdk is None:
+            cv2_sdk = OpenCVSDKAdapter()
+
+        self.cv2_sdk = cv2_sdk
+
         # Linux uses ports for connecting to cameras
         if platform.system() == "Linux":
             if isinstance(self.camera_index, int):
@@ -230,12 +256,21 @@ class OpenCVCamera:
             else:
                 raise ValueError(f"Please check the provided camera_index: {self.camera_index}")
 
+        # Store the raw (capture) resolution from the config.
+        self.capture_width = config.width
+        self.capture_height = config.height
+
+        # If rotated by ±90, swap width and height.
+        if config.rotation in [-90, 90]:
+            self.width = config.height
+            self.height = config.width
+        else:
+            self.width = config.width
+            self.height = config.height
+
         self.fps = config.fps
-        self.width = config.width
-        self.height = config.height
         self.channels = config.channels
         self.color_mode = config.color_mode
-        self.mock = config.mock
 
         self.camera = None
         self.is_connected = False
@@ -244,37 +279,38 @@ class OpenCVCamera:
         self.color_image = None
         self.logs = {}
 
-        if self.mock:
-            import tests.mock_cv2 as cv2
-        else:
-            import cv2
-
-        # TODO(aliberts): Do we keep original width/height or do we define them after rotation?
         self.rotation = None
         if config.rotation == -90:
-            self.rotation = cv2.ROTATE_90_COUNTERCLOCKWISE
+            self.rotation = cv2_sdk.ROTATE_90_COUNTERCLOCKWISE
         elif config.rotation == 90:
-            self.rotation = cv2.ROTATE_90_CLOCKWISE
+            self.rotation = cv2_sdk.ROTATE_90_CLOCKWISE
         elif config.rotation == 180:
-            self.rotation = cv2.ROTATE_180
+            self.rotation = cv2_sdk.ROTATE_180
 
     def connect(self):
         if self.is_connected:
-            raise RobotDeviceAlreadyConnectedError(f"OpenCVCamera({self.camera_index}) is already connected.")
+            raise DeviceAlreadyConnectedError(f"OpenCVCamera({self.camera_index}) is already connected.")
 
-        if self.mock:
-            import tests.mock_cv2 as cv2
-        else:
-            import cv2
+        cv2_sdk = self.cv2_sdk
 
-            # Use 1 thread to avoid blocking the main thread. Especially useful during data collection
-            # when other threads are used to save the images.
-            cv2.setNumThreads(1)
+        # Use 1 thread to avoid blocking the main thread. Especially useful during data collection
+        # when other threads are used to save the images.
+        cv2_sdk.setNumThreads(1)
+
+        backend = (
+            cv2_sdk.CAP_V4L2
+            if platform.system() == "Linux"
+            else cv2_sdk.CAP_DSHOW
+            if platform.system() == "Windows"
+            else cv2_sdk.CAP_AVFOUNDATION
+            if platform.system() == "Darwin"
+            else cv2_sdk.CAP_ANY
+        )
 
         camera_idx = f"/dev/video{self.camera_index}" if platform.system() == "Linux" else self.camera_index
         # First create a temporary camera trying to access `camera_index`,
         # and verify it is a valid camera by calling `isOpened`.
-        tmp_camera = cv2.VideoCapture(camera_idx)
+        tmp_camera = cv2_sdk.VideoCapture(camera_idx, backend)
         is_camera_open = tmp_camera.isOpened()
         # Release camera to make it accessible for `find_camera_indices`
         tmp_camera.release()
@@ -284,7 +320,7 @@ class OpenCVCamera:
         # valid cameras.
         if not is_camera_open:
             # Verify that the provided `camera_index` is valid before printing the traceback
-            cameras_info = find_cameras()
+            cameras_info = find_cameras(cv2_sdk=cv2_sdk)
             available_cam_ids = [cam["index"] for cam in cameras_info]
             if self.camera_index not in available_cam_ids:
                 raise ValueError(
@@ -297,18 +333,18 @@ class OpenCVCamera:
         # Secondly, create the camera that will be used downstream.
         # Note: For some unknown reason, calling `isOpened` blocks the camera which then
         # needs to be re-created.
-        self.camera = cv2.VideoCapture(camera_idx)
+        self.camera = cv2_sdk.VideoCapture(camera_idx, backend)
 
         if self.fps is not None:
-            self.camera.set(cv2.CAP_PROP_FPS, self.fps)
-        if self.width is not None:
-            self.camera.set(cv2.CAP_PROP_FRAME_WIDTH, self.width)
-        if self.height is not None:
-            self.camera.set(cv2.CAP_PROP_FRAME_HEIGHT, self.height)
+            self.camera.set(cv2_sdk.CAP_PROP_FPS, self.fps)
+        if self.capture_width is not None:
+            self.camera.set(cv2_sdk.CAP_PROP_FRAME_WIDTH, self.capture_width)
+        if self.capture_height is not None:
+            self.camera.set(cv2_sdk.CAP_PROP_FRAME_HEIGHT, self.capture_height)
 
-        actual_fps = self.camera.get(cv2.CAP_PROP_FPS)
-        actual_width = self.camera.get(cv2.CAP_PROP_FRAME_WIDTH)
-        actual_height = self.camera.get(cv2.CAP_PROP_FRAME_HEIGHT)
+        actual_fps = self.camera.get(cv2_sdk.CAP_PROP_FPS)
+        actual_width = self.camera.get(cv2_sdk.CAP_PROP_FRAME_WIDTH)
+        actual_height = self.camera.get(cv2_sdk.CAP_PROP_FRAME_HEIGHT)
 
         # Using `math.isclose` since actual fps can be a float (e.g. 29.9 instead of 30)
         if self.fps is not None and not math.isclose(self.fps, actual_fps, rel_tol=1e-3):
@@ -316,19 +352,22 @@ class OpenCVCamera:
             raise OSError(
                 f"Can't set {self.fps=} for OpenCVCamera({self.camera_index}). Actual value is {actual_fps}."
             )
-        if self.width is not None and not math.isclose(self.width, actual_width, rel_tol=1e-3):
+        if self.capture_width is not None and not math.isclose(
+            self.capture_width, actual_width, rel_tol=1e-3
+        ):
             raise OSError(
-                f"Can't set {self.width=} for OpenCVCamera({self.camera_index}). Actual value is {actual_width}."
+                f"Can't set {self.capture_width=} for OpenCVCamera({self.camera_index}). Actual value is {actual_width}."
             )
-        if self.height is not None and not math.isclose(self.height, actual_height, rel_tol=1e-3):
+        if self.capture_height is not None and not math.isclose(
+            self.capture_height, actual_height, rel_tol=1e-3
+        ):
             raise OSError(
-                f"Can't set {self.height=} for OpenCVCamera({self.camera_index}). Actual value is {actual_height}."
+                f"Can't set {self.capture_height=} for OpenCVCamera({self.camera_index}). Actual value is {actual_height}."
             )
 
         self.fps = round(actual_fps)
-        self.width = round(actual_width)
-        self.height = round(actual_height)
-
+        self.capture_width = round(actual_width)
+        self.capture_height = round(actual_height)
         self.is_connected = True
 
     def read(self, temporary_color_mode: str | None = None) -> np.ndarray:
@@ -339,10 +378,12 @@ class OpenCVCamera:
         If you are reading data from other sensors, we advise to use `camera.async_read()` which is non blocking version of `camera.read()`.
         """
         if not self.is_connected:
-            raise RobotDeviceNotConnectedError(
+            raise DeviceNotConnectedError(
                 f"OpenCVCamera({self.camera_index}) is not connected. Try running `camera.connect()` first."
             )
 
+        cv2_sdk = self.cv2_sdk
+
         start_time = time.perf_counter()
 
         ret, color_image = self.camera.read()
@@ -361,21 +402,16 @@ class OpenCVCamera:
         # However, Deep Learning framework such as LeRobot uses RGB format as default to train neural networks,
         # so we convert the image color from BGR to RGB.
         if requested_color_mode == "rgb":
-            if self.mock:
-                import tests.mock_cv2 as cv2
-            else:
-                import cv2
-
-            color_image = cv2.cvtColor(color_image, cv2.COLOR_BGR2RGB)
+            color_image = cv2_sdk.cvtColor(color_image, cv2_sdk.COLOR_BGR2RGB)
 
         h, w, _ = color_image.shape
-        if h != self.height or w != self.width:
+        if h != self.capture_height or w != self.capture_width:
             raise OSError(
                 f"Can't capture color image with expected height and width ({self.height} x {self.width}). ({h} x {w}) returned instead."
             )
 
         if self.rotation is not None:
-            color_image = cv2.rotate(color_image, self.rotation)
+            color_image = cv2_sdk.rotate(color_image, self.rotation)
 
         # log the number of seconds it took to read the image
         self.logs["delta_timestamp_s"] = time.perf_counter() - start_time
@@ -396,13 +432,13 @@ class OpenCVCamera:
 
     def async_read(self):
         if not self.is_connected:
-            raise RobotDeviceNotConnectedError(
+            raise DeviceNotConnectedError(
                 f"OpenCVCamera({self.camera_index}) is not connected. Try running `camera.connect()` first."
             )
 
         if self.thread is None:
             self.stop_event = threading.Event()
-            self.thread = Thread(target=self.read_loop, args=())
+            self.thread = threading.Thread(target=self.read_loop, args=())
             self.thread.daemon = True
             self.thread.start()
 
@@ -418,7 +454,7 @@ class OpenCVCamera:
 
     def disconnect(self):
         if not self.is_connected:
-            raise RobotDeviceNotConnectedError(
+            raise DeviceNotConnectedError(
                 f"OpenCVCamera({self.camera_index}) is not connected. Try running `camera.connect()` first."
             )
 

lerobot/common/cameras/opencv/configuration_opencv.py

@@ -0,0 +1,37 @@
+from dataclasses import dataclass
+
+from ..configs import CameraConfig
+
+
+@CameraConfig.register_subclass("opencv")
+@dataclass
+class OpenCVCameraConfig(CameraConfig):
+    """
+    Example of tested options for Intel Real Sense D405:
+
+    ```python
+    OpenCVCameraConfig(0, 30, 640, 480)
+    OpenCVCameraConfig(0, 60, 640, 480)
+    OpenCVCameraConfig(0, 90, 640, 480)
+    OpenCVCameraConfig(0, 30, 1280, 720)
+    ```
+    """
+
+    camera_index: int
+    fps: int | None = None
+    width: int | None = None
+    height: int | None = None
+    color_mode: str = "rgb"
+    channels: int | None = None
+    rotation: int | None = None
+
+    def __post_init__(self):
+        if self.color_mode not in ["rgb", "bgr"]:
+            raise ValueError(
+                f"`color_mode` is expected to be 'rgb' or 'bgr', but {self.color_mode} is provided."
+            )
+
+        self.channels = 3
+
+        if self.rotation not in [-90, None, 90, 180]:
+            raise ValueError(f"`rotation` must be in [-90, None, 90, 180] (got {self.rotation})")

lerobot/common/cameras/utils.py

@@ -0,0 +1,21 @@
+from .camera import Camera
+from .configs import CameraConfig
+
+
+def make_cameras_from_configs(camera_configs: dict[str, CameraConfig]) -> dict[str, Camera]:
+    cameras = {}
+
+    for key, cfg in camera_configs.items():
+        if cfg.type == "opencv":
+            from .opencv import OpenCVCamera
+
+            cameras[key] = OpenCVCamera(cfg)
+
+        elif cfg.type == "intelrealsense":
+            from .intel.camera_realsense import RealSenseCamera
+
+            cameras[key] = RealSenseCamera(cfg)
+        else:
+            raise ValueError(f"The motor type '{cfg.type}' is not valid.")
+
+    return cameras

lerobot/common/constants.py

@@ -1,15 +1,31 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 # keys
 import os
 from pathlib import Path
 
 from huggingface_hub.constants import HF_HOME
 
-OBS_ENV = "observation.environment_state"
-OBS_ROBOT = "observation.state"
+OBS_ENV_STATE = "observation.environment_state"
+OBS_STATE = "observation.state"
 OBS_IMAGE = "observation.image"
 OBS_IMAGES = "observation.images"
 ACTION = "action"
 
+ROBOTS = "robots"
+TELEOPERATORS = "teleoperators"
+
 # files & directories
 CHECKPOINTS_DIR = "checkpoints"
 LAST_CHECKPOINT_LINK = "last"
@@ -21,12 +37,16 @@ OPTIMIZER_STATE = "optimizer_state.safetensors"
 OPTIMIZER_PARAM_GROUPS = "optimizer_param_groups.json"
 SCHEDULER_STATE = "scheduler_state.json"
 
-# cache dir
-default_cache_path = Path(HF_HOME) / "lerobot"
-HF_LEROBOT_HOME = Path(os.getenv("HF_LEROBOT_HOME", default_cache_path)).expanduser()
-
 if "LEROBOT_HOME" in os.environ:
     raise ValueError(
         f"You have a 'LEROBOT_HOME' environment variable set to '{os.getenv('LEROBOT_HOME')}'.\n"
         "'LEROBOT_HOME' is deprecated, please use 'HF_LEROBOT_HOME' instead."
     )
+
+# cache dir
+default_cache_path = Path(HF_HOME) / "lerobot"
+HF_LEROBOT_HOME = Path(os.getenv("HF_LEROBOT_HOME", default_cache_path)).expanduser()
+
+# calibration dir
+default_calibration_path = HF_LEROBOT_HOME / ".calibration"
+HF_LEROBOT_CALIBRATION = Path(os.getenv("HF_LEROBOT_CALIBRATION", default_calibration_path)).expanduser()

lerobot/common/datasets/backward_compatibility.py

@@ -1,3 +1,17 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import packaging.version
 
 V2_MESSAGE = """

lerobot/common/datasets/lerobot_dataset.py

@@ -13,6 +13,7 @@
 # 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 contextlib
 import logging
 import shutil
 from pathlib import Path
@@ -27,6 +28,7 @@ import torch.utils
 from datasets import concatenate_datasets, load_dataset
 from huggingface_hub import HfApi, snapshot_download
 from huggingface_hub.constants import REPOCARD_NAME
+from huggingface_hub.errors import RevisionNotFoundError
 
 from lerobot.common.constants import HF_LEROBOT_HOME
 from lerobot.common.datasets.compute_stats import aggregate_stats, compute_episode_stats
@@ -65,11 +67,12 @@ from lerobot.common.datasets.utils import (
 )
 from lerobot.common.datasets.video_utils import (
     VideoFrame,
-    decode_video_frames_torchvision,
+    decode_video_frames,
     encode_video_frames,
+    get_safe_default_codec,
     get_video_info,
 )
-from lerobot.common.robot_devices.robots.utils import Robot
+from lerobot.common.robots.utils import Robot
 
 CODEBASE_VERSION = "v2.1"
 
@@ -460,8 +463,8 @@ class LeRobotDataset(torch.utils.data.Dataset):
             download_videos (bool, optional): Flag to download the videos. Note that when set to True but the
                 video files are already present on local disk, they won't be downloaded again. Defaults to
                 True.
-            video_backend (str | None, optional): Video backend to use for decoding videos. There is currently
-                a single option which is the pyav decoder used by Torchvision. Defaults to pyav.
+            video_backend (str | None, optional): Video backend to use for decoding videos. Defaults to torchcodec when available int the platform; otherwise, defaults to 'pyav'.
+                You can also use the 'pyav' decoder used by Torchvision, which used to be the default option, or 'video_reader' which is another decoder of Torchvision.
         """
         super().__init__()
         self.repo_id = repo_id
@@ -471,7 +474,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
         self.episodes = episodes
         self.tolerance_s = tolerance_s
         self.revision = revision if revision else CODEBASE_VERSION
-        self.video_backend = video_backend if video_backend else "pyav"
+        self.video_backend = video_backend if video_backend else get_safe_default_codec()
         self.delta_indices = None
 
         # Unused attributes
@@ -517,6 +520,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
         branch: str | None = None,
         tags: list | None = None,
         license: str | None = "apache-2.0",
+        tag_version: bool = True,
         push_videos: bool = True,
         private: bool = False,
         allow_patterns: list[str] | str | None = None,
@@ -562,6 +566,11 @@ class LeRobotDataset(torch.utils.data.Dataset):
             )
             card.push_to_hub(repo_id=self.repo_id, repo_type="dataset", revision=branch)
 
+        if tag_version:
+            with contextlib.suppress(RevisionNotFoundError):
+                hub_api.delete_tag(self.repo_id, tag=CODEBASE_VERSION, repo_type="dataset")
+            hub_api.create_tag(self.repo_id, tag=CODEBASE_VERSION, revision=branch, repo_type="dataset")
+
     def pull_from_repo(
         self,
         allow_patterns: list[str] | str | None = None,
@@ -699,9 +708,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
         item = {}
         for vid_key, query_ts in query_timestamps.items():
             video_path = self.root / self.meta.get_video_file_path(ep_idx, vid_key)
-            frames = decode_video_frames_torchvision(
-                video_path, query_ts, self.tolerance_s, self.video_backend
-            )
+            frames = decode_video_frames(video_path, query_ts, self.tolerance_s, self.video_backend)
             item[vid_key] = frames.squeeze(0)
 
         return item
@@ -1021,7 +1028,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
         obj.delta_timestamps = None
         obj.delta_indices = None
         obj.episode_data_index = None
-        obj.video_backend = video_backend if video_backend is not None else "pyav"
+        obj.video_backend = video_backend if video_backend is not None else get_safe_default_codec()
         return obj
 
 
@@ -1046,7 +1053,7 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
         super().__init__()
         self.repo_ids = repo_ids
         self.root = Path(root) if root else HF_LEROBOT_HOME
-        self.tolerances_s = tolerances_s if tolerances_s else {repo_id: 1e-4 for repo_id in repo_ids}
+        self.tolerances_s = tolerances_s if tolerances_s else dict.fromkeys(repo_ids, 0.0001)
         # Construct the underlying datasets passing everything but `transform` and `delta_timestamps` which
         # are handled by this class.
         self._datasets = [

lerobot/common/datasets/push_dataset_to_hub/_aloha_raw_urls/mobile_cabinet.txt

@@ -1,85 +0,0 @@
-https://drive.google.com/file/d/1_SOJkgfP5yZyVjMhTt3nwhvyUjcnlI51/view?usp=drive_link
-https://drive.google.com/file/d/1rmgN8UUzph1qwJnzG1d-uOafodn-gLvb/view?usp=drive_link
-https://drive.google.com/file/d/1NYQ-XxsBVinB6dUoZmVWweT83367P3i2/view?usp=drive_link
-https://drive.google.com/file/d/1oAv_j74zxxCJieMG7r5Vl2BeHK1__3s3/view?usp=drive_link
-https://drive.google.com/file/d/1wFUJQROsrTJt64YRuIeExhFjr2wnK5uu/view?usp=drive_link
-https://drive.google.com/file/d/1KzL3Tt0Le7jVl58XVRUcmigmXjyiuhbK/view?usp=drive_link
-https://drive.google.com/file/d/1qy_YBladeHtianSSGtgAPSHtMin7msvf/view?usp=drive_link
-https://drive.google.com/file/d/1rA_F0V_qL_nyuC_0aBKCisF4-0TIkF2Y/view?usp=drive_link
-https://drive.google.com/file/d/1hw-8qMpz9VgSt62XoASqNRuPECpCwJQP/view?usp=drive_link
-https://drive.google.com/file/d/1BpHOl9rKMzdvNGka6js7C0s40hH6vnDA/view?usp=drive_link
-https://drive.google.com/file/d/1PazhkhiDnJ-OUMyDVDFxEZNKQQqHiNWS/view?usp=drive_link
-https://drive.google.com/file/d/1lZ665R6ATl57dypxH4dGJ2NSt6XYnbuz/view?usp=drive_link
-https://drive.google.com/file/d/1V9HzLaf-tlG15wUzT7KrTDCS_z1vi5NV/view?usp=drive_link
-https://drive.google.com/file/d/1aKauWiXoKqbNwn_2xs4MrmLlaNYlVNmO/view?usp=drive_link
-https://drive.google.com/file/d/1WVD5DFhriO1YmmOgiVHhacR6HWoTPxav/view?usp=drive_link
-https://drive.google.com/file/d/1_X43WgeBAsfkhH9EmpyPki8U9joMeAGC/view?usp=drive_link
-https://drive.google.com/file/d/1t8x0GqWoNKWtnBsB7_D40Z34nL9ak4kf/view?usp=drive_link
-https://drive.google.com/file/d/15V_f26WaKOXjKnq2T3HRWAmtQUi4lbu2/view?usp=drive_link
-https://drive.google.com/file/d/11VFIAsiSDsMOBANgrOcZBpKB9AFWnLy7/view?usp=drive_link
-https://drive.google.com/file/d/1M0NS7vVaxJv3FHnuRYtdwTFYF7We4LxP/view?usp=drive_link
-https://drive.google.com/file/d/1mR0OItTNqFnVLoczcyKYlm6drAy778lO/view?usp=drive_link
-https://drive.google.com/file/d/1NbVFWDQAh-z4JJ4D-Zw6Lps9kdvpqh2j/view?usp=drive_link
-https://drive.google.com/file/d/1JQoZGBzl4W3QG26-n39tefcGN0fDRMbB/view?usp=drive_link
-https://drive.google.com/file/d/1VBjHl-TvZpncopvasIP5G9gecbB2a5f6/view?usp=drive_link
-https://drive.google.com/file/d/1VzSf6zaB21nahm7MsPwroXbJ84NIwq0b/view?usp=drive_link
-https://drive.google.com/file/d/1OtNnfMEydNtZOcivs4k6E_uJSpf8PkGy/view?usp=drive_link
-https://drive.google.com/file/d/14nVvpvsrFr_03Pa_N7MKzwnRwibOUYM6/view?usp=drive_link
-https://drive.google.com/file/d/1M8li6duiO2r3lv_9HhF_XJn0oZUIEK5F/view?usp=drive_link
-https://drive.google.com/file/d/1Cpzea6fO14lxAaNfSBifqoa4ekhCiLD1/view?usp=drive_link
-https://drive.google.com/file/d/1mbxRTm5vlbsY9UJ0jfjM6j9D7kPJjBpG/view?usp=drive_link
-https://drive.google.com/file/d/1RXD1i6IfWsHRlCxVmG04h2h5Ycm_WwZN/view?usp=drive_link
-https://drive.google.com/file/d/1QFqFSwDGOk1BkgGmqgCcc2BRWnJ6R3MA/view?usp=drive_link
-https://drive.google.com/file/d/1bFqWR8DQM0ZUxxtS2bl-RANQvukeFLzp/view?usp=drive_link
-https://drive.google.com/file/d/1pR-rH3yNGoyPdD4hJ6-3lXQ-PstBx9du/view?usp=drive_link
-https://drive.google.com/file/d/107OAwLY-hva9HeQLIK7VCh-ytdDabVjr/view?usp=drive_link
-https://drive.google.com/file/d/1Tpl08QOaSZ37GTO4awFWSdD8wBR9xdlT/view?usp=drive_link
-https://drive.google.com/file/d/1MR164AOM-0S1T6RX8xKTV2IHyaCvpqAW/view?usp=drive_link
-https://drive.google.com/file/d/1_wknJfVnStIhJ82lU_QtcrwahsqYIsr8/view?usp=drive_link
-https://drive.google.com/file/d/1ZuEktWrbYkTx0l5pj3WiZ2CJrfbDOHNo/view?usp=drive_link
-https://drive.google.com/file/d/15G_10hkkkq6yxvyI5NGZirlF-RzduR2F/view?usp=drive_link
-https://drive.google.com/file/d/1DBKxg3ONqh7dhLuX6oh1Yyo2x383V1Hp/view?usp=drive_link
-https://drive.google.com/file/d/1B5iDBkTUr5vopDddV_fHud18SqAHhauS/view?usp=drive_link
-https://drive.google.com/file/d/1acwFV0eenRkki1QcjSKH5xqOtys-P3Pr/view?usp=drive_link
-https://drive.google.com/file/d/1S47BI83xyrh-FKXsvAQqer98Biu_p8XK/view?usp=drive_link
-https://drive.google.com/file/d/1JL6DmBZl3uyq9dyLfgSqtGF06e7E9JwM/view?usp=drive_link
-https://drive.google.com/file/d/16WvRS4Kjog8Pxgr0E3sGGnI01YwL9Uql/view?usp=drive_link
-https://drive.google.com/file/d/12ttGqL33IPWg0-s1SD44rr22M6LiSQBr/view?usp=drive_link
-https://drive.google.com/file/d/1OyZqqnldTU_DliRbr6x0C4a_iWPwIN7j/view?usp=drive_link
-https://drive.google.com/file/d/1oYk00IpLnR9fesLfD15Ebe7nVBffEbcS/view?usp=drive_link
-https://drive.google.com/file/d/1eyE2-MQduCEqCd-5_kl5zsoOEERAzpZD/view?usp=drive_link
-https://drive.google.com/file/d/1ir1Ya-vO0d97pfvbePlUeuKTTRc0qIMU/view?usp=drive_link
-https://drive.google.com/file/d/1hOi-JnqlMt47gVnLZHMTqeojyYVErohl/view?usp=drive_link
-https://drive.google.com/file/d/1NFFw5_PqigQ7xGqsL-MNq2B1r5yAscCf/view?usp=drive_link
-https://drive.google.com/file/d/1uftq1-Zlh8d2sNLWrlVcKYQUwZTD7o24/view?usp=drive_link
-https://drive.google.com/file/d/1-ax19dSLPacVgk000T-m3l4flPcg07pM/view?usp=drive_link
-https://drive.google.com/file/d/126y-lgn86-ZmCz8hooF1THKJGGObw3OB/view?usp=drive_link
-https://drive.google.com/file/d/1JiDniK0VmDIkk92AbBILb8J2Ba59PWML/view?usp=drive_link
-https://drive.google.com/file/d/1kr8nPIRljiU0R4J9SMgj80o1FPQxzu9z/view?usp=drive_link
-https://drive.google.com/file/d/1bbThWRij1pKBh_kFgV8FwK0sXtTHBoLX/view?usp=drive_link
-https://drive.google.com/file/d/1WenzDW6lxk1xkOFm-OiGFfc0ROskAuKU/view?usp=drive_link
-https://drive.google.com/file/d/1MiKRzuzUn1yN-k_6kPJJzIGy7dT-nnsD/view?usp=drive_link
-https://drive.google.com/file/d/17rRg2tcmB-gNhQ0KoZJQmNfyFeoij1jH/view?usp=drive_link
-https://drive.google.com/file/d/11mokBpvrY3ld6sY5WztREtJ1jgqfQV70/view?usp=drive_link
-https://drive.google.com/file/d/1Il_6IOx9NDp1bX_KHizJfBwzTufTmn86/view?usp=drive_link
-https://drive.google.com/file/d/1KswtJGsxJ7eeBDAmNA_aeLjOxcH6MIxa/view?usp=drive_link
-https://drive.google.com/file/d/1gzMhi5uWu4C3Y6WbQ3L-08V96GxTZrRR/view?usp=drive_link
-https://drive.google.com/file/d/1nRQFtaBxfUCYc2W90Qibh0kHCt6YQCfc/view?usp=drive_link
-https://drive.google.com/file/d/1vs-gyW-KheqHbUATwAhA2mmR9GOGw7f_/view?usp=drive_link
-https://drive.google.com/file/d/1MuxzGOA2fgLaHryq82KkQumtuRJGcUOC/view?usp=drive_link
-https://drive.google.com/file/d/1IIwxZnGlqrXLUXqG6yMO0r7uhCvhpk9e/view?usp=drive_link
-https://drive.google.com/file/d/1vE7XPyaFcXP4DtTY5Y9WKIt7zWgmX-Cr/view?usp=drive_link
-https://drive.google.com/file/d/1j-bIV09gr21RC3-x1N_pK4RPLV3fmWKz/view?usp=drive_link
-https://drive.google.com/file/d/1t3nW1rD3S-EL0Oymb5U7ZAj5UMkydkln/view?usp=drive_link
-https://drive.google.com/file/d/14hbfHCdMKtJZ41F9CQReMec2jeRFTOqR/view?usp=drive_link
-https://drive.google.com/file/d/1x-hUyOSne5BW0AzQ3W6_Pf4g5yXQWi9M/view?usp=drive_link
-https://drive.google.com/file/d/1sw9JqRg6E-3P84I3ZhzTrJMu0vuiaMmP/view?usp=drive_link
-https://drive.google.com/file/d/1LuqhQlL4MGZhB_6THmkovRxrlP26BbdC/view?usp=drive_link
-https://drive.google.com/file/d/15C5K6v_lkjnMSmUvVyqHQKwh2N166e7K/view?usp=drive_link
-https://drive.google.com/file/d/1ns_9eSsQeeoZ10nlbkLy8tu0GmJFSnkt/view?usp=drive_link
-https://drive.google.com/file/d/1NpzWJeK6CqjxzjIMYe6aYdX8xGsQwD4o/view?usp=drive_link
-https://drive.google.com/file/d/1NMLezwufKJ9_8xTc9KQThSzVVD71B9Ui/view?usp=drive_link
-https://drive.google.com/file/d/1aa71DCUqs6oXlIxX35jgsmsgm-NlDxPV/view?usp=drive_link
-https://drive.google.com/file/d/1UJzkIZzAL0j-D5YQBnoq7mHvttASy12O/view?usp=drive_link
-https://drive.google.com/file/d/1nPgx36HIJFb7oI94VbRzWjpPP2GANxzG/view?usp=drive_link
-https://drive.google.com/file/d/1NovAP-KVJjqcuvWy3d6G4ptGGAIDqcCx/view?usp=drive_link

lerobot/common/datasets/push_dataset_to_hub/_aloha_raw_urls/mobile_chair.txt

@@ -1,55 +0,0 @@
-https://drive.google.com/file/d/11M3Ye0r5agMaaicPbVGD0q2Hb3rGklbb/view?usp=drive_link
-https://drive.google.com/file/d/1-tx7SvYYgSvXCvnf_EI2OVdwK-CkFY6S/view?usp=drive_link
-https://drive.google.com/file/d/1EWJunmOpMHaU1hE106wwpbkGYcjQXYAF/view?usp=drive_link
-https://drive.google.com/file/d/1IDn95Z7FSiCckrSENtGV4u3RyFHNQSDY/view?usp=drive_link
-https://drive.google.com/file/d/1CwzvWj1i7QOtqrZvsCZ6BdZaKNDfpN32/view?usp=drive_link
-https://drive.google.com/file/d/1HvAvlhm77nAD3Td24QPSeq8lw-Rl_aOh/view?usp=drive_link
-https://drive.google.com/file/d/1t-suKYOPhXH666RpAYNRp2QU_DOy3AeM/view?usp=drive_link
-https://drive.google.com/file/d/18xpKgWh7RWyjMN5PkLTOo-AxsAadAuRw/view?usp=drive_link
-https://drive.google.com/file/d/1oci5Eto-ztv-AQNz8EnwZveBIhxvk-xJ/view?usp=drive_link
-https://drive.google.com/file/d/1Y-t_4vxdE6NpHO0DLJR8f3mD0Q-Wj5-c/view?usp=drive_link
-https://drive.google.com/file/d/1lylRqbbbB8bgtpsBWMPACmHJreuKmllv/view?usp=drive_link
-https://drive.google.com/file/d/1yliSyMig_NXShWfQx6qyW7Ijf2Y5lFK6/view?usp=drive_link
-https://drive.google.com/file/d/1XXhwJsJbeb7KXAooGvJapnm9bjnGUmxS/view?usp=drive_link
-https://drive.google.com/file/d/1_xs1f3hW2JArKyvfF7UWubWjyROGTLs6/view?usp=drive_link
-https://drive.google.com/file/d/1WVEHpr6EqKCZbkHapQSTXJq4xE4SWFT-/view?usp=drive_link
-https://drive.google.com/file/d/1RqOHv9pEQGvW8NUA7ynffFmG999TL_Az/view?usp=drive_link
-https://drive.google.com/file/d/1cu5AgD2gh-uA3PFJmzxxzNaF3qOSlYY1/view?usp=drive_link
-https://drive.google.com/file/d/1SsrXqiPclNrnYToPZ9Uq-k3y0C4qdHT1/view?usp=drive_link
-https://drive.google.com/file/d/1-J7EXf0vjkLIfSqT8ICEsP6CTjzSLBop/view?usp=drive_link
-https://drive.google.com/file/d/11O7ewUmoZXfyyKjy_6B5RW4DpjICxqBT/view?usp=drive_link
-https://drive.google.com/file/d/1iic44kZoCsjNsfAz2cMstZ9-WQvAhblF/view?usp=drive_link
-https://drive.google.com/file/d/1yLV1lVX-2WnWQldGlnQZ0x7QBuDiVkL3/view?usp=drive_link
-https://drive.google.com/file/d/1Tybp9ru98TTbGn4eyROpUQwDFuALWXmk/view?usp=drive_link
-https://drive.google.com/file/d/13E9OTMiipVJByDs5-J19oWwAz7l94LTN/view?usp=drive_link
-https://drive.google.com/file/d/1EeTpJQdMSliw4JzSMtJ6CyTvVdexjM4M/view?usp=drive_link
-https://drive.google.com/file/d/1NHyNwoFqzeAu-1_PSpq5JfxaiD_xbpn9/view?usp=drive_link
-https://drive.google.com/file/d/1fJcS0phDp4xm_FyGaJ5wr9Pe4KqtHaxD/view?usp=drive_link
-https://drive.google.com/file/d/12AqrLUaewDPEcFRqPZeZFb_TQ0Lfi3At/view?usp=drive_link
-https://drive.google.com/file/d/1x_hd4Qsq1oJS-aj2t3qM7WbbV7KZj05b/view?usp=drive_link
-https://drive.google.com/file/d/14OUSUArmsB068hs6BuEIXQhI1Cyz8Sf0/view?usp=drive_link
-https://drive.google.com/file/d/16zlzh1T5zeUJQnFf382NXkFEKEnDub4O/view?usp=drive_link
-https://drive.google.com/file/d/1IbDltmN-NEFCNtr1TO4ILxEgQ94rtjWv/view?usp=drive_link
-https://drive.google.com/file/d/15gmlf8Gx9455pZ1AlqcCSwh3nDPxMzSr/view?usp=drive_link
-https://drive.google.com/file/d/1qHpRL1oZfIMo_vxnm8qfwQ-7l0BZIVva/view?usp=drive_link
-https://drive.google.com/file/d/1H1xskIgiFZivkYn23rMzH3xePGOh3VTC/view?usp=drive_link
-https://drive.google.com/file/d/1avls6Pv0kYiCMNVknbc1zQsgy64MUDMM/view?usp=drive_link
-https://drive.google.com/file/d/1MmWVgCj5khc8KMIifmt3EzF1o-CtPyyn/view?usp=drive_link
-https://drive.google.com/file/d/1U0kCc_xqW0WNppf4sbnK14euWKdPZtzB/view?usp=drive_link
-https://drive.google.com/file/d/16CaEyQscOuhLj23PEGDTL9DeyNkohkMn/view?usp=drive_link
-https://drive.google.com/file/d/1Iu8uM6UUJ0zW8tvN-9UiOe_4oSNzEutg/view?usp=drive_link
-https://drive.google.com/file/d/1UImqiBaIxCR-1DNJaZhHqeHhaySOtVIr/view?usp=drive_link
-https://drive.google.com/file/d/1VpU2V_leIoRIyv_lAvE7eLHBG8DxCTnp/view?usp=drive_link
-https://drive.google.com/file/d/1_Q8J27OT3Xby7QY6yHvIJauFRWEMxkRm/view?usp=drive_link
-https://drive.google.com/file/d/1bantmVo1L9Xz4tbiNw_a1UC2Z_HPO1wT/view?usp=drive_link
-https://drive.google.com/file/d/1IRIXMJMCBDkBjbaHvAlEiBogSvZ1jK_3/view?usp=drive_link
-https://drive.google.com/file/d/1mAHXKjiFbjwydypW2t5Lv8_H5x6nHegl/view?usp=drive_link
-https://drive.google.com/file/d/1SfyY796fLrBCMY39OcyuxZafqSCRZPZk/view?usp=drive_link
-https://drive.google.com/file/d/1X-44sZ8CcfzIskc0dvSx882o1yFhHaZB/view?usp=drive_link
-https://drive.google.com/file/d/1BOIWCCCk6DLD4Bmvc75ZbbLi9AQm-1ao/view?usp=drive_link
-https://drive.google.com/file/d/1RuyDtRE1kk76sw-wP8vx5SgLoPF3PA_H/view?usp=drive_link
-https://drive.google.com/file/d/1c4eoQiBbGuy3CTAQDUSkd84Ponh1roAQ/view?usp=drive_link
-https://drive.google.com/file/d/19PXB9z4Ljq6dsbf9TqcOrrP5SRbw2Tc_/view?usp=drive_link
-https://drive.google.com/file/d/1nn1VVZVoIXWdYDozR7XHXE4mPLQG80PQ/view?usp=drive_link
-https://drive.google.com/file/d/1MBdFGOKPV8GUhwoSsJ_Ky3qAMLM2Bv3K/view?usp=drive_link
-https://drive.google.com/file/d/1of3k_M-7Nh3I1TndcWedxK4ca9dn8Sc5/view?usp=drive_link

lerobot/common/datasets/push_dataset_to_hub/_aloha_raw_urls/mobile_elevator.txt

@@ -1,20 +0,0 @@
-https://drive.google.com/file/d/12ctkOAdkCNGN1JLbZb5ww3XTBn2LFpGI/view?usp=drive_link
-https://drive.google.com/file/d/1G_Vd46_4fq6O64gHHjUbJX5Ld44ZZx0y/view?usp=drive_link
-https://drive.google.com/file/d/1uKgUy73B3xBogQAOUhfZjO0X5qZGsi2c/view?usp=drive_link
-https://drive.google.com/file/d/1fu9cIrfI-fE2LhdGUxbx7-8Ci_PF8Ypm/view?usp=drive_link
-https://drive.google.com/file/d/1Ygk9ZPJzx8xw2A9JF3NHbJ44TqnvSTQR/view?usp=drive_link
-https://drive.google.com/file/d/18m5xPuccNsEB20WPshm3zhxmXc6k63ED/view?usp=drive_link
-https://drive.google.com/file/d/1DiqqxC44rriviRQpqogcv0-EB-Y6nr9g/view?usp=drive_link
-https://drive.google.com/file/d/1qPdaoTVDizJXkfXLioWU7iJ8hqCXSyOQ/view?usp=drive_link
-https://drive.google.com/file/d/1Fj9kIA_mG7f67WFfACJEaZ7izcHG7vUm/view?usp=drive_link
-https://drive.google.com/file/d/1WpYehZnI2P7dUdJPfkE-ij1rqCnjZEbB/view?usp=drive_link
-https://drive.google.com/file/d/1_zwWkT4jPyzB38STWb6whlzsPzXmfA9r/view?usp=drive_link
-https://drive.google.com/file/d/1U6-J4I_fPlSFFGfhZPxS5_YzKXwXIZYp/view?usp=drive_link
-https://drive.google.com/file/d/1pRhxxcTfZp5tQo_EScvJUwfc3amiS6Vk/view?usp=drive_link
-https://drive.google.com/file/d/1lWLntqra83RlYU_gN7Vostnfydf6gutd/view?usp=drive_link
-https://drive.google.com/file/d/1vIBKo0x-NYEHV1FvRpco1lQMpRdAWAIL/view?usp=drive_link
-https://drive.google.com/file/d/1pdrLV3JTQou_XH0Aap61Ssf60iVKm1jJ/view?usp=drive_link
-https://drive.google.com/file/d/1QTsLoQ7SwmKdQHjBGVDaR2uTwfFwtrOf/view?usp=drive_link
-https://drive.google.com/file/d/1Gytai8M_12J36GY6L_TulEcOC-035jwS/view?usp=drive_link
-https://drive.google.com/file/d/14LJudNc629NT-i8xreXtzl27ce_DxOFJ/view?usp=drive_link
-https://drive.google.com/file/d/1sBvPCODbzxGAI0S3lgN5cSG9Go3lRi00/view?usp=drive_link

lerobot/common/datasets/push_dataset_to_hub/_aloha_raw_urls/mobile_shrimp.txt

@@ -1,18 +0,0 @@
-https://drive.google.com/file/d/1MJn9GbC8p9lN4gC9KDMLEkTkP_gGpXj0/view?usp=drive_link
-https://drive.google.com/file/d/1-4LXgjl7ZCOgp-8GCJmFRD8OeqN5Jf7-/view?usp=drive_link
-https://drive.google.com/file/d/1Ho06Ce0SPbqU3juaMxNUwAt3zCRLGC8W/view?usp=drive_link
-https://drive.google.com/file/d/1ivHoj7_7olBSxH-Y8kqXEW7ttITK-45j/view?usp=drive_link
-https://drive.google.com/file/d/1qjY4hM_IvZ8cq2II_n9MeJbvyeuN4oBP/view?usp=drive_link
-https://drive.google.com/file/d/1rKVhO_f92-7sw13T8hTVrza3B9oAVgoy/view?usp=drive_link
-https://drive.google.com/file/d/1pcLPHO8fBkc1-CRa88tyQtEueE4xiXNi/view?usp=drive_link
-https://drive.google.com/file/d/1Vev_chCsIeEdvQ8poEYNsOJFGy_QU8kZ/view?usp=drive_link
-https://drive.google.com/file/d/1l5G4zpRkxSLCQjvGPYSN4zfCvVRQuzMz/view?usp=drive_link
-https://drive.google.com/file/d/14vgthE1eoakXkr2-DRw50E6lAqYOiUuE/view?usp=drive_link
-https://drive.google.com/file/d/17nPSmKKmgQ2B7zkzWrZYiLM3RBuFod82/view?usp=drive_link
-https://drive.google.com/file/d/1QcDsxplVvb_ID9BVrihl5FvlC-j7waXi/view?usp=drive_link
-https://drive.google.com/file/d/18pEejBpI-eEVaWAAjBCyC0vgbX3T1Esj/view?usp=drive_link
-https://drive.google.com/file/d/1H8eH6_IRODtEFT6WoM77ltR5OoOrqXmI/view?usp=drive_link
-https://drive.google.com/file/d/1IWlpFRZhoxyG4nS13CWK4leZVk5wbNx4/view?usp=drive_link
-https://drive.google.com/file/d/1PbZA8_OCGmMLxNP9xbkLRSChniL4uGxl/view?usp=drive_link
-https://drive.google.com/file/d/1p9XAdmG2f_WeflNO4DIJ_tr1rK6M9B4B/view?usp=drive_link
-https://drive.google.com/file/d/1nS59Et1cNAvKo3Y4SeSGRuZD5TvBbCF3/view?usp=drive_link

lerobot/common/datasets/push_dataset_to_hub/_aloha_raw_urls/mobile_wash_pan.txt

@@ -1 +0,0 @@
-https://drive.google.com/drive/folders/1S8eFg98IaGAIKVZ8QFWG1bx4mHa-O204

lerobot/common/datasets/push_dataset_to_hub/_aloha_raw_urls/mobile_wipe_wine.txt

@@ -1,4 +0,0 @@
-https://drive.google.com/drive/folders/1tC_g1AJ8lglBLY-fjsQrG6DMBa3Ucp-0
-https://drive.google.com/file/d/1fG_Yi2MJrFjiUVN3XoiWXLtTxHlwwaDv/view?usp=drive_link
-https://drive.google.com/file/d/1WX32VWfzzX3Blmd06DRxLwFbMJfVe7P4/view?usp=drive_link
-https://drive.google.com/file/d/18onsX3vXg3xkFwP5bVUCjdV4n9TRn0C9/view?usp=drive_link

lerobot/common/datasets/push_dataset_to_hub/_aloha_raw_urls/sim_insertion_human.txt

@@ -1,3 +0,0 @@
-https://drive.google.com/drive/folders/1RgyD0JgTX30H4IM5XZn8I3zSV_mr8pyF
-https://drive.google.com/file/d/18Cudl6nikDtgRolea7je8iF_gGKzynOP/view?usp=drive_link
-https://drive.google.com/file/d/1C1kZYyROzs-PrLc0SkDgUgMi4-L3lauE/view?usp=drive_link

lerobot/common/datasets/push_dataset_to_hub/_aloha_raw_urls/sim_insertion_scripted.txt

@@ -1,3 +0,0 @@
-https://drive.google.com/drive/folders/1TsojQQSXtHEoGnqgJ3gmpPQR2DPLtS2N
-https://drive.google.com/file/d/1wfMSZ24oOh5KR_0aaP3Cnu_c4ZCveduB/view?usp=drive_link
-https://drive.google.com/file/d/17EuCUWS6uCCr6yyNzpXdcdE-_TTNCKtf/view?usp=drive_link

lerobot/common/datasets/push_dataset_to_hub/_aloha_raw_urls/sim_transfer_cube_human.txt

@@ -1,3 +0,0 @@
-https://drive.google.com/drive/folders/1sc-E4QYW7A0o23m1u2VWNGVq5smAsfCo
-https://drive.google.com/file/d/18smMymtr8tIxaNUQ61gW6dG50pt3MvGq/view?usp=drive_link
-https://drive.google.com/file/d/1Nk7l53d9sJoGDBKAOnNrExX5nLacATc6/view?usp=drive_link

lerobot/common/datasets/push_dataset_to_hub/_aloha_raw_urls/sim_transfer_cube_scripted.txt

@@ -1,3 +0,0 @@
-https://drive.google.com/drive/folders/1aRyoOhQwxhyt1J8XgEig4s6kzaw__LXj
-https://drive.google.com/file/d/1pnGIOd-E4-rhz2P3VxpknMKRZCoKt6eI/view?usp=drive_link
-https://drive.google.com/file/d/1GKReZHrXU73NMiC5zKCq_UtqPVtYq8eo/view?usp=drive_link

lerobot/common/datasets/push_dataset_to_hub/_aloha_raw_urls/static_battery.txt

@@ -1,2 +0,0 @@
-https://drive.google.com/drive/folders/19qS_n7vKgDcPeTMnvDHQ5-n73xEbJz5D
-https://drive.google.com/file/d/1oC31By0A2bsBeHyUwBdQw1z4ng6yi9Za/view?usp=drive_link

lerobot/common/datasets/push_dataset_to_hub/_aloha_raw_urls/static_candy.txt

@@ -1,2 +0,0 @@
-https://drive.google.com/drive/folders/1m5rQ6UVH8Q9RQp_6c0CxkQ88-L-ScO7q
-https://drive.google.com/file/d/1wHz2qcmwcVG0C0CZ9MjQDQcmj4OY9_a3/view?usp=drive_link

lerobot/common/datasets/push_dataset_to_hub/_aloha_raw_urls/static_coffee.txt

@@ -1,2 +0,0 @@
-https://drive.google.com/drive/folders/1seQGay470nGQ-knBI5TjsTr8iL9Qws5q
-https://drive.google.com/file/d/1T89hSX5U99wLGvGTE7yUBaQPOpyj6Sai/view?usp=drive_link

lerobot/common/datasets/push_dataset_to_hub/_aloha_raw_urls/static_coffee_new.txt

@@ -1,2 +0,0 @@
-https://drive.google.com/drive/folders/1t3eDc5Rg0DveyRe8oTm6Dia_FYU5mXyf
-https://drive.google.com/file/d/1TXFaduTakvS0ZWJqKCX-HIvYglum_5CY/view?usp=drive_link

lerobot/common/datasets/push_dataset_to_hub/_aloha_raw_urls/static_cups_open.txt

@@ -1,2 +0,0 @@
-https://drive.google.com/drive/folders/1Z9X3DNzd6LS0FFjQemNUMoMA5yk5VQOh
-https://drive.google.com/file/d/1Wlyc0vTkjXuWB6zbaVOWhEfD7BmPgUV_/view?usp=drive_link

lerobot/common/datasets/push_dataset_to_hub/_aloha_raw_urls/static_fork_pick_up.txt

@@ -1,53 +0,0 @@
-https://drive.google.com/drive/folders/1DYgB4ifX4uIid9m9jnC0Zdz8Nf7ZC0fc
-https://drive.google.com/file/d/1Eb-NRNk_FmVleCbU_Ng5Y4dfcjTKN7Rv/view?usp=drive_link
-https://drive.google.com/file/d/1dkhjEADakT-44l9jf-nK4x89kr4yG_qb/view?usp=drive_link
-https://drive.google.com/file/d/14hDhgcZkVqNExGb4tIXpSjMshhqZETch/view?usp=drive_link
-https://drive.google.com/file/d/1zVMEHpHbuNyP5A_lYU7RPSLB-4V0yfZw/view?usp=drive_link
-https://drive.google.com/file/d/1JtgDjBvy7FnRpFzrx_foC3quorYQFAR-/view?usp=drive_link
-https://drive.google.com/file/d/1EHdneB6F-PP0dQlX8qPaXbxmKoBy_YwO/view?usp=drive_link
-https://drive.google.com/file/d/17Z0jjVBy1OPKREPu77_n_rQzorDiapji/view?usp=drive_link
-https://drive.google.com/file/d/1F4i23qPJ_qTf5jWjfLo4ARGJChznYWt3/view?usp=drive_link
-https://drive.google.com/file/d/1kZtXWM3uS0-rLblydBfJ0mMcVnMMXw9w/view?usp=drive_link
-https://drive.google.com/file/d/1mNODox87xFfY5Z_o5mcLsr8SHb39jDik/view?usp=drive_link
-https://drive.google.com/file/d/1Ob44VdmEUA93FKDECiRb5Ogz2xQg5IWp/view?usp=drive_link
-https://drive.google.com/file/d/1fdQLdjj3Cwv33R1wZhfrLz9Del8mqgHb/view?usp=drive_link
-https://drive.google.com/file/d/1Yu3L3ft21zP__XL8pCfhb788ZleuW1n5/view?usp=drive_link
-https://drive.google.com/file/d/1ozBBWXVZ9hXDh9ooHUNroHdYm8UDqnhJ/view?usp=drive_link
-https://drive.google.com/file/d/1o0TGqvfWw_Lunxb5ubKDS21Lr_WC0h75/view?usp=drive_link
-https://drive.google.com/file/d/1jZnd5eP5L6BH5l98BPN6OnoQx3fu8e9n/view?usp=drive_link
-https://drive.google.com/file/d/1S5sYbz8wcLYp0V67v13i4PRcBxodn4Hg/view?usp=drive_link
-https://drive.google.com/file/d/1rFeg_x6ftJYwPtBv34D3h2L2cpDLeR4G/view?usp=drive_link
-https://drive.google.com/file/d/1GvS3lcm4o6nm_scUk0XxKeVFNmzjucDZ/view?usp=drive_link
-https://drive.google.com/file/d/1-9i0riphC7NhhDahcQfD1QoBXP5gF90A/view?usp=drive_link
-https://drive.google.com/file/d/15p_IqGsMbKuvzMS872THAZr-3SBtb1Fr/view?usp=drive_link
-https://drive.google.com/file/d/1ToyYcBfJL8gbQn0q_59zPLsFmm7dmMJo/view?usp=drive_link
-https://drive.google.com/file/d/1e_7PNH7CYafE4pAebP7ZdI7XFbmEcy_i/view?usp=drive_link
-https://drive.google.com/file/d/1JoabvGVsIQdug2xOhUIhetEIyDM91y_Y/view?usp=drive_link
-https://drive.google.com/file/d/1kOMw1y0lmnVaCjwZICfzCsx6e0Z8MNGR/view?usp=drive_link
-https://drive.google.com/file/d/16it_wd1JOevUQTK2_CvF_pBACTgpIPgM/view?usp=drive_link
-https://drive.google.com/file/d/1IRcCj9HnJSfbyMgr5XEERGlEnWeZQwOc/view?usp=drive_link
-https://drive.google.com/file/d/1Z2dIJfq_S3liGmPN9Rphvkmucnmw7tlb/view?usp=drive_link
-https://drive.google.com/file/d/1J3NoAjzndGx9yNyaBOJHdNny1epzUoBt/view?usp=drive_link
-https://drive.google.com/file/d/18nOvxV1k8FSmBrhT4TPo2sKKSZXougyx/view?usp=drive_link
-https://drive.google.com/file/d/1CT8FxclafFMjSd7gCWVw3VSeryeiF04i/view?usp=drive_link
-https://drive.google.com/file/d/16M9KVqQMFfSsXfypK0bocFft8Nz3j2Rt/view?usp=drive_link
-https://drive.google.com/file/d/18QPVkw6bj6HW8LTPrQLWrrUX4R6RcF42/view?usp=drive_link
-https://drive.google.com/file/d/1hQTVtA5hBTE_StXpJafTZJ3tgt2VQQ_t/view?usp=drive_link
-https://drive.google.com/file/d/1Dn-d5g69H6EgAWgsFdrcbJKtz7ySsCQ8/view?usp=drive_link
-https://drive.google.com/file/d/13hMr16483P7ALYv73yMRUN37fJdVQM62/view?usp=drive_link
-https://drive.google.com/file/d/1848yN3XMN5zJMEgApt6KzrWgfRPfimtv/view?usp=drive_link
-https://drive.google.com/file/d/1oAD9kSnS0fTgj-CjD4u9VdZ5X67IOIMa/view?usp=drive_link
-https://drive.google.com/file/d/1ilzIWLCCG5b_KgF5s0wdN2I5-lFNpwC1/view?usp=drive_link
-https://drive.google.com/file/d/1rjsT2YBjnidxod1s9s-myAYz8boHr-WB/view?usp=drive_link
-https://drive.google.com/file/d/18Gg48HTub15bd8qzbhiCUufbVy0fbN5G/view?usp=drive_link
-https://drive.google.com/file/d/1WsSnQSqmMTVSRwrhT1Y-v782My2zcjLm/view?usp=drive_link
-https://drive.google.com/file/d/1ea9ZCvoyc-xqiFXgeDcA_mOWsw7VUuoi/view?usp=drive_link
-https://drive.google.com/file/d/1wv1v3-XhPgbNzp62BXbJTDzMPu2tlDUc/view?usp=drive_link
-https://drive.google.com/file/d/18-ikzt8LoZ83Gi3goKCELs4U4z8hrRoF/view?usp=drive_link
-https://drive.google.com/file/d/16Bjhp7JNCXkGuLvyNcZowAx3W-Y-15DV/view?usp=drive_link
-https://drive.google.com/file/d/1Gc-KRI-xwcp1fMR55ugbrLg_5y3SPde-/view?usp=drive_link
-https://drive.google.com/file/d/1oP72Q386Z4Sy5MMm-t5yNogIe5Van_9k/view?usp=drive_link
-https://drive.google.com/file/d/112T90eDUDVH-SyOV7UnZl5bscAH2hcfq/view?usp=drive_link
-https://drive.google.com/file/d/1y-uKOesRRhjgDtFbG_j65f4SGg0v8XDg/view?usp=drive_link
-https://drive.google.com/file/d/1LOP05OagoI3km-ZKQBrS204A85UVk7Ok/view?usp=drive_link
-https://drive.google.com/file/d/1QkHQKgasVzWsmdPvkXgGhWyQ84d93_Az/view?usp=drive_link

lerobot/common/datasets/push_dataset_to_hub/_aloha_raw_urls/static_pingpong_test.txt

@@ -1 +0,0 @@
-https://drive.google.com/drive/folders/1Ut2cv6o6Pkfgg46DgwVUM7Z5PkNG8eJ-

lerobot/common/datasets/push_dataset_to_hub/_aloha_raw_urls/static_pro_pencil.txt

@@ -1 +0,0 @@
-https://drive.google.com/drive/folders/1FqxPV0PgvgIu8XFjtvZSPSExuNcxVVAY

lerobot/common/datasets/push_dataset_to_hub/_aloha_raw_urls/static_screw_driver.txt

@@ -1,2 +0,0 @@
-https://drive.google.com/drive/folders/1SKtG0ct9q0nVdYssJNMWSOjikcXliT58
-https://drive.google.com/file/d/1nchD21O30B3i3LDoqramo1zgW5YvpJIN/view?usp=drive_link

lerobot/common/datasets/push_dataset_to_hub/_aloha_raw_urls/static_tape.txt

@@ -1,2 +0,0 @@
-https://drive.google.com/drive/folders/1_4DHf2cma0xsChLQFghwigX6Ukti5-zQ
-https://drive.google.com/file/d/1_8vS4hDNDgUQY-SmekrNaa7dF67QJYU-/view?usp=drive_link

lerobot/common/datasets/push_dataset_to_hub/_aloha_raw_urls/static_thread_velcro.txt

@@ -1,2 +0,0 @@
-https://drive.google.com/drive/folders/1_4DHf2cma0xsChLQFghwigX6Ukti5-zQ
-https://drive.google.com/file/d/1_8vS4hDNDgUQY-SmekrNaa7dF67QJYU-/view?usp=drive_link

lerobot/common/datasets/push_dataset_to_hub/_aloha_raw_urls/static_towel.txt

@@ -1,2 +0,0 @@
-https://drive.google.com/drive/folders/1fAD7vkyTGTFB_nGXIKofCU1U05oE3MFv
-https://drive.google.com/file/d/1XzyQ2B6LLvcurIonOpEu4nij2qwNWshH/view?usp=drive_link

lerobot/common/datasets/push_dataset_to_hub/_aloha_raw_urls/static_vinh_cup.txt

@@ -1,53 +0,0 @@
-https://drive.google.com/drive/folders/13EQsVsnxT86K20QAoyE_YpsFbQ7fZQdu
-https://drive.google.com/file/d/1-W_JHghZG65FNTVhw1SXhtQrazdLL3Ue/view?usp=drive_link
-https://drive.google.com/file/d/1VwRJgdWUo-2nQaNM7Bs77-fsm8iwUxEo/view?usp=drive_link
-https://drive.google.com/file/d/1wFzGRo5iYA13WLi6IV1ry64RyahQBFio/view?usp=drive_link
-https://drive.google.com/file/d/1IKtQzQ-n-UTv64hYpReu2R4cqUvmNQqD/view?usp=drive_link
-https://drive.google.com/file/d/1GicVci9OiuuZZH79i5Mg7AtWod94MzwT/view?usp=drive_link
-https://drive.google.com/file/d/1JVnIoR7EIQp70T4eAf9RX65JcTrzsjQc/view?usp=drive_link
-https://drive.google.com/file/d/1W2xr4h23ucjPrc-mBEeqnACsfaImpc0p/view?usp=drive_link
-https://drive.google.com/file/d/10xj_0V7A07o3uCa7v5omUrTC0YlPW8H3/view?usp=drive_link
-https://drive.google.com/file/d/1FOc3EMaCy8Mb0_a7PuXLAwKwvxkbKmwU/view?usp=drive_link
-https://drive.google.com/file/d/143PgDXBcf2GQ0Q07ZPMVMfBgZDd5sLJG/view?usp=drive_link
-https://drive.google.com/file/d/1pE5Tyj0LlGbGWvUzuhixp86Ibu55Ez3I/view?usp=drive_link
-https://drive.google.com/file/d/141668b1VzX80ncrVJPzhkoAeIFB4MEK9/view?usp=drive_link
-https://drive.google.com/file/d/1bw12lo37p1ZvRvErHsll7cEYi2OxscvZ/view?usp=drive_link
-https://drive.google.com/file/d/1zfnMFvbgBjl6SzYhksbaOzfbwLrCN6tb/view?usp=drive_link
-https://drive.google.com/file/d/1-GIszA6mUJMaNB-tdh9r9skc77SWA0VX/view?usp=drive_link
-https://drive.google.com/file/d/1fTB0zWFYU6zh4IIUFT2zX_OkwYqmElwY/view?usp=drive_link
-https://drive.google.com/file/d/1gPIPNKGmrO9c7gKF7SP0SuUYbIBBq8z1/view?usp=drive_link
-https://drive.google.com/file/d/12JeJ-dQd5lYyn6PlDOGdE-ChVeiZ-Uv0/view?usp=drive_link
-https://drive.google.com/file/d/100_20cgCqerU6qoh3TfTbwLy9mlDAFEG/view?usp=drive_link
-https://drive.google.com/file/d/111oAGJ76ku_pYgbBoIdZAC1_XEQcPI__/view?usp=drive_link
-https://drive.google.com/file/d/1UhC8L-354ZQ2gblPFGI35EMsVwfpuKa0/view?usp=drive_link
-https://drive.google.com/file/d/1sIXQSgUR_xdrNtGrL6QGBnkLMKErsIp1/view?usp=drive_link
-https://drive.google.com/file/d/16Ax77bDSIXnsn4GFL8XYKKT1P6bPpfMd/view?usp=drive_link
-https://drive.google.com/file/d/1pgRVYwwVIsWq_qsWqZpe1UBzZfF5Fa9D/view?usp=drive_link
-https://drive.google.com/file/d/1jtimaZkWsY1P5gC2bbS64H_WCUU7HXN2/view?usp=drive_link
-https://drive.google.com/file/d/1N6Bh02P-RiTEgtx1YH1Db_X3TGpP-X_r/view?usp=drive_link
-https://drive.google.com/file/d/14Fy8EwJ8d9Vh97Yt1VOvUChSCrfIjBij/view?usp=drive_link
-https://drive.google.com/file/d/1IRuv42dvIMPuKhcMZmuXaBjJ-lPFOmQd/view?usp=drive_link
-https://drive.google.com/file/d/16XWzNY2D8ucVVn5geBgsVdhm3ppO4que/view?usp=drive_link
-https://drive.google.com/file/d/1xsVOoQgthK_L_SDrmq_JvQgUpAvPEAY8/view?usp=drive_link
-https://drive.google.com/file/d/1bZbw66DyEMvnJnzkdUUNbKjvNKg8KFYM/view?usp=drive_link
-https://drive.google.com/file/d/1CyTVkdrNGGpouCXr4CfhKbMzE6Ah3oo3/view?usp=drive_link
-https://drive.google.com/file/d/1hDRyeM-XEDpHXpptbT8LvNnlQUR3PWOh/view?usp=drive_link
-https://drive.google.com/file/d/1XhHWxbra8Iy5irQZ83IvxwaJqHq9x4s1/view?usp=drive_link
-https://drive.google.com/file/d/1haZcn6aM1o4JlmP9tJj3x2enrxiPaDSD/view?usp=drive_link
-https://drive.google.com/file/d/1ypDyuUTbljaBZ34f-t7lj3O_0bRmyX2n/view?usp=drive_link
-https://drive.google.com/file/d/1ILEEZo_tA9_ChIAprr2mPaNVKZi5vXsO/view?usp=drive_link
-https://drive.google.com/file/d/1U7nVYFaGE8vVTfLCW33D74xOjDcqfgyJ/view?usp=drive_link
-https://drive.google.com/file/d/1rZ93_rmCov5SMDxPkfM3qthcRELZrQX6/view?usp=drive_link
-https://drive.google.com/file/d/1mYO1b_csddtyE3qT6cwLiw-m2w2_1Lxh/view?usp=drive_link
-https://drive.google.com/file/d/1xz7Q5x2jikY8wJQjMRQpRws6AnfWlHm5/view?usp=drive_link
-https://drive.google.com/file/d/1OO8GaO-0FrSZRd1kxMYwBmubyiLOWnbl/view?usp=drive_link
-https://drive.google.com/file/d/1EXn4NVDmf-4_HCy34mYwT-vwK2CFI9ev/view?usp=drive_link
-https://drive.google.com/file/d/10hH70XhXRL9C5SnAG4toHtfHqfJUJo4H/view?usp=drive_link
-https://drive.google.com/file/d/18tiBcxea0guUai4lwsXQvt0q2LZ8ZnnJ/view?usp=drive_link
-https://drive.google.com/file/d/1Q8R8qv37vk5PQ5kQ2ibx6BFLOySD0VpX/view?usp=drive_link
-https://drive.google.com/file/d/17aNriHzjhdibCyuUjQoMFZqjybJZtggG/view?usp=drive_link
-https://drive.google.com/file/d/1LVjEYHSdeKm6CotU1QguIeNEPaIaFl_1/view?usp=drive_link
-https://drive.google.com/file/d/1ufAhE_EkgJ85slg2EW8aW_grOzE_Lmxd/view?usp=drive_link
-https://drive.google.com/file/d/1wtzLtXrkw9eXRGESTPIOlpl1tInu-b2m/view?usp=drive_link
-https://drive.google.com/file/d/1Mk5qvVtD_QHwGOUApRq76TUw2T5THu6f/view?usp=drive_link
-https://drive.google.com/file/d/1y1WQ3hboWVJ68KEYQQ3OhreGuaUpSgwc/view?usp=drive_link

lerobot/common/datasets/push_dataset_to_hub/_aloha_raw_urls/static_vinh_cup_left.txt

@@ -1,52 +0,0 @@
-https://drive.google.com/drive/folders/1dxWh6YFZUDt6qXIoxgD9bla3CiFjZ11C
-https://drive.google.com/file/d/1hNBJN00SCAlOl0ZEgm7RRGbAGDjyBs0p/view?usp=drive_link
-https://drive.google.com/file/d/17He0CVwXGeoMmXg4SHKo-osNn7YPKVL7/view?usp=drive_link
-https://drive.google.com/file/d/1laNKUVID1x2CV6a2O2WQjwFewKu4lidL/view?usp=drive_link
-https://drive.google.com/file/d/1pNf36xbZJGRArYLmNAvRj5y6CoqdC6kB/view?usp=drive_link
-https://drive.google.com/file/d/1_4E1-y3JXk5I0ebycLYM70YDPK9g52gZ/view?usp=drive_link
-https://drive.google.com/file/d/1PHfzhGPdbolKyOpS3FnR2w7Q8zUlJXSk/view?usp=drive_link
-https://drive.google.com/file/d/17ls2PPN-Pi3tEuK059cwV2_iDT8aGhOO/view?usp=drive_link
-https://drive.google.com/file/d/1LWsg6PmCT00Kv_N_slrmcwKmQPGoBT3k/view?usp=drive_link
-https://drive.google.com/file/d/12LckrchoHTUVH7rxi8J7zD9dA19GXvoW/view?usp=drive_link
-https://drive.google.com/file/d/1VqrJKjAIkj5gtFXL69grdSeu9CyaqnSw/view?usp=drive_link
-https://drive.google.com/file/d/1g5rQYDBZvW-kUtYPeyF3qmd53v6k7kXu/view?usp=drive_link
-https://drive.google.com/file/d/10kUgaSJ0TS7teaG83G3Rf_DG4XGrBt6A/view?usp=drive_link
-https://drive.google.com/file/d/1je9XmneZQZvTma5adMJICUPDovW3ppei/view?usp=drive_link
-https://drive.google.com/file/d/1v28r6bedwZGbUPVVTVImXhK-42XdtGfj/view?usp=drive_link
-https://drive.google.com/file/d/1-TEEx9sGVvzMMaNXYfQMtY2JJ6cvl0dT/view?usp=drive_link
-https://drive.google.com/file/d/1YdBKdJFP9rJWBUX7qrOYL_gfUA8o6J9M/view?usp=drive_link
-https://drive.google.com/file/d/1X9vffwQHNUSKLXr2RlYNtbWDIFCIDfdF/view?usp=drive_link
-https://drive.google.com/file/d/11hqesqa5kvEe5FABUnZRcvmOhR373cYM/view?usp=drive_link
-https://drive.google.com/file/d/1ltTTECjEcbQPgS3UPRgMzaE2x9n6H7dC/view?usp=drive_link
-https://drive.google.com/file/d/1Zxqfa29JdwT-bfMpivi6IG2vz34d21dD/view?usp=drive_link
-https://drive.google.com/file/d/11LQlVxS5hz494dYUJ_PNRPx2NHIJbQns/view?usp=drive_link
-https://drive.google.com/file/d/1i1JhNtnZpO_E8rAv8gxBP3ZTZRvcvsZi/view?usp=drive_link
-https://drive.google.com/file/d/11jOXAr2EULUO4Qkm748634lg4UUFho5U/view?usp=drive_link
-https://drive.google.com/file/d/1rj67wur8DdB_Pipwx24bY43xu4X1eQ5e/view?usp=drive_link
-https://drive.google.com/file/d/15ZTm6lO6f_JQy_4SNfrOu3iPYn1Ro8mh/view?usp=drive_link
-https://drive.google.com/file/d/1q4gBtqWPJtCwXEvknGgN0WHGp7Vfn1b9/view?usp=drive_link
-https://drive.google.com/file/d/1t17keyre47AYqm8GgXiQ7EcvcUkeSiDQ/view?usp=drive_link
-https://drive.google.com/file/d/1OYUPGxtZgOF86Ng_BEOTXm_XOYpuQPsO/view?usp=drive_link
-https://drive.google.com/file/d/1cBjbGHi3dwWHtx6r9EQJi0JT_CE3LuHt/view?usp=drive_link
-https://drive.google.com/file/d/14qaMyF0mcbCB-fCYKNyo5_2NahSC6D5u/view?usp=drive_link
-https://drive.google.com/file/d/12FgX86eA7Y5co9ULBVK80XMsiKQSs-Ri/view?usp=drive_link
-https://drive.google.com/file/d/1yvoHWidf-jdBVw6qCCXOFfkVwKj_2hPk/view?usp=drive_link
-https://drive.google.com/file/d/1a2SugsSDlC8UtUrFzp-_KAwyZckQOvdQ/view?usp=drive_link
-https://drive.google.com/file/d/1l8pILBFSAosypWJMza2K09Vm7rug9axm/view?usp=drive_link
-https://drive.google.com/file/d/1hfPQ8dBCk97PnOhq6_MIISm3IEzcOxJG/view?usp=drive_link
-https://drive.google.com/file/d/1PPAUwlJCFKpms8cqF_k1v2_fCgDBOc3S/view?usp=drive_link
-https://drive.google.com/file/d/1lVKQZeqFfK3amEmLuFhYLUFQ2eyE8rOW/view?usp=drive_link
-https://drive.google.com/file/d/1K9iPMLfDowcIFoyzpvgn88dQ6x6kVwNG/view?usp=drive_link
-https://drive.google.com/file/d/1PNvMqG9tL7QxeLaYBGHiWYR6SYb5iIct/view?usp=drive_link
-https://drive.google.com/file/d/1xkRtzbvIkUsylx9hrFLGQsJn0h1EYu-5/view?usp=drive_link
-https://drive.google.com/file/d/1nxMRrJlSayjDIfr5CmHO1NzAw3COhsLi/view?usp=drive_link
-https://drive.google.com/file/d/1Qs3WEyMGrmagiHIkkFEueWNnJhkUeR1s/view?usp=drive_link
-https://drive.google.com/file/d/1D-G2_Q0SS3M8zyJbg_XzkF2ANPw1HTuX/view?usp=drive_link
-https://drive.google.com/file/d/1mdmJsDGO-YtJAOF_yPKl6lq4PJOIbQhT/view?usp=drive_link
-https://drive.google.com/file/d/11m9bwfop_sPmnQr_8amB6EEsrbAeG_z5/view?usp=drive_link
-https://drive.google.com/file/d/19tyYt5FMn5kru0g9o2nMJhKPnsDqkIZv/view?usp=drive_link
-https://drive.google.com/file/d/1XvTpUdsVTZ-vydvdYYmynbma--HfUGSl/view?usp=drive_link
-https://drive.google.com/file/d/1MO3hFu68J6NohTzr9aB_fY02VA6QSOqj/view?usp=drive_link
-https://drive.google.com/file/d/1Lh-UjwAk__04YOTWINF_QGVU8SjetVaY/view?usp=drive_link
-https://drive.google.com/file/d/1jkSOUwZV5GJ7rZlVeErjcu0DBQs8Np0d/view?usp=drive_link
-https://drive.google.com/file/d/1VIN1eLI-93WrVQwCjsv6XQr353DqqBYA/view?usp=drive_link

lerobot/common/datasets/push_dataset_to_hub/_aloha_raw_urls/static_ziploc_slide.txt

@@ -1,8 +0,0 @@
-https://drive.google.com/drive/folders/1EgKar7rWBmTIRmeJYZciSwjZx3uP2mHO
-https://drive.google.com/file/d/12eYWQO15atK2hBjXhynPJd9MKAj_42pz/view?usp=drive_link
-https://drive.google.com/file/d/1Ul4oEeICJDjgfYTl4H1uaisTzVYIM6wd/view?usp=drive_link
-https://drive.google.com/file/d/1WSF-OG8lKSe2wVYCv5D1aJNipxpgddk-/view?usp=drive_link
-https://drive.google.com/file/d/1_ppD5j5sFh26aWW0JmhLzJMeNB-lCArk/view?usp=drive_link
-https://drive.google.com/file/d/1WUp846dgWXYhu4oJfhHxiU6YL_7N6s4W/view?usp=drive_link
-https://drive.google.com/file/d/1HRZNAIoAQw_uYiPwnBvtBioQoqiqoXdA/view?usp=drive_link
-https://drive.google.com/file/d/1hedGq-QDMnIn8GlXXBC3GiEJ_Y-LTxyt/view?usp=drive_link

lerobot/common/datasets/push_dataset_to_hub/_diffusion_policy_replay_buffer.py

@@ -1,634 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Helper code for loading PushT dataset from Diffusion Policy (https://diffusion-policy.cs.columbia.edu/)
-
-Copied from the original Diffusion Policy repository and used in our `download_and_upload_dataset.py` script.
-"""
-
-from __future__ import annotations
-
-import math
-import numbers
-import os
-from functools import cached_property
-
-import numcodecs
-import numpy as np
-import zarr
-
-
-def check_chunks_compatible(chunks: tuple, shape: tuple):
-    assert len(shape) == len(chunks)
-    for c in chunks:
-        assert isinstance(c, numbers.Integral)
-        assert c > 0
-
-
-def rechunk_recompress_array(group, name, chunks=None, chunk_length=None, compressor=None, tmp_key="_temp"):
-    old_arr = group[name]
-    if chunks is None:
-        chunks = (chunk_length,) + old_arr.chunks[1:] if chunk_length is not None else old_arr.chunks
-    check_chunks_compatible(chunks, old_arr.shape)
-
-    if compressor is None:
-        compressor = old_arr.compressor
-
-    if (chunks == old_arr.chunks) and (compressor == old_arr.compressor):
-        # no change
-        return old_arr
-
-    # rechunk recompress
-    group.move(name, tmp_key)
-    old_arr = group[tmp_key]
-    n_copied, n_skipped, n_bytes_copied = zarr.copy(
-        source=old_arr,
-        dest=group,
-        name=name,
-        chunks=chunks,
-        compressor=compressor,
-    )
-    del group[tmp_key]
-    arr = group[name]
-    return arr
-
-
-def get_optimal_chunks(shape, dtype, target_chunk_bytes=2e6, max_chunk_length=None):
-    """
-    Common shapes
-    T,D
-    T,N,D
-    T,H,W,C
-    T,N,H,W,C
-    """
-    itemsize = np.dtype(dtype).itemsize
-    # reversed
-    rshape = list(shape[::-1])
-    if max_chunk_length is not None:
-        rshape[-1] = int(max_chunk_length)
-    split_idx = len(shape) - 1
-    for i in range(len(shape) - 1):
-        this_chunk_bytes = itemsize * np.prod(rshape[:i])
-        next_chunk_bytes = itemsize * np.prod(rshape[: i + 1])
-        if this_chunk_bytes <= target_chunk_bytes and next_chunk_bytes > target_chunk_bytes:
-            split_idx = i
-
-    rchunks = rshape[:split_idx]
-    item_chunk_bytes = itemsize * np.prod(rshape[:split_idx])
-    this_max_chunk_length = rshape[split_idx]
-    next_chunk_length = min(this_max_chunk_length, math.ceil(target_chunk_bytes / item_chunk_bytes))
-    rchunks.append(next_chunk_length)
-    len_diff = len(shape) - len(rchunks)
-    rchunks.extend([1] * len_diff)
-    chunks = tuple(rchunks[::-1])
-    # print(np.prod(chunks) * itemsize / target_chunk_bytes)
-    return chunks
-
-
-class ReplayBuffer:
-    """
-    Zarr-based temporal datastructure.
-    Assumes first dimension to be time. Only chunk in time dimension.
-    """
-
-    def __init__(self, root: zarr.Group | dict[str, dict]):
-        """
-        Dummy constructor. Use copy_from* and create_from* class methods instead.
-        """
-        assert "data" in root
-        assert "meta" in root
-        assert "episode_ends" in root["meta"]
-        for value in root["data"].values():
-            assert value.shape[0] == root["meta"]["episode_ends"][-1]
-        self.root = root
-
-    # ============= create constructors ===============
-    @classmethod
-    def create_empty_zarr(cls, storage=None, root=None):
-        if root is None:
-            if storage is None:
-                storage = zarr.MemoryStore()
-            root = zarr.group(store=storage)
-        root.require_group("data", overwrite=False)
-        meta = root.require_group("meta", overwrite=False)
-        if "episode_ends" not in meta:
-            meta.zeros("episode_ends", shape=(0,), dtype=np.int64, compressor=None, overwrite=False)
-        return cls(root=root)
-
-    @classmethod
-    def create_empty_numpy(cls):
-        root = {"data": {}, "meta": {"episode_ends": np.zeros((0,), dtype=np.int64)}}
-        return cls(root=root)
-
-    @classmethod
-    def create_from_group(cls, group, **kwargs):
-        if "data" not in group:
-            # create from stratch
-            buffer = cls.create_empty_zarr(root=group, **kwargs)
-        else:
-            # already exist
-            buffer = cls(root=group, **kwargs)
-        return buffer
-
-    @classmethod
-    def create_from_path(cls, zarr_path, mode="r", **kwargs):
-        """
-        Open a on-disk zarr directly (for dataset larger than memory).
-        Slower.
-        """
-        group = zarr.open(os.path.expanduser(zarr_path), mode)
-        return cls.create_from_group(group, **kwargs)
-
-    # ============= copy constructors ===============
-    @classmethod
-    def copy_from_store(
-        cls,
-        src_store,
-        store=None,
-        keys=None,
-        chunks: dict[str, tuple] | None = None,
-        compressors: dict | str | numcodecs.abc.Codec | None = None,
-        if_exists="replace",
-        **kwargs,
-    ):
-        """
-        Load to memory.
-        """
-        src_root = zarr.group(src_store)
-        if chunks is None:
-            chunks = {}
-        if compressors is None:
-            compressors = {}
-        root = None
-        if store is None:
-            # numpy backend
-            meta = {}
-            for key, value in src_root["meta"].items():
-                if len(value.shape) == 0:
-                    meta[key] = np.array(value)
-                else:
-                    meta[key] = value[:]
-
-            if keys is None:
-                keys = src_root["data"].keys()
-            data = {}
-            for key in keys:
-                arr = src_root["data"][key]
-                data[key] = arr[:]
-
-            root = {"meta": meta, "data": data}
-        else:
-            root = zarr.group(store=store)
-            # copy without recompression
-            n_copied, n_skipped, n_bytes_copied = zarr.copy_store(
-                source=src_store, dest=store, source_path="/meta", dest_path="/meta", if_exists=if_exists
-            )
-            data_group = root.create_group("data", overwrite=True)
-            if keys is None:
-                keys = src_root["data"].keys()
-            for key in keys:
-                value = src_root["data"][key]
-                cks = cls._resolve_array_chunks(chunks=chunks, key=key, array=value)
-                cpr = cls._resolve_array_compressor(compressors=compressors, key=key, array=value)
-                if cks == value.chunks and cpr == value.compressor:
-                    # copy without recompression
-                    this_path = "/data/" + key
-                    n_copied, n_skipped, n_bytes_copied = zarr.copy_store(
-                        source=src_store,
-                        dest=store,
-                        source_path=this_path,
-                        dest_path=this_path,
-                        if_exists=if_exists,
-                    )
-                else:
-                    # copy with recompression
-                    n_copied, n_skipped, n_bytes_copied = zarr.copy(
-                        source=value,
-                        dest=data_group,
-                        name=key,
-                        chunks=cks,
-                        compressor=cpr,
-                        if_exists=if_exists,
-                    )
-        buffer = cls(root=root)
-        return buffer
-
-    @classmethod
-    def copy_from_path(
-        cls,
-        zarr_path,
-        backend=None,
-        store=None,
-        keys=None,
-        chunks: dict[str, tuple] | None = None,
-        compressors: dict | str | numcodecs.abc.Codec | None = None,
-        if_exists="replace",
-        **kwargs,
-    ):
-        """
-        Copy a on-disk zarr to in-memory compressed.
-        Recommended
-        """
-        if chunks is None:
-            chunks = {}
-        if compressors is None:
-            compressors = {}
-        if backend == "numpy":
-            print("backend argument is deprecated!")
-            store = None
-        group = zarr.open(os.path.expanduser(zarr_path), "r")
-        return cls.copy_from_store(
-            src_store=group.store,
-            store=store,
-            keys=keys,
-            chunks=chunks,
-            compressors=compressors,
-            if_exists=if_exists,
-            **kwargs,
-        )
-
-    # ============= save methods ===============
-    def save_to_store(
-        self,
-        store,
-        chunks: dict[str, tuple] | None = None,
-        compressors: str | numcodecs.abc.Codec | dict | None = None,
-        if_exists="replace",
-        **kwargs,
-    ):
-        root = zarr.group(store)
-        if chunks is None:
-            chunks = {}
-        if compressors is None:
-            compressors = {}
-        if self.backend == "zarr":
-            # recompression free copy
-            n_copied, n_skipped, n_bytes_copied = zarr.copy_store(
-                source=self.root.store,
-                dest=store,
-                source_path="/meta",
-                dest_path="/meta",
-                if_exists=if_exists,
-            )
-        else:
-            meta_group = root.create_group("meta", overwrite=True)
-            # save meta, no chunking
-            for key, value in self.root["meta"].items():
-                _ = meta_group.array(name=key, data=value, shape=value.shape, chunks=value.shape)
-
-        # save data, chunk
-        data_group = root.create_group("data", overwrite=True)
-        for key, value in self.root["data"].items():
-            cks = self._resolve_array_chunks(chunks=chunks, key=key, array=value)
-            cpr = self._resolve_array_compressor(compressors=compressors, key=key, array=value)
-            if isinstance(value, zarr.Array):
-                if cks == value.chunks and cpr == value.compressor:
-                    # copy without recompression
-                    this_path = "/data/" + key
-                    n_copied, n_skipped, n_bytes_copied = zarr.copy_store(
-                        source=self.root.store,
-                        dest=store,
-                        source_path=this_path,
-                        dest_path=this_path,
-                        if_exists=if_exists,
-                    )
-                else:
-                    # copy with recompression
-                    n_copied, n_skipped, n_bytes_copied = zarr.copy(
-                        source=value,
-                        dest=data_group,
-                        name=key,
-                        chunks=cks,
-                        compressor=cpr,
-                        if_exists=if_exists,
-                    )
-            else:
-                # numpy
-                _ = data_group.array(name=key, data=value, chunks=cks, compressor=cpr)
-        return store
-
-    def save_to_path(
-        self,
-        zarr_path,
-        chunks: dict[str, tuple] | None = None,
-        compressors: str | numcodecs.abc.Codec | dict | None = None,
-        if_exists="replace",
-        **kwargs,
-    ):
-        if chunks is None:
-            chunks = {}
-        if compressors is None:
-            compressors = {}
-        store = zarr.DirectoryStore(os.path.expanduser(zarr_path))
-        return self.save_to_store(
-            store, chunks=chunks, compressors=compressors, if_exists=if_exists, **kwargs
-        )
-
-    @staticmethod
-    def resolve_compressor(compressor="default"):
-        if compressor == "default":
-            compressor = numcodecs.Blosc(cname="lz4", clevel=5, shuffle=numcodecs.Blosc.NOSHUFFLE)
-        elif compressor == "disk":
-            compressor = numcodecs.Blosc("zstd", clevel=5, shuffle=numcodecs.Blosc.BITSHUFFLE)
-        return compressor
-
-    @classmethod
-    def _resolve_array_compressor(cls, compressors: dict | str | numcodecs.abc.Codec, key, array):
-        # allows compressor to be explicitly set to None
-        cpr = "nil"
-        if isinstance(compressors, dict):
-            if key in compressors:
-                cpr = cls.resolve_compressor(compressors[key])
-            elif isinstance(array, zarr.Array):
-                cpr = array.compressor
-        else:
-            cpr = cls.resolve_compressor(compressors)
-        # backup default
-        if cpr == "nil":
-            cpr = cls.resolve_compressor("default")
-        return cpr
-
-    @classmethod
-    def _resolve_array_chunks(cls, chunks: dict | tuple, key, array):
-        cks = None
-        if isinstance(chunks, dict):
-            if key in chunks:
-                cks = chunks[key]
-            elif isinstance(array, zarr.Array):
-                cks = array.chunks
-        elif isinstance(chunks, tuple):
-            cks = chunks
-        else:
-            raise TypeError(f"Unsupported chunks type {type(chunks)}")
-        # backup default
-        if cks is None:
-            cks = get_optimal_chunks(shape=array.shape, dtype=array.dtype)
-        # check
-        check_chunks_compatible(chunks=cks, shape=array.shape)
-        return cks
-
-    # ============= properties =================
-    @cached_property
-    def data(self):
-        return self.root["data"]
-
-    @cached_property
-    def meta(self):
-        return self.root["meta"]
-
-    def update_meta(self, data):
-        # sanitize data
-        np_data = {}
-        for key, value in data.items():
-            if isinstance(value, np.ndarray):
-                np_data[key] = value
-            else:
-                arr = np.array(value)
-                if arr.dtype == object:
-                    raise TypeError(f"Invalid value type {type(value)}")
-                np_data[key] = arr
-
-        meta_group = self.meta
-        if self.backend == "zarr":
-            for key, value in np_data.items():
-                _ = meta_group.array(
-                    name=key, data=value, shape=value.shape, chunks=value.shape, overwrite=True
-                )
-        else:
-            meta_group.update(np_data)
-
-        return meta_group
-
-    @property
-    def episode_ends(self):
-        return self.meta["episode_ends"]
-
-    def get_episode_idxs(self):
-        import numba
-
-        numba.jit(nopython=True)
-
-        def _get_episode_idxs(episode_ends):
-            result = np.zeros((episode_ends[-1],), dtype=np.int64)
-            for i in range(len(episode_ends)):
-                start = 0
-                if i > 0:
-                    start = episode_ends[i - 1]
-                end = episode_ends[i]
-                for idx in range(start, end):
-                    result[idx] = i
-            return result
-
-        return _get_episode_idxs(self.episode_ends)
-
-    @property
-    def backend(self):
-        backend = "numpy"
-        if isinstance(self.root, zarr.Group):
-            backend = "zarr"
-        return backend
-
-    # =========== dict-like API ==============
-    def __repr__(self) -> str:
-        if self.backend == "zarr":
-            return str(self.root.tree())
-        else:
-            return super().__repr__()
-
-    def keys(self):
-        return self.data.keys()
-
-    def values(self):
-        return self.data.values()
-
-    def items(self):
-        return self.data.items()
-
-    def __getitem__(self, key):
-        return self.data[key]
-
-    def __contains__(self, key):
-        return key in self.data
-
-    # =========== our API ==============
-    @property
-    def n_steps(self):
-        if len(self.episode_ends) == 0:
-            return 0
-        return self.episode_ends[-1]
-
-    @property
-    def n_episodes(self):
-        return len(self.episode_ends)
-
-    @property
-    def chunk_size(self):
-        if self.backend == "zarr":
-            return next(iter(self.data.arrays()))[-1].chunks[0]
-        return None
-
-    @property
-    def episode_lengths(self):
-        ends = self.episode_ends[:]
-        ends = np.insert(ends, 0, 0)
-        lengths = np.diff(ends)
-        return lengths
-
-    def add_episode(
-        self,
-        data: dict[str, np.ndarray],
-        chunks: dict[str, tuple] | None = None,
-        compressors: str | numcodecs.abc.Codec | dict | None = None,
-    ):
-        if chunks is None:
-            chunks = {}
-        if compressors is None:
-            compressors = {}
-        assert len(data) > 0
-        is_zarr = self.backend == "zarr"
-
-        curr_len = self.n_steps
-        episode_length = None
-        for value in data.values():
-            assert len(value.shape) >= 1
-            if episode_length is None:
-                episode_length = len(value)
-            else:
-                assert episode_length == len(value)
-        new_len = curr_len + episode_length
-
-        for key, value in data.items():
-            new_shape = (new_len,) + value.shape[1:]
-            # create array
-            if key not in self.data:
-                if is_zarr:
-                    cks = self._resolve_array_chunks(chunks=chunks, key=key, array=value)
-                    cpr = self._resolve_array_compressor(compressors=compressors, key=key, array=value)
-                    arr = self.data.zeros(
-                        name=key, shape=new_shape, chunks=cks, dtype=value.dtype, compressor=cpr
-                    )
-                else:
-                    # copy data to prevent modify
-                    arr = np.zeros(shape=new_shape, dtype=value.dtype)
-                    self.data[key] = arr
-            else:
-                arr = self.data[key]
-                assert value.shape[1:] == arr.shape[1:]
-                # same method for both zarr and numpy
-                if is_zarr:
-                    arr.resize(new_shape)
-                else:
-                    arr.resize(new_shape, refcheck=False)
-            # copy data
-            arr[-value.shape[0] :] = value
-
-        # append to episode ends
-        episode_ends = self.episode_ends
-        if is_zarr:
-            episode_ends.resize(episode_ends.shape[0] + 1)
-        else:
-            episode_ends.resize(episode_ends.shape[0] + 1, refcheck=False)
-        episode_ends[-1] = new_len
-
-        # rechunk
-        if is_zarr and episode_ends.chunks[0] < episode_ends.shape[0]:
-            rechunk_recompress_array(self.meta, "episode_ends", chunk_length=int(episode_ends.shape[0] * 1.5))
-
-    def drop_episode(self):
-        is_zarr = self.backend == "zarr"
-        episode_ends = self.episode_ends[:].copy()
-        assert len(episode_ends) > 0
-        start_idx = 0
-        if len(episode_ends) > 1:
-            start_idx = episode_ends[-2]
-        for value in self.data.values():
-            new_shape = (start_idx,) + value.shape[1:]
-            if is_zarr:
-                value.resize(new_shape)
-            else:
-                value.resize(new_shape, refcheck=False)
-        if is_zarr:
-            self.episode_ends.resize(len(episode_ends) - 1)
-        else:
-            self.episode_ends.resize(len(episode_ends) - 1, refcheck=False)
-
-    def pop_episode(self):
-        assert self.n_episodes > 0
-        episode = self.get_episode(self.n_episodes - 1, copy=True)
-        self.drop_episode()
-        return episode
-
-    def extend(self, data):
-        self.add_episode(data)
-
-    def get_episode(self, idx, copy=False):
-        idx = list(range(len(self.episode_ends)))[idx]
-        start_idx = 0
-        if idx > 0:
-            start_idx = self.episode_ends[idx - 1]
-        end_idx = self.episode_ends[idx]
-        result = self.get_steps_slice(start_idx, end_idx, copy=copy)
-        return result
-
-    def get_episode_slice(self, idx):
-        start_idx = 0
-        if idx > 0:
-            start_idx = self.episode_ends[idx - 1]
-        end_idx = self.episode_ends[idx]
-        return slice(start_idx, end_idx)
-
-    def get_steps_slice(self, start, stop, step=None, copy=False):
-        _slice = slice(start, stop, step)
-
-        result = {}
-        for key, value in self.data.items():
-            x = value[_slice]
-            if copy and isinstance(value, np.ndarray):
-                x = x.copy()
-            result[key] = x
-        return result
-
-    # =========== chunking =============
-    def get_chunks(self) -> dict:
-        assert self.backend == "zarr"
-        chunks = {}
-        for key, value in self.data.items():
-            chunks[key] = value.chunks
-        return chunks
-
-    def set_chunks(self, chunks: dict):
-        assert self.backend == "zarr"
-        for key, value in chunks.items():
-            if key in self.data:
-                arr = self.data[key]
-                if value != arr.chunks:
-                    check_chunks_compatible(chunks=value, shape=arr.shape)
-                    rechunk_recompress_array(self.data, key, chunks=value)
-
-    def get_compressors(self) -> dict:
-        assert self.backend == "zarr"
-        compressors = {}
-        for key, value in self.data.items():
-            compressors[key] = value.compressor
-        return compressors
-
-    def set_compressors(self, compressors: dict):
-        assert self.backend == "zarr"
-        for key, value in compressors.items():
-            if key in self.data:
-                arr = self.data[key]
-                compressor = self.resolve_compressor(value)
-                if compressor != arr.compressor:
-                    rechunk_recompress_array(self.data, key, compressor=compressor)

lerobot/common/datasets/push_dataset_to_hub/_download_raw.py

@@ -1,202 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""
-This file contains download scripts for raw datasets.
-
-Example of usage:
-```
-python lerobot/common/datasets/push_dataset_to_hub/_download_raw.py \
---raw-dir data/lerobot-raw/pusht_raw \
---repo-id lerobot-raw/pusht_raw
-```
-"""
-
-import argparse
-import logging
-import warnings
-from pathlib import Path
-
-from huggingface_hub import snapshot_download
-
-from lerobot.common.datasets.push_dataset_to_hub.utils import check_repo_id
-
-# {raw_repo_id: raw_format}
-AVAILABLE_RAW_REPO_IDS = {
-    "lerobot-raw/aloha_mobile_cabinet_raw": "aloha_hdf5",
-    "lerobot-raw/aloha_mobile_chair_raw": "aloha_hdf5",
-    "lerobot-raw/aloha_mobile_elevator_raw": "aloha_hdf5",
-    "lerobot-raw/aloha_mobile_shrimp_raw": "aloha_hdf5",
-    "lerobot-raw/aloha_mobile_wash_pan_raw": "aloha_hdf5",
-    "lerobot-raw/aloha_mobile_wipe_wine_raw": "aloha_hdf5",
-    "lerobot-raw/aloha_sim_insertion_human_raw": "aloha_hdf5",
-    "lerobot-raw/aloha_sim_insertion_scripted_raw": "aloha_hdf5",
-    "lerobot-raw/aloha_sim_transfer_cube_human_raw": "aloha_hdf5",
-    "lerobot-raw/aloha_sim_transfer_cube_scripted_raw": "aloha_hdf5",
-    "lerobot-raw/aloha_static_battery_raw": "aloha_hdf5",
-    "lerobot-raw/aloha_static_candy_raw": "aloha_hdf5",
-    "lerobot-raw/aloha_static_coffee_new_raw": "aloha_hdf5",
-    "lerobot-raw/aloha_static_coffee_raw": "aloha_hdf5",
-    "lerobot-raw/aloha_static_cups_open_raw": "aloha_hdf5",
-    "lerobot-raw/aloha_static_fork_pick_up_raw": "aloha_hdf5",
-    "lerobot-raw/aloha_static_pingpong_test_raw": "aloha_hdf5",
-    "lerobot-raw/aloha_static_pro_pencil_raw": "aloha_hdf5",
-    "lerobot-raw/aloha_static_screw_driver_raw": "aloha_hdf5",
-    "lerobot-raw/aloha_static_tape_raw": "aloha_hdf5",
-    "lerobot-raw/aloha_static_thread_velcro_raw": "aloha_hdf5",
-    "lerobot-raw/aloha_static_towel_raw": "aloha_hdf5",
-    "lerobot-raw/aloha_static_vinh_cup_left_raw": "aloha_hdf5",
-    "lerobot-raw/aloha_static_vinh_cup_raw": "aloha_hdf5",
-    "lerobot-raw/aloha_static_ziploc_slide_raw": "aloha_hdf5",
-    "lerobot-raw/umi_cup_in_the_wild_raw": "umi_zarr",
-    "lerobot-raw/pusht_raw": "pusht_zarr",
-    "lerobot-raw/unitreeh1_fold_clothes_raw": "aloha_hdf5",
-    "lerobot-raw/unitreeh1_rearrange_objects_raw": "aloha_hdf5",
-    "lerobot-raw/unitreeh1_two_robot_greeting_raw": "aloha_hdf5",
-    "lerobot-raw/unitreeh1_warehouse_raw": "aloha_hdf5",
-    "lerobot-raw/xarm_lift_medium_raw": "xarm_pkl",
-    "lerobot-raw/xarm_lift_medium_replay_raw": "xarm_pkl",
-    "lerobot-raw/xarm_push_medium_raw": "xarm_pkl",
-    "lerobot-raw/xarm_push_medium_replay_raw": "xarm_pkl",
-    "lerobot-raw/fractal20220817_data_raw": "openx_rlds.fractal20220817_data",
-    "lerobot-raw/kuka_raw": "openx_rlds.kuka",
-    "lerobot-raw/bridge_openx_raw": "openx_rlds.bridge_openx",
-    "lerobot-raw/taco_play_raw": "openx_rlds.taco_play",
-    "lerobot-raw/jaco_play_raw": "openx_rlds.jaco_play",
-    "lerobot-raw/berkeley_cable_routing_raw": "openx_rlds.berkeley_cable_routing",
-    "lerobot-raw/roboturk_raw": "openx_rlds.roboturk",
-    "lerobot-raw/nyu_door_opening_surprising_effectiveness_raw": "openx_rlds.nyu_door_opening_surprising_effectiveness",
-    "lerobot-raw/viola_raw": "openx_rlds.viola",
-    "lerobot-raw/berkeley_autolab_ur5_raw": "openx_rlds.berkeley_autolab_ur5",
-    "lerobot-raw/toto_raw": "openx_rlds.toto",
-    "lerobot-raw/language_table_raw": "openx_rlds.language_table",
-    "lerobot-raw/columbia_cairlab_pusht_real_raw": "openx_rlds.columbia_cairlab_pusht_real",
-    "lerobot-raw/stanford_kuka_multimodal_dataset_raw": "openx_rlds.stanford_kuka_multimodal_dataset",
-    "lerobot-raw/nyu_rot_dataset_raw": "openx_rlds.nyu_rot_dataset",
-    "lerobot-raw/io_ai_tech_raw": "openx_rlds.io_ai_tech",
-    "lerobot-raw/stanford_hydra_dataset_raw": "openx_rlds.stanford_hydra_dataset",
-    "lerobot-raw/austin_buds_dataset_raw": "openx_rlds.austin_buds_dataset",
-    "lerobot-raw/nyu_franka_play_dataset_raw": "openx_rlds.nyu_franka_play_dataset",
-    "lerobot-raw/maniskill_dataset_raw": "openx_rlds.maniskill_dataset",
-    "lerobot-raw/furniture_bench_dataset_raw": "openx_rlds.furniture_bench_dataset",
-    "lerobot-raw/cmu_franka_exploration_dataset_raw": "openx_rlds.cmu_franka_exploration_dataset",
-    "lerobot-raw/ucsd_kitchen_dataset_raw": "openx_rlds.ucsd_kitchen_dataset",
-    "lerobot-raw/ucsd_pick_and_place_dataset_raw": "openx_rlds.ucsd_pick_and_place_dataset",
-    "lerobot-raw/spoc_raw": "openx_rlds.spoc",
-    "lerobot-raw/austin_sailor_dataset_raw": "openx_rlds.austin_sailor_dataset",
-    "lerobot-raw/austin_sirius_dataset_raw": "openx_rlds.austin_sirius_dataset",
-    "lerobot-raw/bc_z_raw": "openx_rlds.bc_z",
-    "lerobot-raw/utokyo_pr2_opening_fridge_raw": "openx_rlds.utokyo_pr2_opening_fridge",
-    "lerobot-raw/utokyo_pr2_tabletop_manipulation_raw": "openx_rlds.utokyo_pr2_tabletop_manipulation",
-    "lerobot-raw/utokyo_xarm_pick_and_place_raw": "openx_rlds.utokyo_xarm_pick_and_place",
-    "lerobot-raw/utokyo_xarm_bimanual_raw": "openx_rlds.utokyo_xarm_bimanual",
-    "lerobot-raw/utokyo_saytap_raw": "openx_rlds.utokyo_saytap",
-    "lerobot-raw/robo_net_raw": "openx_rlds.robo_net",
-    "lerobot-raw/robo_set_raw": "openx_rlds.robo_set",
-    "lerobot-raw/berkeley_mvp_raw": "openx_rlds.berkeley_mvp",
-    "lerobot-raw/berkeley_rpt_raw": "openx_rlds.berkeley_rpt",
-    "lerobot-raw/kaist_nonprehensile_raw": "openx_rlds.kaist_nonprehensile",
-    "lerobot-raw/stanford_mask_vit_raw": "openx_rlds.stanford_mask_vit",
-    "lerobot-raw/tokyo_u_lsmo_raw": "openx_rlds.tokyo_u_lsmo",
-    "lerobot-raw/dlr_sara_pour_raw": "openx_rlds.dlr_sara_pour",
-    "lerobot-raw/dlr_sara_grid_clamp_raw": "openx_rlds.dlr_sara_grid_clamp",
-    "lerobot-raw/dlr_edan_shared_control_raw": "openx_rlds.dlr_edan_shared_control",
-    "lerobot-raw/asu_table_top_raw": "openx_rlds.asu_table_top",
-    "lerobot-raw/stanford_robocook_raw": "openx_rlds.stanford_robocook",
-    "lerobot-raw/imperialcollege_sawyer_wrist_cam_raw": "openx_rlds.imperialcollege_sawyer_wrist_cam",
-    "lerobot-raw/iamlab_cmu_pickup_insert_raw": "openx_rlds.iamlab_cmu_pickup_insert",
-    "lerobot-raw/uiuc_d3field_raw": "openx_rlds.uiuc_d3field",
-    "lerobot-raw/utaustin_mutex_raw": "openx_rlds.utaustin_mutex",
-    "lerobot-raw/berkeley_fanuc_manipulation_raw": "openx_rlds.berkeley_fanuc_manipulation",
-    "lerobot-raw/cmu_playing_with_food_raw": "openx_rlds.cmu_playing_with_food",
-    "lerobot-raw/cmu_play_fusion_raw": "openx_rlds.cmu_play_fusion",
-    "lerobot-raw/cmu_stretch_raw": "openx_rlds.cmu_stretch",
-    "lerobot-raw/berkeley_gnm_recon_raw": "openx_rlds.berkeley_gnm_recon",
-    "lerobot-raw/berkeley_gnm_cory_hall_raw": "openx_rlds.berkeley_gnm_cory_hall",
-    "lerobot-raw/berkeley_gnm_sac_son_raw": "openx_rlds.berkeley_gnm_sac_son",
-    "lerobot-raw/droid_raw": "openx_rlds.droid",
-    "lerobot-raw/droid_100_raw": "openx_rlds.droid100",
-    "lerobot-raw/fmb_raw": "openx_rlds.fmb",
-    "lerobot-raw/dobbe_raw": "openx_rlds.dobbe",
-    "lerobot-raw/usc_cloth_sim_raw": "openx_rlds.usc_cloth_sim",
-    "lerobot-raw/plex_robosuite_raw": "openx_rlds.plex_robosuite",
-    "lerobot-raw/conq_hose_manipulation_raw": "openx_rlds.conq_hose_manipulation",
-    "lerobot-raw/vima_raw": "openx_rlds.vima",
-    "lerobot-raw/robot_vqa_raw": "openx_rlds.robot_vqa",
-    "lerobot-raw/mimic_play_raw": "openx_rlds.mimic_play",
-    "lerobot-raw/tidybot_raw": "openx_rlds.tidybot",
-    "lerobot-raw/eth_agent_affordances_raw": "openx_rlds.eth_agent_affordances",
-}
-
-
-def download_raw(raw_dir: Path, repo_id: str):
-    check_repo_id(repo_id)
-    user_id, dataset_id = repo_id.split("/")
-
-    if not dataset_id.endswith("_raw"):
-        warnings.warn(
-            f"""`dataset_id` ({dataset_id}) doesn't end with '_raw' (e.g. 'lerobot/pusht_raw'). Following this
-             naming convention by renaming your repository is advised, but not mandatory.""",
-            stacklevel=1,
-        )
-
-    # Send warning if raw_dir isn't well formatted
-    if raw_dir.parts[-2] != user_id or raw_dir.parts[-1] != dataset_id:
-        warnings.warn(
-            f"""`raw_dir` ({raw_dir}) doesn't contain a community or user id `/` the name of the dataset that
-             match the `repo_id` (e.g. 'data/lerobot/pusht_raw'). Following this naming convention is advised,
-             but not mandatory.""",
-            stacklevel=1,
-        )
-    raw_dir.mkdir(parents=True, exist_ok=True)
-
-    logging.info(f"Start downloading from huggingface.co/{user_id} for {dataset_id}")
-    snapshot_download(repo_id, repo_type="dataset", local_dir=raw_dir)
-    logging.info(f"Finish downloading from huggingface.co/{user_id} for {dataset_id}")
-
-
-def download_all_raw_datasets(data_dir: Path | None = None):
-    if data_dir is None:
-        data_dir = Path("data")
-    for repo_id in AVAILABLE_RAW_REPO_IDS:
-        raw_dir = data_dir / repo_id
-        download_raw(raw_dir, repo_id)
-
-
-def main():
-    parser = argparse.ArgumentParser(
-        description=f"""A script to download raw datasets from Hugging Face hub to a local directory. Here is a
-            non exhaustive list of available repositories to use in `--repo-id`: {list(AVAILABLE_RAW_REPO_IDS.keys())}""",
-    )
-
-    parser.add_argument(
-        "--raw-dir",
-        type=Path,
-        required=True,
-        help="Directory containing input raw datasets (e.g. `data/aloha_mobile_chair_raw` or `data/pusht_raw).",
-    )
-    parser.add_argument(
-        "--repo-id",
-        type=str,
-        required=True,
-        help="""Repositery identifier on Hugging Face: a community or a user name `/` the name of
-        the dataset (e.g. `lerobot/pusht_raw`, `cadene/aloha_sim_insertion_human_raw`).""",
-    )
-    args = parser.parse_args()
-    download_raw(**vars(args))
-
-
-if __name__ == "__main__":
-    main()

lerobot/common/datasets/push_dataset_to_hub/_encode_datasets.py

@@ -1,184 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""
-Use this script to batch encode lerobot dataset from their raw format to LeRobotDataset and push their updated
-version to the hub. Under the hood, this script reuses 'push_dataset_to_hub.py'. It assumes that you already
-downloaded raw datasets, which you can do with the related '_download_raw.py' script.
-
-For instance, for codebase_version = 'v1.6', the following command was run, assuming raw datasets from
-lerobot-raw were downloaded in 'raw/datasets/directory':
-```bash
-python lerobot/common/datasets/push_dataset_to_hub/_encode_datasets.py \
-  --raw-dir raw/datasets/directory \
-  --raw-repo-ids lerobot-raw \
-  --local-dir push/datasets/directory \
-  --tests-data-dir tests/data \
-  --push-repo lerobot \
-  --vcodec libsvtav1 \
-  --pix-fmt yuv420p \
-  --g 2 \
-  --crf 30
-```
-"""
-
-import argparse
-from pathlib import Path
-
-from lerobot.common.datasets.lerobot_dataset import CODEBASE_VERSION
-from lerobot.common.datasets.push_dataset_to_hub._download_raw import AVAILABLE_RAW_REPO_IDS
-from lerobot.common.datasets.push_dataset_to_hub.utils import check_repo_id
-from lerobot.scripts.push_dataset_to_hub import push_dataset_to_hub
-
-
-def get_push_repo_id_from_raw(raw_repo_id: str, push_repo: str) -> str:
-    dataset_id_raw = raw_repo_id.split("/")[1]
-    dataset_id = dataset_id_raw.removesuffix("_raw")
-    return f"{push_repo}/{dataset_id}"
-
-
-def encode_datasets(
-    raw_dir: Path,
-    raw_repo_ids: list[str],
-    push_repo: str,
-    vcodec: str,
-    pix_fmt: str,
-    g: int,
-    crf: int,
-    local_dir: Path | None = None,
-    tests_data_dir: Path | None = None,
-    raw_format: str | None = None,
-    dry_run: bool = False,
-) -> None:
-    if len(raw_repo_ids) == 1 and raw_repo_ids[0].lower() == "lerobot-raw":
-        raw_repo_ids_format = AVAILABLE_RAW_REPO_IDS
-    else:
-        if raw_format is None:
-            raise ValueError(raw_format)
-        raw_repo_ids_format = {id_: raw_format for id_ in raw_repo_ids}
-
-    for raw_repo_id, repo_raw_format in raw_repo_ids_format.items():
-        check_repo_id(raw_repo_id)
-        dataset_repo_id_push = get_push_repo_id_from_raw(raw_repo_id, push_repo)
-        dataset_raw_dir = raw_dir / raw_repo_id
-        dataset_dir = local_dir / dataset_repo_id_push if local_dir is not None else None
-        encoding = {
-            "vcodec": vcodec,
-            "pix_fmt": pix_fmt,
-            "g": g,
-            "crf": crf,
-        }
-
-        if not (dataset_raw_dir).is_dir():
-            raise NotADirectoryError(dataset_raw_dir)
-
-        if not dry_run:
-            push_dataset_to_hub(
-                dataset_raw_dir,
-                raw_format=repo_raw_format,
-                repo_id=dataset_repo_id_push,
-                local_dir=dataset_dir,
-                resume=True,
-                encoding=encoding,
-                tests_data_dir=tests_data_dir,
-            )
-        else:
-            print(
-                f"DRY RUN: {dataset_raw_dir}  -->  {dataset_dir}  -->  {dataset_repo_id_push}@{CODEBASE_VERSION}"
-            )
-
-
-def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--raw-dir",
-        type=Path,
-        default=Path("data"),
-        help="Directory where raw datasets are located.",
-    )
-    parser.add_argument(
-        "--raw-repo-ids",
-        type=str,
-        nargs="*",
-        default=["lerobot-raw"],
-        help="""Raw dataset repo ids. if 'lerobot-raw', the keys from `AVAILABLE_RAW_REPO_IDS` will be
-            used and raw datasets will be fetched from the 'lerobot-raw/' repo and pushed with their
-            associated format. It is assumed that each dataset is located at `raw_dir / raw_repo_id` """,
-    )
-    parser.add_argument(
-        "--raw-format",
-        type=str,
-        default=None,
-        help="""Raw format to use for the raw repo-ids. Must be specified if --raw-repo-ids is not
-            'lerobot-raw'""",
-    )
-    parser.add_argument(
-        "--local-dir",
-        type=Path,
-        default=None,
-        help="""When provided, writes the dataset converted to LeRobotDataset format in this directory
-        (e.g. `data/lerobot/aloha_mobile_chair`).""",
-    )
-    parser.add_argument(
-        "--push-repo",
-        type=str,
-        default="lerobot",
-        help="Repo to upload datasets to",
-    )
-    parser.add_argument(
-        "--vcodec",
-        type=str,
-        default="libsvtav1",
-        help="Codec to use for encoding videos",
-    )
-    parser.add_argument(
-        "--pix-fmt",
-        type=str,
-        default="yuv420p",
-        help="Pixel formats (chroma subsampling) to be used for encoding",
-    )
-    parser.add_argument(
-        "--g",
-        type=int,
-        default=2,
-        help="Group of pictures sizes to be used for encoding.",
-    )
-    parser.add_argument(
-        "--crf",
-        type=int,
-        default=30,
-        help="Constant rate factors to be used for encoding.",
-    )
-    parser.add_argument(
-        "--tests-data-dir",
-        type=Path,
-        default=None,
-        help=(
-            "When provided, save tests artifacts into the given directory "
-            "(e.g. `--tests-data-dir tests/data` will save to tests/data/{--repo-id})."
-        ),
-    )
-    parser.add_argument(
-        "--dry-run",
-        type=int,
-        default=0,
-        help="If not set to 0, this script won't download or upload anything.",
-    )
-    args = parser.parse_args()
-    encode_datasets(**vars(args))
-
-
-if __name__ == "__main__":
-    main()

lerobot/common/datasets/push_dataset_to_hub/_umi_imagecodecs_numcodecs.py

@@ -1,326 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# imagecodecs/numcodecs.py
-
-# Copyright (c) 2021-2022, Christoph Gohlke
-# All rights reserved.
-#
-# Redistribution and use in source and binary forms, with or without
-# modification, are permitted provided that the following conditions are met:
-#
-# 1. Redistributions of source code must retain the above copyright notice,
-#    this list of conditions and the following disclaimer.
-#
-# 2. Redistributions in binary form must reproduce the above copyright notice,
-#    this list of conditions and the following disclaimer in the documentation
-#    and/or other materials provided with the distribution.
-#
-# 3. Neither the name of the copyright holder nor the names of its
-#    contributors may be used to endorse or promote products derived from
-#    this software without specific prior written permission.
-#
-# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
-# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
-# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
-# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
-# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
-# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
-# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
-# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
-# POSSIBILITY OF SUCH DAMAGE.
-
-# Copied from: https://github.com/real-stanford/universal_manipulation_interface/blob/298776ce251f33b6b3185a98d6e7d1f9ad49168b/diffusion_policy/codecs/imagecodecs_numcodecs.py#L1
-"""Additional numcodecs implemented using imagecodecs."""
-
-__version__ = "2022.9.26"
-
-__all__ = ("register_codecs",)
-
-import imagecodecs
-import numpy
-from numcodecs.abc import Codec
-from numcodecs.registry import get_codec, register_codec
-
-# TODO (azouitine): Remove useless codecs
-
-
-def protective_squeeze(x: numpy.ndarray):
-    """
-    Squeeze dim only if it's not the last dim.
-    Image dim expected to be *, H, W, C
-    """
-    img_shape = x.shape[-3:]
-    if len(x.shape) > 3:
-        n_imgs = numpy.prod(x.shape[:-3])
-        if n_imgs > 1:
-            img_shape = (-1,) + img_shape
-    return x.reshape(img_shape)
-
-
-def get_default_image_compressor(**kwargs):
-    if imagecodecs.JPEGXL:
-        # has JPEGXL
-        this_kwargs = {
-            "effort": 3,
-            "distance": 0.3,
-            # bug in libjxl, invalid codestream for non-lossless
-            # when decoding speed > 1
-            "decodingspeed": 1,
-        }
-        this_kwargs.update(kwargs)
-        return JpegXl(**this_kwargs)
-    else:
-        this_kwargs = {"level": 50}
-        this_kwargs.update(kwargs)
-        return Jpeg2k(**this_kwargs)
-
-
-class Jpeg2k(Codec):
-    """JPEG 2000 codec for numcodecs."""
-
-    codec_id = "imagecodecs_jpeg2k"
-
-    def __init__(
-        self,
-        level=None,
-        codecformat=None,
-        colorspace=None,
-        tile=None,
-        reversible=None,
-        bitspersample=None,
-        resolutions=None,
-        numthreads=None,
-        verbose=0,
-    ):
-        self.level = level
-        self.codecformat = codecformat
-        self.colorspace = colorspace
-        self.tile = None if tile is None else tuple(tile)
-        self.reversible = reversible
-        self.bitspersample = bitspersample
-        self.resolutions = resolutions
-        self.numthreads = numthreads
-        self.verbose = verbose
-
-    def encode(self, buf):
-        buf = protective_squeeze(numpy.asarray(buf))
-        return imagecodecs.jpeg2k_encode(
-            buf,
-            level=self.level,
-            codecformat=self.codecformat,
-            colorspace=self.colorspace,
-            tile=self.tile,
-            reversible=self.reversible,
-            bitspersample=self.bitspersample,
-            resolutions=self.resolutions,
-            numthreads=self.numthreads,
-            verbose=self.verbose,
-        )
-
-    def decode(self, buf, out=None):
-        return imagecodecs.jpeg2k_decode(buf, verbose=self.verbose, numthreads=self.numthreads, out=out)
-
-
-class JpegXl(Codec):
-    """JPEG XL codec for numcodecs."""
-
-    codec_id = "imagecodecs_jpegxl"
-
-    def __init__(
-        self,
-        # encode
-        level=None,
-        effort=None,
-        distance=None,
-        lossless=None,
-        decodingspeed=None,
-        photometric=None,
-        planar=None,
-        usecontainer=None,
-        # decode
-        index=None,
-        keeporientation=None,
-        # both
-        numthreads=None,
-    ):
-        """
-        Return JPEG XL image from numpy array.
-        Float must be in nominal range 0..1.
-
-        Currently L, LA, RGB, RGBA images are supported in contig mode.
-        Extra channels are only supported for grayscale images in planar mode.
-
-        Parameters
-        ----------
-        level : Default to None, i.e. not overwriting lossess and decodingspeed options.
-            When < 0: Use lossless compression
-            When in [0,1,2,3,4]: Sets the decoding speed tier for the provided options.
-                Minimum is 0 (slowest to decode, best quality/density), and maximum
-                is 4 (fastest to decode, at the cost of some quality/density).
-        effort : Default to 3.
-            Sets encoder effort/speed level without affecting decoding speed.
-            Valid values are, from faster to slower speed: 1:lightning 2:thunder
-                3:falcon 4:cheetah 5:hare 6:wombat 7:squirrel 8:kitten 9:tortoise.
-            Speed: lightning, thunder, falcon, cheetah, hare, wombat, squirrel, kitten, tortoise
-            control the encoder effort in ascending order.
-            This also affects memory usage: using lower effort will typically reduce memory
-            consumption during encoding.
-            lightning and thunder are fast modes useful for lossless mode (modular).
-            falcon disables all of the following tools.
-            cheetah enables coefficient reordering, context clustering, and heuristics for selecting DCT sizes and quantization steps.
-            hare enables Gaborish filtering, chroma from luma, and an initial estimate of quantization steps.
-            wombat enables error diffusion quantization and full DCT size selection heuristics.
-            squirrel (default) enables dots, patches, and spline detection, and full context clustering.
-            kitten optimizes the adaptive quantization for a psychovisual metric.
-            tortoise enables a more thorough adaptive quantization search.
-        distance : Default to 1.0
-            Sets the distance level for lossy compression: target max butteraugli distance,
-            lower = higher quality. Range: 0 .. 15. 0.0 = mathematically lossless
-            (however, use JxlEncoderSetFrameLossless instead to use true lossless,
-            as setting distance to 0 alone is not the only requirement).
-            1.0 = visually lossless. Recommended range: 0.5 .. 3.0.
-        lossess : Default to False.
-            Use lossess encoding.
-        decodingspeed : Default to 0.
-            Duplicate to level. [0,4]
-        photometric : Return JxlColorSpace value.
-            Default logic is quite complicated but works most of the time.
-            Accepted value:
-                int: [-1,3]
-                str: ['RGB',
-                    'WHITEISZERO', 'MINISWHITE',
-                    'BLACKISZERO', 'MINISBLACK', 'GRAY',
-                    'XYB', 'KNOWN']
-        planar : Enable multi-channel mode.
-            Default to false.
-        usecontainer :
-            Forces the encoder to use the box-based container format (BMFF)
-            even when not necessary.
-            When using JxlEncoderUseBoxes, JxlEncoderStoreJPEGMetadata or
-            JxlEncoderSetCodestreamLevel with level 10, the encoder will
-            automatically also use the container format, it is not necessary
-            to use JxlEncoderUseContainer for those use cases.
-            By default this setting is disabled.
-        index : Selectively decode frames for animation.
-            Default to 0, decode all frames.
-            When set to > 0, decode that frame index only.
-        keeporientation :
-            Enables or disables preserving of as-in-bitstream pixeldata orientation.
-            Some images are encoded with an Orientation tag indicating that the
-            decoder must perform a rotation and/or mirroring to the encoded image data.
-
-            If skip_reorientation is JXL_FALSE (the default): the decoder will apply
-            the transformation from the orientation setting, hence rendering the image
-            according to its specified intent. When producing a JxlBasicInfo, the decoder
-            will always set the orientation field to JXL_ORIENT_IDENTITY (matching the
-            returned pixel data) and also align xsize and ysize so that they correspond
-            to the width and the height of the returned pixel data.
-
-            If skip_reorientation is JXL_TRUE: the decoder will skip applying the
-            transformation from the orientation setting, returning the image in
-            the as-in-bitstream pixeldata orientation. This may be faster to decode
-            since the decoder doesnt have to apply the transformation, but can
-            cause wrong display of the image if the orientation tag is not correctly
-            taken into account by the user.
-
-            By default, this option is disabled, and the returned pixel data is
-            re-oriented according to the images Orientation setting.
-        threads : Default to 1.
-            If <= 0, use all cores.
-            If > 32, clipped to 32.
-        """
-
-        self.level = level
-        self.effort = effort
-        self.distance = distance
-        self.lossless = bool(lossless)
-        self.decodingspeed = decodingspeed
-        self.photometric = photometric
-        self.planar = planar
-        self.usecontainer = usecontainer
-        self.index = index
-        self.keeporientation = keeporientation
-        self.numthreads = numthreads
-
-    def encode(self, buf):
-        # TODO: only squeeze all but last dim
-        buf = protective_squeeze(numpy.asarray(buf))
-        return imagecodecs.jpegxl_encode(
-            buf,
-            level=self.level,
-            effort=self.effort,
-            distance=self.distance,
-            lossless=self.lossless,
-            decodingspeed=self.decodingspeed,
-            photometric=self.photometric,
-            planar=self.planar,
-            usecontainer=self.usecontainer,
-            numthreads=self.numthreads,
-        )
-
-    def decode(self, buf, out=None):
-        return imagecodecs.jpegxl_decode(
-            buf,
-            index=self.index,
-            keeporientation=self.keeporientation,
-            numthreads=self.numthreads,
-            out=out,
-        )
-
-
-def _flat(out):
-    """Return numpy array as contiguous view of bytes if possible."""
-    if out is None:
-        return None
-    view = memoryview(out)
-    if view.readonly or not view.contiguous:
-        return None
-    return view.cast("B")
-
-
-def register_codecs(codecs=None, force=False, verbose=True):
-    """Register codecs in this module with numcodecs."""
-    for name, cls in globals().items():
-        if not hasattr(cls, "codec_id") or name == "Codec":
-            continue
-        if codecs is not None and cls.codec_id not in codecs:
-            continue
-        try:
-            try:  # noqa: SIM105
-                get_codec({"id": cls.codec_id})
-            except TypeError:
-                # registered, but failed
-                pass
-        except ValueError:
-            # not registered yet
-            pass
-        else:
-            if not force:
-                if verbose:
-                    log_warning(f"numcodec {cls.codec_id!r} already registered")
-                continue
-            if verbose:
-                log_warning(f"replacing registered numcodec {cls.codec_id!r}")
-        register_codec(cls)
-
-
-def log_warning(msg, *args, **kwargs):
-    """Log message with level WARNING."""
-    import logging
-
-    logging.getLogger(__name__).warning(msg, *args, **kwargs)

lerobot/common/datasets/push_dataset_to_hub/aloha_hdf5_format.py

@@ -1,233 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""
-Contains utilities to process raw data format of HDF5 files like in: https://github.com/tonyzhaozh/act
-"""
-
-import gc
-import shutil
-from pathlib import Path
-
-import h5py
-import numpy as np
-import torch
-import tqdm
-from datasets import Dataset, Features, Image, Sequence, Value
-from PIL import Image as PILImage
-
-from lerobot.common.datasets.lerobot_dataset import CODEBASE_VERSION
-from lerobot.common.datasets.push_dataset_to_hub.utils import (
-    calculate_episode_data_index,
-    concatenate_episodes,
-    get_default_encoding,
-    save_images_concurrently,
-)
-from lerobot.common.datasets.utils import (
-    hf_transform_to_torch,
-)
-from lerobot.common.datasets.video_utils import VideoFrame, encode_video_frames
-
-
-def get_cameras(hdf5_data):
-    # ignore depth channel, not currently handled
-    # TODO(rcadene): add depth
-    rgb_cameras = [key for key in hdf5_data["/observations/images"].keys() if "depth" not in key]  # noqa: SIM118
-    return rgb_cameras
-
-
-def check_format(raw_dir) -> bool:
-    # only frames from simulation are uncompressed
-    compressed_images = "sim" not in raw_dir.name
-
-    hdf5_paths = list(raw_dir.glob("episode_*.hdf5"))
-    assert len(hdf5_paths) != 0
-    for hdf5_path in hdf5_paths:
-        with h5py.File(hdf5_path, "r") as data:
-            assert "/action" in data
-            assert "/observations/qpos" in data
-
-            assert data["/action"].ndim == 2
-            assert data["/observations/qpos"].ndim == 2
-
-            num_frames = data["/action"].shape[0]
-            assert num_frames == data["/observations/qpos"].shape[0]
-
-            for camera in get_cameras(data):
-                assert num_frames == data[f"/observations/images/{camera}"].shape[0]
-
-                if compressed_images:
-                    assert data[f"/observations/images/{camera}"].ndim == 2
-                else:
-                    assert data[f"/observations/images/{camera}"].ndim == 4
-                    b, h, w, c = data[f"/observations/images/{camera}"].shape
-                    assert c < h and c < w, f"Expect (h,w,c) image format but ({h=},{w=},{c=}) provided."
-
-
-def load_from_raw(
-    raw_dir: Path,
-    videos_dir: Path,
-    fps: int,
-    video: bool,
-    episodes: list[int] | None = None,
-    encoding: dict | None = None,
-):
-    # only frames from simulation are uncompressed
-    compressed_images = "sim" not in raw_dir.name
-
-    hdf5_files = sorted(raw_dir.glob("episode_*.hdf5"))
-    num_episodes = len(hdf5_files)
-
-    ep_dicts = []
-    ep_ids = episodes if episodes else range(num_episodes)
-    for ep_idx in tqdm.tqdm(ep_ids):
-        ep_path = hdf5_files[ep_idx]
-        with h5py.File(ep_path, "r") as ep:
-            num_frames = ep["/action"].shape[0]
-
-            # last step of demonstration is considered done
-            done = torch.zeros(num_frames, dtype=torch.bool)
-            done[-1] = True
-
-            state = torch.from_numpy(ep["/observations/qpos"][:])
-            action = torch.from_numpy(ep["/action"][:])
-            if "/observations/qvel" in ep:
-                velocity = torch.from_numpy(ep["/observations/qvel"][:])
-            if "/observations/effort" in ep:
-                effort = torch.from_numpy(ep["/observations/effort"][:])
-
-            ep_dict = {}
-
-            for camera in get_cameras(ep):
-                img_key = f"observation.images.{camera}"
-
-                if compressed_images:
-                    import cv2
-
-                    # load one compressed image after the other in RAM and uncompress
-                    imgs_array = []
-                    for data in ep[f"/observations/images/{camera}"]:
-                        imgs_array.append(cv2.imdecode(data, 1))
-                    imgs_array = np.array(imgs_array)
-
-                else:
-                    # load all images in RAM
-                    imgs_array = ep[f"/observations/images/{camera}"][:]
-
-                if video:
-                    # save png images in temporary directory
-                    tmp_imgs_dir = videos_dir / "tmp_images"
-                    save_images_concurrently(imgs_array, tmp_imgs_dir)
-
-                    # encode images to a mp4 video
-                    fname = f"{img_key}_episode_{ep_idx:06d}.mp4"
-                    video_path = videos_dir / fname
-                    encode_video_frames(tmp_imgs_dir, video_path, fps, **(encoding or {}))
-
-                    # clean temporary images directory
-                    shutil.rmtree(tmp_imgs_dir)
-
-                    # store the reference to the video frame
-                    ep_dict[img_key] = [
-                        {"path": f"videos/{fname}", "timestamp": i / fps} for i in range(num_frames)
-                    ]
-                else:
-                    ep_dict[img_key] = [PILImage.fromarray(x) for x in imgs_array]
-
-            ep_dict["observation.state"] = state
-            if "/observations/velocity" in ep:
-                ep_dict["observation.velocity"] = velocity
-            if "/observations/effort" in ep:
-                ep_dict["observation.effort"] = effort
-            ep_dict["action"] = action
-            ep_dict["episode_index"] = torch.tensor([ep_idx] * num_frames)
-            ep_dict["frame_index"] = torch.arange(0, num_frames, 1)
-            ep_dict["timestamp"] = torch.arange(0, num_frames, 1) / fps
-            ep_dict["next.done"] = done
-            # TODO(rcadene): add reward and success by computing them in sim
-
-            assert isinstance(ep_idx, int)
-            ep_dicts.append(ep_dict)
-
-        gc.collect()
-
-    data_dict = concatenate_episodes(ep_dicts)
-
-    total_frames = data_dict["frame_index"].shape[0]
-    data_dict["index"] = torch.arange(0, total_frames, 1)
-    return data_dict
-
-
-def to_hf_dataset(data_dict, video) -> Dataset:
-    features = {}
-
-    keys = [key for key in data_dict if "observation.images." in key]
-    for key in keys:
-        if video:
-            features[key] = VideoFrame()
-        else:
-            features[key] = Image()
-
-    features["observation.state"] = Sequence(
-        length=data_dict["observation.state"].shape[1], feature=Value(dtype="float32", id=None)
-    )
-    if "observation.velocity" in data_dict:
-        features["observation.velocity"] = Sequence(
-            length=data_dict["observation.velocity"].shape[1], feature=Value(dtype="float32", id=None)
-        )
-    if "observation.effort" in data_dict:
-        features["observation.effort"] = Sequence(
-            length=data_dict["observation.effort"].shape[1], feature=Value(dtype="float32", id=None)
-        )
-    features["action"] = Sequence(
-        length=data_dict["action"].shape[1], feature=Value(dtype="float32", id=None)
-    )
-    features["episode_index"] = Value(dtype="int64", id=None)
-    features["frame_index"] = Value(dtype="int64", id=None)
-    features["timestamp"] = Value(dtype="float32", id=None)
-    features["next.done"] = Value(dtype="bool", id=None)
-    features["index"] = Value(dtype="int64", id=None)
-
-    hf_dataset = Dataset.from_dict(data_dict, features=Features(features))
-    hf_dataset.set_transform(hf_transform_to_torch)
-    return hf_dataset
-
-
-def from_raw_to_lerobot_format(
-    raw_dir: Path,
-    videos_dir: Path,
-    fps: int | None = None,
-    video: bool = True,
-    episodes: list[int] | None = None,
-    encoding: dict | None = None,
-):
-    # sanity check
-    check_format(raw_dir)
-
-    if fps is None:
-        fps = 50
-
-    data_dict = load_from_raw(raw_dir, videos_dir, fps, video, episodes, encoding)
-    hf_dataset = to_hf_dataset(data_dict, video)
-    episode_data_index = calculate_episode_data_index(hf_dataset)
-    info = {
-        "codebase_version": CODEBASE_VERSION,
-        "fps": fps,
-        "video": video,
-    }
-    if video:
-        info["encoding"] = get_default_encoding()
-
-    return hf_dataset, episode_data_index, info

lerobot/common/datasets/push_dataset_to_hub/cam_png_format.py

@@ -1,107 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""
-Contains utilities to process raw data format of png images files recorded with capture_camera_feed.py
-"""
-
-from pathlib import Path
-
-import torch
-from datasets import Dataset, Features, Image, Value
-from PIL import Image as PILImage
-
-from lerobot.common.datasets.lerobot_dataset import CODEBASE_VERSION
-from lerobot.common.datasets.push_dataset_to_hub.utils import (
-    calculate_episode_data_index,
-    concatenate_episodes,
-)
-from lerobot.common.datasets.utils import hf_transform_to_torch
-from lerobot.common.datasets.video_utils import VideoFrame
-
-
-def check_format(raw_dir: Path) -> bool:
-    image_paths = list(raw_dir.glob("frame_*.png"))
-    if len(image_paths) == 0:
-        raise ValueError
-
-
-def load_from_raw(raw_dir: Path, fps: int, episodes: list[int] | None = None):
-    if episodes is not None:
-        # TODO(aliberts): add support for multi-episodes.
-        raise NotImplementedError()
-
-    ep_dict = {}
-    ep_idx = 0
-
-    image_paths = sorted(raw_dir.glob("frame_*.png"))
-    num_frames = len(image_paths)
-
-    ep_dict["observation.image"] = [PILImage.open(x) for x in image_paths]
-    ep_dict["episode_index"] = torch.tensor([ep_idx] * num_frames)
-    ep_dict["frame_index"] = torch.arange(0, num_frames, 1)
-    ep_dict["timestamp"] = torch.arange(0, num_frames, 1) / fps
-
-    ep_dicts = [ep_dict]
-    data_dict = concatenate_episodes(ep_dicts)
-    total_frames = data_dict["frame_index"].shape[0]
-    data_dict["index"] = torch.arange(0, total_frames, 1)
-    return data_dict
-
-
-def to_hf_dataset(data_dict, video) -> Dataset:
-    features = {}
-    if video:
-        features["observation.image"] = VideoFrame()
-    else:
-        features["observation.image"] = Image()
-
-    features["episode_index"] = Value(dtype="int64", id=None)
-    features["frame_index"] = Value(dtype="int64", id=None)
-    features["timestamp"] = Value(dtype="float32", id=None)
-    features["index"] = Value(dtype="int64", id=None)
-
-    hf_dataset = Dataset.from_dict(data_dict, features=Features(features))
-    hf_dataset.set_transform(hf_transform_to_torch)
-    return hf_dataset
-
-
-def from_raw_to_lerobot_format(
-    raw_dir: Path,
-    videos_dir: Path,
-    fps: int | None = None,
-    video: bool = True,
-    episodes: list[int] | None = None,
-    encoding: dict | None = None,
-):
-    if video or episodes or encoding is not None:
-        # TODO(aliberts): support this
-        raise NotImplementedError
-
-    # sanity check
-    check_format(raw_dir)
-
-    if fps is None:
-        fps = 30
-
-    data_dict = load_from_raw(raw_dir, videos_dir, fps, video, episodes)
-    hf_dataset = to_hf_dataset(data_dict, video)
-    episode_data_index = calculate_episode_data_index(hf_dataset)
-    info = {
-        "codebase_version": CODEBASE_VERSION,
-        "fps": fps,
-        "video": video,
-    }
-    return hf_dataset, episode_data_index, info

lerobot/common/datasets/push_dataset_to_hub/dora_parquet_format.py

@@ -1,233 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""
-Contains utilities to process raw data format from dora-record
-"""
-
-import re
-import warnings
-from pathlib import Path
-
-import pandas as pd
-import torch
-from datasets import Dataset, Features, Image, Sequence, Value
-
-from lerobot.common.datasets.lerobot_dataset import CODEBASE_VERSION
-from lerobot.common.datasets.push_dataset_to_hub.utils import calculate_episode_data_index
-from lerobot.common.datasets.utils import (
-    hf_transform_to_torch,
-)
-from lerobot.common.datasets.video_utils import VideoFrame
-
-
-def check_format(raw_dir) -> bool:
-    assert raw_dir.exists()
-
-    leader_file = list(raw_dir.glob("*.parquet"))
-    if len(leader_file) == 0:
-        raise ValueError(f"Missing parquet files in '{raw_dir}'")
-    return True
-
-
-def load_from_raw(raw_dir: Path, videos_dir: Path, fps: int, video: bool, episodes: list[int] | None = None):
-    # Load data stream that will be used as reference for the timestamps synchronization
-    reference_files = list(raw_dir.glob("observation.images.cam_*.parquet"))
-    if len(reference_files) == 0:
-        raise ValueError(f"Missing reference files for camera, starting with  in '{raw_dir}'")
-    # select first camera in alphanumeric order
-    reference_key = sorted(reference_files)[0].stem
-    reference_df = pd.read_parquet(raw_dir / f"{reference_key}.parquet")
-    reference_df = reference_df[["timestamp_utc", reference_key]]
-
-    # Merge all data stream using nearest backward strategy
-    df = reference_df
-    for path in raw_dir.glob("*.parquet"):
-        key = path.stem  # action or observation.state or ...
-        if key == reference_key:
-            continue
-        if "failed_episode_index" in key:
-            # TODO(rcadene): add support for removing episodes that are tagged as "failed"
-            continue
-        modality_df = pd.read_parquet(path)
-        modality_df = modality_df[["timestamp_utc", key]]
-        df = pd.merge_asof(
-            df,
-            modality_df,
-            on="timestamp_utc",
-            # "nearest" is the best option over "backward", since the latter can desynchronizes camera timestamps by
-            # matching timestamps that are too far apart, in order to fit the backward constraints. It's not the case for "nearest".
-            # However, note that "nearest" might synchronize the reference camera with other cameras on slightly future timestamps.
-            # are too far apart.
-            direction="nearest",
-            tolerance=pd.Timedelta(f"{1 / fps} seconds"),
-        )
-    # Remove rows with episode_index -1 which indicates data that correspond to in-between episodes
-    df = df[df["episode_index"] != -1]
-
-    image_keys = [key for key in df if "observation.images." in key]
-
-    def get_episode_index(row):
-        episode_index_per_cam = {}
-        for key in image_keys:
-            path = row[key][0]["path"]
-            match = re.search(r"_(\d{6}).mp4", path)
-            if not match:
-                raise ValueError(path)
-            episode_index = int(match.group(1))
-            episode_index_per_cam[key] = episode_index
-        if len(set(episode_index_per_cam.values())) != 1:
-            raise ValueError(
-                f"All cameras are expected to belong to the same episode, but getting {episode_index_per_cam}"
-            )
-        return episode_index
-
-    df["episode_index"] = df.apply(get_episode_index, axis=1)
-
-    # dora only use arrays, so single values are encapsulated into a list
-    df["frame_index"] = df.groupby("episode_index").cumcount()
-    df = df.reset_index()
-    df["index"] = df.index
-
-    # set 'next.done' to True for the last frame of each episode
-    df["next.done"] = False
-    df.loc[df.groupby("episode_index").tail(1).index, "next.done"] = True
-
-    df["timestamp"] = df["timestamp_utc"].map(lambda x: x.timestamp())
-    # each episode starts with timestamp 0 to match the ones from the video
-    df["timestamp"] = df.groupby("episode_index")["timestamp"].transform(lambda x: x - x.iloc[0])
-
-    del df["timestamp_utc"]
-
-    # sanity check
-    has_nan = df.isna().any().any()
-    if has_nan:
-        raise ValueError("Dataset contains Nan values.")
-
-    # sanity check episode indices go from 0 to n-1
-    ep_ids = [ep_idx for ep_idx, _ in df.groupby("episode_index")]
-    expected_ep_ids = list(range(df["episode_index"].max() + 1))
-    if ep_ids != expected_ep_ids:
-        raise ValueError(f"Episodes indices go from {ep_ids} instead of {expected_ep_ids}")
-
-    # Create symlink to raw videos directory (that needs to be absolute not relative)
-    videos_dir.parent.mkdir(parents=True, exist_ok=True)
-    videos_dir.symlink_to((raw_dir / "videos").absolute())
-
-    # sanity check the video paths are well formatted
-    for key in df:
-        if "observation.images." not in key:
-            continue
-        for ep_idx in ep_ids:
-            video_path = videos_dir / f"{key}_episode_{ep_idx:06d}.mp4"
-            if not video_path.exists():
-                raise ValueError(f"Video file not found in {video_path}")
-
-    data_dict = {}
-    for key in df:
-        # is video frame
-        if "observation.images." in key:
-            # we need `[0] because dora only use arrays, so single values are encapsulated into a list.
-            # it is the case for video_frame dictionary = [{"path": ..., "timestamp": ...}]
-            data_dict[key] = [video_frame[0] for video_frame in df[key].values]
-
-            # sanity check the video path is well formatted
-            video_path = videos_dir.parent / data_dict[key][0]["path"]
-            if not video_path.exists():
-                raise ValueError(f"Video file not found in {video_path}")
-        # is number
-        elif df[key].iloc[0].ndim == 0 or df[key].iloc[0].shape[0] == 1:
-            data_dict[key] = torch.from_numpy(df[key].values)
-        # is vector
-        elif df[key].iloc[0].shape[0] > 1:
-            data_dict[key] = torch.stack([torch.from_numpy(x.copy()) for x in df[key].values])
-        else:
-            raise ValueError(key)
-
-    return data_dict
-
-
-def to_hf_dataset(data_dict, video) -> Dataset:
-    features = {}
-
-    keys = [key for key in data_dict if "observation.images." in key]
-    for key in keys:
-        if video:
-            features[key] = VideoFrame()
-        else:
-            features[key] = Image()
-
-    features["observation.state"] = Sequence(
-        length=data_dict["observation.state"].shape[1], feature=Value(dtype="float32", id=None)
-    )
-    if "observation.velocity" in data_dict:
-        features["observation.velocity"] = Sequence(
-            length=data_dict["observation.velocity"].shape[1], feature=Value(dtype="float32", id=None)
-        )
-    if "observation.effort" in data_dict:
-        features["observation.effort"] = Sequence(
-            length=data_dict["observation.effort"].shape[1], feature=Value(dtype="float32", id=None)
-        )
-    features["action"] = Sequence(
-        length=data_dict["action"].shape[1], feature=Value(dtype="float32", id=None)
-    )
-    features["episode_index"] = Value(dtype="int64", id=None)
-    features["frame_index"] = Value(dtype="int64", id=None)
-    features["timestamp"] = Value(dtype="float32", id=None)
-    features["next.done"] = Value(dtype="bool", id=None)
-    features["index"] = Value(dtype="int64", id=None)
-
-    hf_dataset = Dataset.from_dict(data_dict, features=Features(features))
-    hf_dataset.set_transform(hf_transform_to_torch)
-    return hf_dataset
-
-
-def from_raw_to_lerobot_format(
-    raw_dir: Path,
-    videos_dir: Path,
-    fps: int | None = None,
-    video: bool = True,
-    episodes: list[int] | None = None,
-    encoding: dict | None = None,
-):
-    # sanity check
-    check_format(raw_dir)
-
-    if fps is None:
-        fps = 30
-    else:
-        raise NotImplementedError()
-
-    if not video:
-        raise NotImplementedError()
-
-    if encoding is not None:
-        warnings.warn(
-            "Video encoding is currently done outside of LeRobot for the dora_parquet format.",
-            stacklevel=1,
-        )
-
-    data_df = load_from_raw(raw_dir, videos_dir, fps, episodes)
-    hf_dataset = to_hf_dataset(data_df, video)
-    episode_data_index = calculate_episode_data_index(hf_dataset)
-    info = {
-        "codebase_version": CODEBASE_VERSION,
-        "fps": fps,
-        "video": video,
-    }
-    if video:
-        info["encoding"] = "unknown"
-
-    return hf_dataset, episode_data_index, info

lerobot/common/datasets/push_dataset_to_hub/openx_rlds_format.py

@@ -1,312 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""
-For all datasets in the RLDS format.
-For https://github.com/google-deepmind/open_x_embodiment (OPENX) datasets.
-
-NOTE: You need to install tensorflow and tensorflow_datasets before running this script.
-
-Example:
-    python lerobot/scripts/push_dataset_to_hub.py \
-        --raw-dir /path/to/data/bridge_dataset/1.0.0/ \
-        --repo-id your_hub/sampled_bridge_data_v2 \
-        --raw-format rlds \
-        --episodes 3 4 5 8 9
-
-Exact dataset fps defined in openx/config.py, obtained from:
-    https://docs.google.com/spreadsheets/d/1rPBD77tk60AEIGZrGSODwyyzs5FgCU9Uz3h-3_t2A9g/edit?gid=0#gid=0&range=R:R
-"""
-
-import shutil
-from pathlib import Path
-
-import numpy as np
-import tensorflow as tf
-import tensorflow_datasets as tfds
-import torch
-import tqdm
-from datasets import Dataset, Features, Image, Sequence, Value
-from PIL import Image as PILImage
-
-from lerobot.common.datasets.lerobot_dataset import CODEBASE_VERSION
-from lerobot.common.datasets.push_dataset_to_hub.utils import (
-    calculate_episode_data_index,
-    concatenate_episodes,
-    get_default_encoding,
-    save_images_concurrently,
-)
-from lerobot.common.datasets.utils import (
-    hf_transform_to_torch,
-)
-from lerobot.common.datasets.video_utils import VideoFrame, encode_video_frames
-
-np.set_printoptions(precision=2)
-
-
-def tf_to_torch(data):
-    return torch.from_numpy(data.numpy())
-
-
-def tf_img_convert(img):
-    if img.dtype == tf.string:
-        img = tf.io.decode_image(img, expand_animations=False, dtype=tf.uint8)
-    elif img.dtype != tf.uint8:
-        raise ValueError(f"Unsupported image dtype: found with dtype {img.dtype}")
-    return img.numpy()
-
-
-def _broadcast_metadata_rlds(i: tf.Tensor, traj: dict) -> dict:
-    """
-    In the RLDS format, each trajectory has some top-level metadata that is explicitly separated out, and a "steps"
-    entry. This function moves the "steps" entry to the top level, broadcasting any metadata to the length of the
-    trajectory. This function also adds the extra metadata fields `_len`, `_traj_index`, and `_frame_index`.
-
-    NOTE: adapted from DLimp library https://github.com/kvablack/dlimp/
-    """
-    steps = traj.pop("steps")
-
-    traj_len = tf.shape(tf.nest.flatten(steps)[0])[0]
-
-    # broadcast metadata to the length of the trajectory
-    metadata = tf.nest.map_structure(lambda x: tf.repeat(x, traj_len), traj)
-
-    # put steps back in
-    assert "traj_metadata" not in steps
-    traj = {**steps, "traj_metadata": metadata}
-
-    assert "_len" not in traj
-    assert "_traj_index" not in traj
-    assert "_frame_index" not in traj
-    traj["_len"] = tf.repeat(traj_len, traj_len)
-    traj["_traj_index"] = tf.repeat(i, traj_len)
-    traj["_frame_index"] = tf.range(traj_len)
-
-    return traj
-
-
-def load_from_raw(
-    raw_dir: Path,
-    videos_dir: Path,
-    fps: int,
-    video: bool,
-    episodes: list[int] | None = None,
-    encoding: dict | None = None,
-):
-    """
-    Args:
-        raw_dir (Path): _description_
-        videos_dir (Path): _description_
-        fps (int): _description_
-        video (bool): _description_
-        episodes (list[int] | None, optional): _description_. Defaults to None.
-    """
-    ds_builder = tfds.builder_from_directory(str(raw_dir))
-    dataset = ds_builder.as_dataset(
-        split="all",
-        decoders={"steps": tfds.decode.SkipDecoding()},
-    )
-
-    dataset_info = ds_builder.info
-    print("dataset_info: ", dataset_info)
-
-    ds_length = len(dataset)
-    dataset = dataset.take(ds_length)
-    # "flatten" the dataset as such we can apply trajectory level map() easily
-    # each [obs][key] has a shape of (frame_size, ...)
-    dataset = dataset.enumerate().map(_broadcast_metadata_rlds)
-
-    # we will apply the standardization transform if the dataset_name is provided
-    # if the dataset name is not provided and the goal is to convert any rlds formatted dataset
-    # search for 'image' keys in the observations
-    image_keys = []
-    state_keys = []
-    observation_info = dataset_info.features["steps"]["observation"]
-    for key in observation_info:
-        # check whether the key is for an image or a vector observation
-        if len(observation_info[key].shape) == 3:
-            # only adding uint8 images discards depth images
-            if observation_info[key].dtype == tf.uint8:
-                image_keys.append(key)
-        else:
-            state_keys.append(key)
-
-    lang_key = "language_instruction" if "language_instruction" in dataset.element_spec else None
-
-    print(" - image_keys: ", image_keys)
-    print(" - lang_key: ", lang_key)
-
-    it = iter(dataset)
-
-    ep_dicts = []
-    # Init temp path to save ep_dicts in case of crash
-    tmp_ep_dicts_dir = videos_dir.parent.joinpath("ep_dicts")
-    tmp_ep_dicts_dir.mkdir(parents=True, exist_ok=True)
-
-    # check if ep_dicts have already been saved in /tmp
-    starting_ep_idx = 0
-    saved_ep_dicts = [ep.__str__() for ep in tmp_ep_dicts_dir.iterdir()]
-    if len(saved_ep_dicts) > 0:
-        saved_ep_dicts.sort()
-        # get last ep_idx number
-        starting_ep_idx = int(saved_ep_dicts[-1][-13:-3]) + 1
-        for i in range(starting_ep_idx):
-            episode = next(it)
-            ep_dicts.append(torch.load(saved_ep_dicts[i]))
-
-    # if we user specified episodes, skip the ones not in the list
-    if episodes is not None:
-        if ds_length == 0:
-            raise ValueError("No episodes found.")
-        # convert episodes index to sorted list
-        episodes = sorted(episodes)
-
-    for ep_idx in tqdm.tqdm(range(starting_ep_idx, ds_length)):
-        episode = next(it)
-
-        # if user specified episodes, skip the ones not in the list
-        if episodes is not None:
-            if len(episodes) == 0:
-                break
-            if ep_idx == episodes[0]:
-                # process this episode
-                print(" selecting episode idx: ", ep_idx)
-                episodes.pop(0)
-            else:
-                continue  # skip
-
-        num_frames = episode["action"].shape[0]
-
-        ep_dict = {}
-        for key in state_keys:
-            ep_dict[f"observation.{key}"] = tf_to_torch(episode["observation"][key])
-
-        ep_dict["action"] = tf_to_torch(episode["action"])
-        ep_dict["next.reward"] = tf_to_torch(episode["reward"]).float()
-        ep_dict["next.done"] = tf_to_torch(episode["is_last"])
-        ep_dict["is_terminal"] = tf_to_torch(episode["is_terminal"])
-        ep_dict["is_first"] = tf_to_torch(episode["is_first"])
-        ep_dict["discount"] = tf_to_torch(episode["discount"])
-
-        # If lang_key is present, convert the entire tensor at once
-        if lang_key is not None:
-            ep_dict["language_instruction"] = [x.numpy().decode("utf-8") for x in episode[lang_key]]
-
-        ep_dict["timestamp"] = torch.arange(0, num_frames, 1) / fps
-        ep_dict["episode_index"] = torch.tensor([ep_idx] * num_frames)
-        ep_dict["frame_index"] = torch.arange(0, num_frames, 1)
-
-        image_array_dict = {key: [] for key in image_keys}
-
-        for im_key in image_keys:
-            imgs = episode["observation"][im_key]
-            image_array_dict[im_key] = [tf_img_convert(img) for img in imgs]
-
-        # loop through all cameras
-        for im_key in image_keys:
-            img_key = f"observation.images.{im_key}"
-            imgs_array = image_array_dict[im_key]
-            imgs_array = np.array(imgs_array)
-            if video:
-                # save png images in temporary directory
-                tmp_imgs_dir = videos_dir / "tmp_images"
-                save_images_concurrently(imgs_array, tmp_imgs_dir)
-
-                # encode images to a mp4 video
-                fname = f"{img_key}_episode_{ep_idx:06d}.mp4"
-                video_path = videos_dir / fname
-                encode_video_frames(tmp_imgs_dir, video_path, fps, **(encoding or {}))
-
-                # clean temporary images directory
-                shutil.rmtree(tmp_imgs_dir)
-
-                # store the reference to the video frame
-                ep_dict[img_key] = [
-                    {"path": f"videos/{fname}", "timestamp": i / fps} for i in range(num_frames)
-                ]
-            else:
-                ep_dict[img_key] = [PILImage.fromarray(x) for x in imgs_array]
-
-        path_ep_dict = tmp_ep_dicts_dir.joinpath(
-            "ep_dict_" + "0" * (10 - len(str(ep_idx))) + str(ep_idx) + ".pt"
-        )
-        torch.save(ep_dict, path_ep_dict)
-
-        ep_dicts.append(ep_dict)
-
-    data_dict = concatenate_episodes(ep_dicts)
-
-    total_frames = data_dict["frame_index"].shape[0]
-    data_dict["index"] = torch.arange(0, total_frames, 1)
-    return data_dict
-
-
-def to_hf_dataset(data_dict, video) -> Dataset:
-    features = {}
-
-    for key in data_dict:
-        # check if vector state obs
-        if key.startswith("observation.") and "observation.images." not in key:
-            features[key] = Sequence(length=data_dict[key].shape[1], feature=Value(dtype="float32", id=None))
-        # check if image obs
-        elif "observation.images." in key:
-            if video:
-                features[key] = VideoFrame()
-            else:
-                features[key] = Image()
-
-    if "language_instruction" in data_dict:
-        features["language_instruction"] = Value(dtype="string", id=None)
-
-    features["action"] = Sequence(
-        length=data_dict["action"].shape[1], feature=Value(dtype="float32", id=None)
-    )
-
-    features["is_terminal"] = Value(dtype="bool", id=None)
-    features["is_first"] = Value(dtype="bool", id=None)
-    features["discount"] = Value(dtype="float32", id=None)
-
-    features["episode_index"] = Value(dtype="int64", id=None)
-    features["frame_index"] = Value(dtype="int64", id=None)
-    features["timestamp"] = Value(dtype="float32", id=None)
-    features["next.reward"] = Value(dtype="float32", id=None)
-    features["next.done"] = Value(dtype="bool", id=None)
-    features["index"] = Value(dtype="int64", id=None)
-
-    hf_dataset = Dataset.from_dict(data_dict, features=Features(features))
-    hf_dataset.set_transform(hf_transform_to_torch)
-    return hf_dataset
-
-
-def from_raw_to_lerobot_format(
-    raw_dir: Path,
-    videos_dir: Path,
-    fps: int | None = None,
-    video: bool = True,
-    episodes: list[int] | None = None,
-    encoding: dict | None = None,
-):
-    data_dict = load_from_raw(raw_dir, videos_dir, fps, video, episodes, encoding)
-    hf_dataset = to_hf_dataset(data_dict, video)
-    episode_data_index = calculate_episode_data_index(hf_dataset)
-    info = {
-        "codebase_version": CODEBASE_VERSION,
-        "fps": fps,
-        "video": video,
-    }
-    if video:
-        info["encoding"] = get_default_encoding()
-
-    return hf_dataset, episode_data_index, info

lerobot/common/datasets/push_dataset_to_hub/pusht_zarr_format.py

@@ -1,275 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Process zarr files formatted like in: https://github.com/real-stanford/diffusion_policy"""
-
-import shutil
-from pathlib import Path
-
-import numpy as np
-import torch
-import tqdm
-import zarr
-from datasets import Dataset, Features, Image, Sequence, Value
-from PIL import Image as PILImage
-
-from lerobot.common.datasets.lerobot_dataset import CODEBASE_VERSION
-from lerobot.common.datasets.push_dataset_to_hub.utils import (
-    calculate_episode_data_index,
-    concatenate_episodes,
-    get_default_encoding,
-    save_images_concurrently,
-)
-from lerobot.common.datasets.utils import (
-    hf_transform_to_torch,
-)
-from lerobot.common.datasets.video_utils import VideoFrame, encode_video_frames
-
-
-def check_format(raw_dir):
-    zarr_path = raw_dir / "pusht_cchi_v7_replay.zarr"
-    zarr_data = zarr.open(zarr_path, mode="r")
-
-    required_datasets = {
-        "data/action",
-        "data/img",
-        "data/keypoint",
-        "data/n_contacts",
-        "data/state",
-        "meta/episode_ends",
-    }
-    for dataset in required_datasets:
-        assert dataset in zarr_data
-    nb_frames = zarr_data["data/img"].shape[0]
-
-    required_datasets.remove("meta/episode_ends")
-
-    assert all(nb_frames == zarr_data[dataset].shape[0] for dataset in required_datasets)
-
-
-def load_from_raw(
-    raw_dir: Path,
-    videos_dir: Path,
-    fps: int,
-    video: bool,
-    episodes: list[int] | None = None,
-    keypoints_instead_of_image: bool = False,
-    encoding: dict | None = None,
-):
-    try:
-        import pymunk
-        from gym_pusht.envs.pusht import PushTEnv, pymunk_to_shapely
-
-        from lerobot.common.datasets.push_dataset_to_hub._diffusion_policy_replay_buffer import (
-            ReplayBuffer as DiffusionPolicyReplayBuffer,
-        )
-    except ModuleNotFoundError as e:
-        print("`gym_pusht` is not installed. Please install it with `pip install 'lerobot[gym_pusht]'`")
-        raise e
-    # as define in gmy-pusht env: https://github.com/huggingface/gym-pusht/blob/e0684ff988d223808c0a9dcfaba9dc4991791370/gym_pusht/envs/pusht.py#L174
-    success_threshold = 0.95  # 95% coverage,
-
-    zarr_path = raw_dir / "pusht_cchi_v7_replay.zarr"
-    zarr_data = DiffusionPolicyReplayBuffer.copy_from_path(zarr_path)
-
-    episode_ids = torch.from_numpy(zarr_data.get_episode_idxs())
-    assert len(
-        {zarr_data[key].shape[0] for key in zarr_data.keys()}  # noqa: SIM118
-    ), "Some data type dont have the same number of total frames."
-
-    # TODO(rcadene): verify that goal pose is expected to be fixed
-    goal_pos_angle = np.array([256, 256, np.pi / 4])  # x, y, theta (in radians)
-    goal_body = PushTEnv.get_goal_pose_body(goal_pos_angle)
-
-    imgs = torch.from_numpy(zarr_data["img"])  # b h w c
-    states = torch.from_numpy(zarr_data["state"])
-    actions = torch.from_numpy(zarr_data["action"])
-
-    # load data indices from which each episode starts and ends
-    from_ids, to_ids = [], []
-    from_idx = 0
-    for to_idx in zarr_data.meta["episode_ends"]:
-        from_ids.append(from_idx)
-        to_ids.append(to_idx)
-        from_idx = to_idx
-
-    num_episodes = len(from_ids)
-
-    ep_dicts = []
-    ep_ids = episodes if episodes else range(num_episodes)
-    for ep_idx, selected_ep_idx in tqdm.tqdm(enumerate(ep_ids)):
-        from_idx = from_ids[selected_ep_idx]
-        to_idx = to_ids[selected_ep_idx]
-        num_frames = to_idx - from_idx
-
-        # sanity check
-        assert (episode_ids[from_idx:to_idx] == ep_idx).all()
-
-        # get image
-        if not keypoints_instead_of_image:
-            image = imgs[from_idx:to_idx]
-            assert image.min() >= 0.0
-            assert image.max() <= 255.0
-            image = image.type(torch.uint8)
-
-        # get state
-        state = states[from_idx:to_idx]
-        agent_pos = state[:, :2]
-        block_pos = state[:, 2:4]
-        block_angle = state[:, 4]
-
-        # get reward, success, done, and (maybe) keypoints
-        reward = torch.zeros(num_frames)
-        success = torch.zeros(num_frames, dtype=torch.bool)
-        if keypoints_instead_of_image:
-            keypoints = torch.zeros(num_frames, 16)  # 8 keypoints each with 2 coords
-        done = torch.zeros(num_frames, dtype=torch.bool)
-        for i in range(num_frames):
-            space = pymunk.Space()
-            space.gravity = 0, 0
-            space.damping = 0
-
-            # Add walls.
-            walls = [
-                PushTEnv.add_segment(space, (5, 506), (5, 5), 2),
-                PushTEnv.add_segment(space, (5, 5), (506, 5), 2),
-                PushTEnv.add_segment(space, (506, 5), (506, 506), 2),
-                PushTEnv.add_segment(space, (5, 506), (506, 506), 2),
-            ]
-            space.add(*walls)
-
-            block_body, block_shapes = PushTEnv.add_tee(space, block_pos[i].tolist(), block_angle[i].item())
-            goal_geom = pymunk_to_shapely(goal_body, block_body.shapes)
-            block_geom = pymunk_to_shapely(block_body, block_body.shapes)
-            intersection_area = goal_geom.intersection(block_geom).area
-            goal_area = goal_geom.area
-            coverage = intersection_area / goal_area
-            reward[i] = np.clip(coverage / success_threshold, 0, 1)
-            success[i] = coverage > success_threshold
-            if keypoints_instead_of_image:
-                keypoints[i] = torch.from_numpy(PushTEnv.get_keypoints(block_shapes).flatten())
-
-        # last step of demonstration is considered done
-        done[-1] = True
-
-        ep_dict = {}
-
-        if not keypoints_instead_of_image:
-            imgs_array = [x.numpy() for x in image]
-            img_key = "observation.image"
-            if video:
-                # save png images in temporary directory
-                tmp_imgs_dir = videos_dir / "tmp_images"
-                save_images_concurrently(imgs_array, tmp_imgs_dir)
-
-                # encode images to a mp4 video
-                fname = f"{img_key}_episode_{ep_idx:06d}.mp4"
-                video_path = videos_dir / fname
-                encode_video_frames(tmp_imgs_dir, video_path, fps, **(encoding or {}))
-
-                # clean temporary images directory
-                shutil.rmtree(tmp_imgs_dir)
-
-                # store the reference to the video frame
-                ep_dict[img_key] = [
-                    {"path": f"videos/{fname}", "timestamp": i / fps} for i in range(num_frames)
-                ]
-            else:
-                ep_dict[img_key] = [PILImage.fromarray(x) for x in imgs_array]
-
-        ep_dict["observation.state"] = agent_pos
-        if keypoints_instead_of_image:
-            ep_dict["observation.environment_state"] = keypoints
-        ep_dict["action"] = actions[from_idx:to_idx]
-        ep_dict["episode_index"] = torch.tensor([ep_idx] * num_frames, dtype=torch.int64)
-        ep_dict["frame_index"] = torch.arange(0, num_frames, 1)
-        ep_dict["timestamp"] = torch.arange(0, num_frames, 1) / fps
-        # ep_dict["next.observation.image"] = image[1:],
-        # ep_dict["next.observation.state"] = agent_pos[1:],
-        # TODO(rcadene)] = verify that reward and done are aligned with image and agent_pos
-        ep_dict["next.reward"] = torch.cat([reward[1:], reward[[-1]]])
-        ep_dict["next.done"] = torch.cat([done[1:], done[[-1]]])
-        ep_dict["next.success"] = torch.cat([success[1:], success[[-1]]])
-        ep_dicts.append(ep_dict)
-    data_dict = concatenate_episodes(ep_dicts)
-
-    total_frames = data_dict["frame_index"].shape[0]
-    data_dict["index"] = torch.arange(0, total_frames, 1)
-    return data_dict
-
-
-def to_hf_dataset(data_dict, video, keypoints_instead_of_image: bool = False):
-    features = {}
-
-    if not keypoints_instead_of_image:
-        if video:
-            features["observation.image"] = VideoFrame()
-        else:
-            features["observation.image"] = Image()
-
-    features["observation.state"] = Sequence(
-        length=data_dict["observation.state"].shape[1], feature=Value(dtype="float32", id=None)
-    )
-    if keypoints_instead_of_image:
-        features["observation.environment_state"] = Sequence(
-            length=data_dict["observation.environment_state"].shape[1],
-            feature=Value(dtype="float32", id=None),
-        )
-    features["action"] = Sequence(
-        length=data_dict["action"].shape[1], feature=Value(dtype="float32", id=None)
-    )
-    features["episode_index"] = Value(dtype="int64", id=None)
-    features["frame_index"] = Value(dtype="int64", id=None)
-    features["timestamp"] = Value(dtype="float32", id=None)
-    features["next.reward"] = Value(dtype="float32", id=None)
-    features["next.done"] = Value(dtype="bool", id=None)
-    features["next.success"] = Value(dtype="bool", id=None)
-    features["index"] = Value(dtype="int64", id=None)
-
-    hf_dataset = Dataset.from_dict(data_dict, features=Features(features))
-    hf_dataset.set_transform(hf_transform_to_torch)
-    return hf_dataset
-
-
-def from_raw_to_lerobot_format(
-    raw_dir: Path,
-    videos_dir: Path,
-    fps: int | None = None,
-    video: bool = True,
-    episodes: list[int] | None = None,
-    encoding: dict | None = None,
-):
-    # Manually change this to True to use keypoints of the T instead of an image observation (but don't merge
-    # with True). Also make sure to use video = 0 in the `push_dataset_to_hub.py` script.
-    keypoints_instead_of_image = False
-
-    # sanity check
-    check_format(raw_dir)
-
-    if fps is None:
-        fps = 10
-
-    data_dict = load_from_raw(raw_dir, videos_dir, fps, video, episodes, keypoints_instead_of_image, encoding)
-    hf_dataset = to_hf_dataset(data_dict, video, keypoints_instead_of_image)
-    episode_data_index = calculate_episode_data_index(hf_dataset)
-    info = {
-        "codebase_version": CODEBASE_VERSION,
-        "fps": fps,
-        "video": video if not keypoints_instead_of_image else 0,
-    }
-    if video:
-        info["encoding"] = get_default_encoding()
-
-    return hf_dataset, episode_data_index, info

lerobot/common/datasets/push_dataset_to_hub/umi_zarr_format.py

@@ -1,234 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Process UMI (Universal Manipulation Interface) data stored in Zarr format like in: https://github.com/real-stanford/universal_manipulation_interface"""
-
-import logging
-import shutil
-from pathlib import Path
-
-import torch
-import tqdm
-import zarr
-from datasets import Dataset, Features, Image, Sequence, Value
-from PIL import Image as PILImage
-
-from lerobot.common.datasets.lerobot_dataset import CODEBASE_VERSION
-from lerobot.common.datasets.push_dataset_to_hub._umi_imagecodecs_numcodecs import register_codecs
-from lerobot.common.datasets.push_dataset_to_hub.utils import (
-    calculate_episode_data_index,
-    concatenate_episodes,
-    get_default_encoding,
-    save_images_concurrently,
-)
-from lerobot.common.datasets.utils import (
-    hf_transform_to_torch,
-)
-from lerobot.common.datasets.video_utils import VideoFrame, encode_video_frames
-
-
-def check_format(raw_dir) -> bool:
-    zarr_path = raw_dir / "cup_in_the_wild.zarr"
-    zarr_data = zarr.open(zarr_path, mode="r")
-
-    required_datasets = {
-        "data/robot0_demo_end_pose",
-        "data/robot0_demo_start_pose",
-        "data/robot0_eef_pos",
-        "data/robot0_eef_rot_axis_angle",
-        "data/robot0_gripper_width",
-        "meta/episode_ends",
-        "data/camera0_rgb",
-    }
-    for dataset in required_datasets:
-        if dataset not in zarr_data:
-            return False
-
-    # mandatory to access zarr_data
-    register_codecs()
-    nb_frames = zarr_data["data/camera0_rgb"].shape[0]
-
-    required_datasets.remove("meta/episode_ends")
-    assert all(nb_frames == zarr_data[dataset].shape[0] for dataset in required_datasets)
-
-
-def load_from_raw(
-    raw_dir: Path,
-    videos_dir: Path,
-    fps: int,
-    video: bool,
-    episodes: list[int] | None = None,
-    encoding: dict | None = None,
-):
-    zarr_path = raw_dir / "cup_in_the_wild.zarr"
-    zarr_data = zarr.open(zarr_path, mode="r")
-
-    # We process the image data separately because it is too large to fit in memory
-    end_pose = torch.from_numpy(zarr_data["data/robot0_demo_end_pose"][:])
-    start_pos = torch.from_numpy(zarr_data["data/robot0_demo_start_pose"][:])
-    eff_pos = torch.from_numpy(zarr_data["data/robot0_eef_pos"][:])
-    eff_rot_axis_angle = torch.from_numpy(zarr_data["data/robot0_eef_rot_axis_angle"][:])
-    gripper_width = torch.from_numpy(zarr_data["data/robot0_gripper_width"][:])
-
-    states_pos = torch.cat([eff_pos, eff_rot_axis_angle], dim=1)
-    states = torch.cat([states_pos, gripper_width], dim=1)
-
-    episode_ends = zarr_data["meta/episode_ends"][:]
-    num_episodes = episode_ends.shape[0]
-
-    # We convert it in torch tensor later because the jit function does not support torch tensors
-    episode_ends = torch.from_numpy(episode_ends)
-
-    # load data indices from which each episode starts and ends
-    from_ids, to_ids = [], []
-    from_idx = 0
-    for to_idx in episode_ends:
-        from_ids.append(from_idx)
-        to_ids.append(to_idx)
-        from_idx = to_idx
-
-    ep_dicts_dir = videos_dir / "ep_dicts"
-    ep_dicts_dir.mkdir(exist_ok=True, parents=True)
-    ep_dicts = []
-
-    ep_ids = episodes if episodes else range(num_episodes)
-    for ep_idx, selected_ep_idx in tqdm.tqdm(enumerate(ep_ids)):
-        ep_dict_path = ep_dicts_dir / f"{ep_idx}"
-        if not ep_dict_path.is_file():
-            from_idx = from_ids[selected_ep_idx]
-            to_idx = to_ids[selected_ep_idx]
-            num_frames = to_idx - from_idx
-
-            # TODO(rcadene): save temporary images of the episode?
-
-            state = states[from_idx:to_idx]
-
-            ep_dict = {}
-
-            # load 57MB of images in RAM (400x224x224x3 uint8)
-            imgs_array = zarr_data["data/camera0_rgb"][from_idx:to_idx]
-            img_key = "observation.image"
-            if video:
-                fname = f"{img_key}_episode_{ep_idx:06d}.mp4"
-                video_path = videos_dir / fname
-                if not video_path.is_file():
-                    # save png images in temporary directory
-                    tmp_imgs_dir = videos_dir / "tmp_images"
-                    save_images_concurrently(imgs_array, tmp_imgs_dir)
-
-                    # encode images to a mp4 video
-                    encode_video_frames(tmp_imgs_dir, video_path, fps, **(encoding or {}))
-
-                    # clean temporary images directory
-                    shutil.rmtree(tmp_imgs_dir)
-
-                # store the reference to the video frame
-                ep_dict[img_key] = [
-                    {"path": f"videos/{fname}", "timestamp": i / fps} for i in range(num_frames)
-                ]
-            else:
-                ep_dict[img_key] = [PILImage.fromarray(x) for x in imgs_array]
-
-            ep_dict["observation.state"] = state
-            ep_dict["episode_index"] = torch.tensor([ep_idx] * num_frames, dtype=torch.int64)
-            ep_dict["frame_index"] = torch.arange(0, num_frames, 1)
-            ep_dict["timestamp"] = torch.arange(0, num_frames, 1) / fps
-            ep_dict["episode_data_index_from"] = torch.tensor([from_idx] * num_frames)
-            ep_dict["episode_data_index_to"] = torch.tensor([from_idx + num_frames] * num_frames)
-            ep_dict["end_pose"] = end_pose[from_idx:to_idx]
-            ep_dict["start_pos"] = start_pos[from_idx:to_idx]
-            ep_dict["gripper_width"] = gripper_width[from_idx:to_idx]
-            torch.save(ep_dict, ep_dict_path)
-        else:
-            ep_dict = torch.load(ep_dict_path)
-
-        ep_dicts.append(ep_dict)
-
-    data_dict = concatenate_episodes(ep_dicts)
-
-    total_frames = data_dict["frame_index"].shape[0]
-    data_dict["index"] = torch.arange(0, total_frames, 1)
-    return data_dict
-
-
-def to_hf_dataset(data_dict, video):
-    features = {}
-
-    if video:
-        features["observation.image"] = VideoFrame()
-    else:
-        features["observation.image"] = Image()
-
-    features["observation.state"] = Sequence(
-        length=data_dict["observation.state"].shape[1], feature=Value(dtype="float32", id=None)
-    )
-    features["episode_index"] = Value(dtype="int64", id=None)
-    features["frame_index"] = Value(dtype="int64", id=None)
-    features["timestamp"] = Value(dtype="float32", id=None)
-    features["index"] = Value(dtype="int64", id=None)
-    features["episode_data_index_from"] = Value(dtype="int64", id=None)
-    features["episode_data_index_to"] = Value(dtype="int64", id=None)
-    # `start_pos` and `end_pos` respectively represent the positions of the end-effector
-    # at the beginning and the end of the episode.
-    # `gripper_width` indicates the distance between the grippers, and this value is included
-    # in the state vector, which comprises the concatenation of the end-effector position
-    # and gripper width.
-    features["end_pose"] = Sequence(
-        length=data_dict["end_pose"].shape[1], feature=Value(dtype="float32", id=None)
-    )
-    features["start_pos"] = Sequence(
-        length=data_dict["start_pos"].shape[1], feature=Value(dtype="float32", id=None)
-    )
-    features["gripper_width"] = Sequence(
-        length=data_dict["gripper_width"].shape[1], feature=Value(dtype="float32", id=None)
-    )
-
-    hf_dataset = Dataset.from_dict(data_dict, features=Features(features))
-    hf_dataset.set_transform(hf_transform_to_torch)
-    return hf_dataset
-
-
-def from_raw_to_lerobot_format(
-    raw_dir: Path,
-    videos_dir: Path,
-    fps: int | None = None,
-    video: bool = True,
-    episodes: list[int] | None = None,
-    encoding: dict | None = None,
-):
-    # sanity check
-    check_format(raw_dir)
-
-    if fps is None:
-        # For umi cup in the wild: https://arxiv.org/pdf/2402.10329#table.caption.16
-        fps = 10
-
-    if not video:
-        logging.warning(
-            "Generating UMI dataset without `video=True` creates ~150GB on disk and requires ~80GB in RAM."
-        )
-
-    data_dict = load_from_raw(raw_dir, videos_dir, fps, video, episodes, encoding)
-    hf_dataset = to_hf_dataset(data_dict, video)
-    episode_data_index = calculate_episode_data_index(hf_dataset)
-    info = {
-        "codebase_version": CODEBASE_VERSION,
-        "fps": fps,
-        "video": video,
-    }
-    if video:
-        info["encoding"] = get_default_encoding()
-
-    return hf_dataset, episode_data_index, info

lerobot/common/datasets/push_dataset_to_hub/xarm_pkl_format.py

@@ -1,200 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Process pickle files formatted like in: https://github.com/fyhMer/fowm"""
-
-import pickle
-import shutil
-from pathlib import Path
-
-import einops
-import torch
-import tqdm
-from datasets import Dataset, Features, Image, Sequence, Value
-from PIL import Image as PILImage
-
-from lerobot.common.datasets.lerobot_dataset import CODEBASE_VERSION
-from lerobot.common.datasets.push_dataset_to_hub.utils import (
-    calculate_episode_data_index,
-    concatenate_episodes,
-    get_default_encoding,
-    save_images_concurrently,
-)
-from lerobot.common.datasets.utils import (
-    hf_transform_to_torch,
-)
-from lerobot.common.datasets.video_utils import VideoFrame, encode_video_frames
-
-
-def check_format(raw_dir):
-    keys = {"actions", "rewards", "dones"}
-    nested_keys = {"observations": {"rgb", "state"}, "next_observations": {"rgb", "state"}}
-
-    xarm_files = list(raw_dir.glob("*.pkl"))
-    assert len(xarm_files) > 0
-
-    with open(xarm_files[0], "rb") as f:
-        dataset_dict = pickle.load(f)
-
-    assert isinstance(dataset_dict, dict)
-    assert all(k in dataset_dict for k in keys)
-
-    # Check for consistent lengths in nested keys
-    expected_len = len(dataset_dict["actions"])
-    assert all(len(dataset_dict[key]) == expected_len for key in keys if key in dataset_dict)
-
-    for key, subkeys in nested_keys.items():
-        nested_dict = dataset_dict.get(key, {})
-        assert all(len(nested_dict[subkey]) == expected_len for subkey in subkeys if subkey in nested_dict)
-
-
-def load_from_raw(
-    raw_dir: Path,
-    videos_dir: Path,
-    fps: int,
-    video: bool,
-    episodes: list[int] | None = None,
-    encoding: dict | None = None,
-):
-    pkl_path = raw_dir / "buffer.pkl"
-
-    with open(pkl_path, "rb") as f:
-        pkl_data = pickle.load(f)
-
-    # load data indices from which each episode starts and ends
-    from_ids, to_ids = [], []
-    from_idx, to_idx = 0, 0
-    for done in pkl_data["dones"]:
-        to_idx += 1
-        if not done:
-            continue
-        from_ids.append(from_idx)
-        to_ids.append(to_idx)
-        from_idx = to_idx
-
-    num_episodes = len(from_ids)
-
-    ep_dicts = []
-    ep_ids = episodes if episodes else range(num_episodes)
-    for ep_idx, selected_ep_idx in tqdm.tqdm(enumerate(ep_ids)):
-        from_idx = from_ids[selected_ep_idx]
-        to_idx = to_ids[selected_ep_idx]
-        num_frames = to_idx - from_idx
-
-        image = torch.tensor(pkl_data["observations"]["rgb"][from_idx:to_idx])
-        image = einops.rearrange(image, "b c h w -> b h w c")
-        state = torch.tensor(pkl_data["observations"]["state"][from_idx:to_idx])
-        action = torch.tensor(pkl_data["actions"][from_idx:to_idx])
-        # TODO(rcadene): we have a missing last frame which is the observation when the env is done
-        # it is critical to have this frame for tdmpc to predict a "done observation/state"
-        # next_image = torch.tensor(pkl_data["next_observations"]["rgb"][from_idx:to_idx])
-        # next_state = torch.tensor(pkl_data["next_observations"]["state"][from_idx:to_idx])
-        next_reward = torch.tensor(pkl_data["rewards"][from_idx:to_idx])
-        next_done = torch.tensor(pkl_data["dones"][from_idx:to_idx])
-
-        ep_dict = {}
-
-        imgs_array = [x.numpy() for x in image]
-        img_key = "observation.image"
-        if video:
-            # save png images in temporary directory
-            tmp_imgs_dir = videos_dir / "tmp_images"
-            save_images_concurrently(imgs_array, tmp_imgs_dir)
-
-            # encode images to a mp4 video
-            fname = f"{img_key}_episode_{ep_idx:06d}.mp4"
-            video_path = videos_dir / fname
-            encode_video_frames(tmp_imgs_dir, video_path, fps, **(encoding or {}))
-
-            # clean temporary images directory
-            shutil.rmtree(tmp_imgs_dir)
-
-            # store the reference to the video frame
-            ep_dict[img_key] = [{"path": f"videos/{fname}", "timestamp": i / fps} for i in range(num_frames)]
-        else:
-            ep_dict[img_key] = [PILImage.fromarray(x) for x in imgs_array]
-
-        ep_dict["observation.state"] = state
-        ep_dict["action"] = action
-        ep_dict["episode_index"] = torch.tensor([ep_idx] * num_frames, dtype=torch.int64)
-        ep_dict["frame_index"] = torch.arange(0, num_frames, 1)
-        ep_dict["timestamp"] = torch.arange(0, num_frames, 1) / fps
-        # ep_dict["next.observation.image"] = next_image
-        # ep_dict["next.observation.state"] = next_state
-        ep_dict["next.reward"] = next_reward
-        ep_dict["next.done"] = next_done
-        ep_dicts.append(ep_dict)
-
-    data_dict = concatenate_episodes(ep_dicts)
-
-    total_frames = data_dict["frame_index"].shape[0]
-    data_dict["index"] = torch.arange(0, total_frames, 1)
-    return data_dict
-
-
-def to_hf_dataset(data_dict, video):
-    features = {}
-
-    if video:
-        features["observation.image"] = VideoFrame()
-    else:
-        features["observation.image"] = Image()
-
-    features["observation.state"] = Sequence(
-        length=data_dict["observation.state"].shape[1], feature=Value(dtype="float32", id=None)
-    )
-    features["action"] = Sequence(
-        length=data_dict["action"].shape[1], feature=Value(dtype="float32", id=None)
-    )
-    features["episode_index"] = Value(dtype="int64", id=None)
-    features["frame_index"] = Value(dtype="int64", id=None)
-    features["timestamp"] = Value(dtype="float32", id=None)
-    features["next.reward"] = Value(dtype="float32", id=None)
-    features["next.done"] = Value(dtype="bool", id=None)
-    features["index"] = Value(dtype="int64", id=None)
-    # TODO(rcadene): add success
-    # features["next.success"] = Value(dtype='bool', id=None)
-
-    hf_dataset = Dataset.from_dict(data_dict, features=Features(features))
-    hf_dataset.set_transform(hf_transform_to_torch)
-    return hf_dataset
-
-
-def from_raw_to_lerobot_format(
-    raw_dir: Path,
-    videos_dir: Path,
-    fps: int | None = None,
-    video: bool = True,
-    episodes: list[int] | None = None,
-    encoding: dict | None = None,
-):
-    # sanity check
-    check_format(raw_dir)
-
-    if fps is None:
-        fps = 15
-
-    data_dict = load_from_raw(raw_dir, videos_dir, fps, video, episodes, encoding)
-    hf_dataset = to_hf_dataset(data_dict, video)
-    episode_data_index = calculate_episode_data_index(hf_dataset)
-    info = {
-        "codebase_version": CODEBASE_VERSION,
-        "fps": fps,
-        "video": video,
-    }
-    if video:
-        info["encoding"] = get_default_encoding()
-
-    return hf_dataset, episode_data_index, info

lerobot/common/datasets/utils.py

@@ -31,6 +31,7 @@ import packaging.version
 import torch
 from datasets.table import embed_table_storage
 from huggingface_hub import DatasetCard, DatasetCardData, HfApi
+from huggingface_hub.errors import RevisionNotFoundError
 from PIL import Image as PILImage
 from torchvision import transforms
 
@@ -39,7 +40,7 @@ from lerobot.common.datasets.backward_compatibility import (
     BackwardCompatibilityError,
     ForwardCompatibilityError,
 )
-from lerobot.common.robot_devices.robots.utils import Robot
+from lerobot.common.robots.utils import Robot
 from lerobot.common.utils.utils import is_valid_numpy_dtype_string
 from lerobot.configs.types import DictLike, FeatureType, PolicyFeature
 
@@ -239,7 +240,7 @@ def load_episodes_stats(local_dir: Path) -> dict:
 def backward_compatible_episodes_stats(
     stats: dict[str, dict[str, np.ndarray]], episodes: list[int]
 ) -> dict[str, dict[str, np.ndarray]]:
-    return {ep_idx: stats for ep_idx in episodes}
+    return dict.fromkeys(episodes, stats)
 
 
 def load_image_as_numpy(
@@ -325,6 +326,19 @@ def get_safe_version(repo_id: str, version: str | packaging.version.Version) ->
     )
     hub_versions = get_repo_versions(repo_id)
 
+    if not hub_versions:
+        raise RevisionNotFoundError(
+            f"""Your dataset must be tagged with a codebase version.
+            Assuming _version_ is the codebase_version value in the info.json, you can run this:
+            ```python
+            from huggingface_hub import HfApi
+
+            hub_api = HfApi()
+            hub_api.create_tag("{repo_id}", tag="_version_", repo_type="dataset")
+            ```
+            """
+        )
+
     if target_version in hub_versions:
         return f"v{target_version}"
 

lerobot/common/datasets/v2/batch_convert_dataset_v1_to_v2.py

@@ -27,7 +27,7 @@ from textwrap import dedent
 
 from lerobot import available_datasets
 from lerobot.common.datasets.v2.convert_dataset_v1_to_v2 import convert_dataset
-from lerobot.common.robot_devices.robots.configs import AlohaRobotConfig
+from lerobot.common.robots.aloha.configuration_aloha import AlohaRobotConfig
 
 LOCAL_DIR = Path("data/")
 

lerobot/common/datasets/v2/convert_dataset_v1_to_v2.py

@@ -141,8 +141,8 @@ from lerobot.common.datasets.video_utils import (
     get_image_pixel_channels,
     get_video_info,
 )
-from lerobot.common.robot_devices.robots.configs import RobotConfig
-from lerobot.common.robot_devices.robots.utils import make_robot_config
+from lerobot.common.robots import RobotConfig
+from lerobot.common.robots.utils import make_robot_config
 
 V16 = "v1.6"
 V20 = "v2.0"
@@ -481,7 +481,7 @@ def convert_dataset(
 
     # Tasks
     if single_task:
-        tasks_by_episodes = {ep_idx: single_task for ep_idx in episode_indices}
+        tasks_by_episodes = dict.fromkeys(episode_indices, single_task)
         dataset, tasks = add_task_index_by_episodes(dataset, tasks_by_episodes)
         tasks_by_episodes = {ep_idx: [task] for ep_idx, task in tasks_by_episodes.items()}
     elif tasks_path:

lerobot/common/datasets/v21/_remove_language_instruction.py

@@ -1,3 +1,17 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import logging
 import traceback
 from pathlib import Path

lerobot/common/datasets/v21/convert_dataset_v20_to_v21.py

@@ -1,3 +1,17 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 """
 This script will help you convert any LeRobot dataset already pushed to the hub from codebase version 2.0 to
 2.1. It will:
@@ -57,7 +71,7 @@ def convert_dataset(
     dataset.meta.info["codebase_version"] = CODEBASE_VERSION
     write_info(dataset.meta.info, dataset.root)
 
-    dataset.push_to_hub(branch=branch, allow_patterns="meta/")
+    dataset.push_to_hub(branch=branch, tag_version=False, allow_patterns="meta/")
 
     # delete old stats.json file
     if (dataset.root / STATS_PATH).is_file:

lerobot/common/datasets/v21/convert_stats.py

@@ -1,3 +1,17 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from concurrent.futures import ThreadPoolExecutor, as_completed
 
 import numpy as np

lerobot/common/datasets/video_utils.py

@@ -13,6 +13,7 @@
 # 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 importlib
 import json
 import logging
 import subprocess
@@ -29,6 +30,46 @@ from datasets.features.features import register_feature
 from PIL import Image
 
 
+def get_safe_default_codec():
+    if importlib.util.find_spec("torchcodec"):
+        return "torchcodec"
+    else:
+        logging.warning(
+            "'torchcodec' is not available in your platform, falling back to 'pyav' as a default decoder"
+        )
+        return "pyav"
+
+
+def decode_video_frames(
+    video_path: Path | str,
+    timestamps: list[float],
+    tolerance_s: float,
+    backend: str | None = None,
+) -> torch.Tensor:
+    """
+    Decodes video frames using the specified backend.
+
+    Args:
+        video_path (Path): Path to the video file.
+        timestamps (list[float]): List of timestamps to extract frames.
+        tolerance_s (float): Allowed deviation in seconds for frame retrieval.
+        backend (str, optional): Backend to use for decoding. Defaults to "torchcodec" when available in the platform; otherwise, defaults to "pyav"..
+
+    Returns:
+        torch.Tensor: Decoded frames.
+
+    Currently supports torchcodec on cpu and pyav.
+    """
+    if backend is None:
+        backend = get_safe_default_codec()
+    if backend == "torchcodec":
+        return decode_video_frames_torchcodec(video_path, timestamps, tolerance_s)
+    elif backend in ["pyav", "video_reader"]:
+        return decode_video_frames_torchvision(video_path, timestamps, tolerance_s, backend)
+    else:
+        raise ValueError(f"Unsupported video backend: {backend}")
+
+
 def decode_video_frames_torchvision(
     video_path: Path | str,
     timestamps: list[float],
@@ -127,6 +168,81 @@ def decode_video_frames_torchvision(
     return closest_frames
 
 
+def decode_video_frames_torchcodec(
+    video_path: Path | str,
+    timestamps: list[float],
+    tolerance_s: float,
+    device: str = "cpu",
+    log_loaded_timestamps: bool = False,
+) -> torch.Tensor:
+    """Loads frames associated with the requested timestamps of a video using torchcodec.
+
+    Note: Setting device="cuda" outside the main process, e.g. in data loader workers, will lead to CUDA initialization errors.
+
+    Note: Video benefits from inter-frame compression. Instead of storing every frame individually,
+    the encoder stores a reference frame (or a key frame) and subsequent frames as differences relative to
+    that key frame. As a consequence, to access a requested frame, we need to load the preceding key frame,
+    and all subsequent frames until reaching the requested frame. The number of key frames in a video
+    can be adjusted during encoding to take into account decoding time and video size in bytes.
+    """
+
+    if importlib.util.find_spec("torchcodec"):
+        from torchcodec.decoders import VideoDecoder
+    else:
+        raise ImportError("torchcodec is required but not available.")
+
+    # initialize video decoder
+    decoder = VideoDecoder(video_path, device=device, seek_mode="approximate")
+    loaded_frames = []
+    loaded_ts = []
+    # get metadata for frame information
+    metadata = decoder.metadata
+    average_fps = metadata.average_fps
+
+    # convert timestamps to frame indices
+    frame_indices = [round(ts * average_fps) for ts in timestamps]
+
+    # retrieve frames based on indices
+    frames_batch = decoder.get_frames_at(indices=frame_indices)
+
+    for frame, pts in zip(frames_batch.data, frames_batch.pts_seconds, strict=False):
+        loaded_frames.append(frame)
+        loaded_ts.append(pts.item())
+        if log_loaded_timestamps:
+            logging.info(f"Frame loaded at timestamp={pts:.4f}")
+
+    query_ts = torch.tensor(timestamps)
+    loaded_ts = torch.tensor(loaded_ts)
+
+    # compute distances between each query timestamp and loaded timestamps
+    dist = torch.cdist(query_ts[:, None], loaded_ts[:, None], p=1)
+    min_, argmin_ = dist.min(1)
+
+    is_within_tol = min_ < tolerance_s
+    assert is_within_tol.all(), (
+        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."
+        "This might be due to synchronization issues with timestamps during data collection."
+        "To be safe, we advise to ignore this item during training."
+        f"\nqueried timestamps: {query_ts}"
+        f"\nloaded timestamps: {loaded_ts}"
+        f"\nvideo: {video_path}"
+    )
+
+    # get closest frames to the query timestamps
+    closest_frames = torch.stack([loaded_frames[idx] for idx in argmin_])
+    closest_ts = loaded_ts[argmin_]
+
+    if log_loaded_timestamps:
+        logging.info(f"{closest_ts=}")
+
+    # convert to float32 in [0,1] range (channel first)
+    closest_frames = closest_frames.type(torch.float32) / 255
+
+    assert len(timestamps) == len(closest_frames)
+    return closest_frames
+
+
 def encode_video_frames(
     imgs_dir: Path | str,
     video_path: Path | str,
@@ -141,6 +257,7 @@ def encode_video_frames(
 ) -> None:
     """More info on ffmpeg arguments tuning on `benchmark/video/README.md`"""
     video_path = Path(video_path)
+    imgs_dir = Path(imgs_dir)
     video_path.parent.mkdir(parents=True, exist_ok=True)
 
     ffmpeg_args = OrderedDict(

lerobot/common/envs/__init__.py

@@ -1 +1,15 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from .configs import AlohaEnv, EnvConfig, PushtEnv, XarmEnv  # noqa: F401

lerobot/common/envs/configs.py

@@ -1,9 +1,23 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import abc
 from dataclasses import dataclass, field
 
 import draccus
 
-from lerobot.common.constants import ACTION, OBS_ENV, OBS_IMAGE, OBS_IMAGES, OBS_ROBOT
+from lerobot.common.constants import ACTION, OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
 from lerobot.configs.types import FeatureType, PolicyFeature
 
 
@@ -39,7 +53,7 @@ class AlohaEnv(EnvConfig):
     features_map: dict[str, str] = field(
         default_factory=lambda: {
             "action": ACTION,
-            "agent_pos": OBS_ROBOT,
+            "agent_pos": OBS_STATE,
             "top": f"{OBS_IMAGE}.top",
             "pixels/top": f"{OBS_IMAGES}.top",
         }
@@ -80,8 +94,8 @@ class PushtEnv(EnvConfig):
     features_map: dict[str, str] = field(
         default_factory=lambda: {
             "action": ACTION,
-            "agent_pos": OBS_ROBOT,
-            "environment_state": OBS_ENV,
+            "agent_pos": OBS_STATE,
+            "environment_state": OBS_ENV_STATE,
             "pixels": OBS_IMAGE,
         }
     )
@@ -122,7 +136,7 @@ class XarmEnv(EnvConfig):
     features_map: dict[str, str] = field(
         default_factory=lambda: {
             "action": ACTION,
-            "agent_pos": OBS_ROBOT,
+            "agent_pos": OBS_STATE,
             "pixels": OBS_IMAGE,
         }
     )

lerobot/common/envs/utils.py

@@ -13,7 +13,11 @@
 # 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 warnings
+from typing import Any
+
 import einops
+import gymnasium as gym
 import numpy as np
 import torch
 from torch import Tensor
@@ -86,3 +90,38 @@ def env_to_policy_features(env_cfg: EnvConfig) -> dict[str, PolicyFeature]:
         policy_features[policy_key] = feature
 
     return policy_features
+
+
+def are_all_envs_same_type(env: gym.vector.VectorEnv) -> bool:
+    first_type = type(env.envs[0])  # Get type of first env
+    return all(type(e) is first_type for e in env.envs)  # Fast type check
+
+
+def check_env_attributes_and_types(env: gym.vector.VectorEnv) -> None:
+    with warnings.catch_warnings():
+        warnings.simplefilter("once", UserWarning)  # Apply filter only in this function
+
+        if not (hasattr(env.envs[0], "task_description") and hasattr(env.envs[0], "task")):
+            warnings.warn(
+                "The environment does not have 'task_description' and 'task'. Some policies require these features.",
+                UserWarning,
+                stacklevel=2,
+            )
+        if not are_all_envs_same_type(env):
+            warnings.warn(
+                "The environments have different types. Make sure you infer the right task from each environment. Empty task will be passed instead.",
+                UserWarning,
+                stacklevel=2,
+            )
+
+
+def add_envs_task(env: gym.vector.VectorEnv, observation: dict[str, Any]) -> dict[str, Any]:
+    """Adds task feature to the observation dict with respect to the first environment attribute."""
+    if hasattr(env.envs[0], "task_description"):
+        observation["task"] = env.call("task_description")
+    elif hasattr(env.envs[0], "task"):
+        observation["task"] = env.call("task")
+    else:  #  For envs without language instructions, e.g. aloha transfer cube and etc.
+        num_envs = observation[list(observation.keys())[0]].shape[0]
+        observation["task"] = ["" for _ in range(num_envs)]
+    return observation

lerobot/common/errors.py

@@ -0,0 +1,17 @@
+class DeviceNotConnectedError(ConnectionError):
+    """Exception raised when the device is not connected."""
+
+    def __init__(self, message="This device is not connected. Try calling `connect()` first."):
+        self.message = message
+        super().__init__(self.message)
+
+
+class DeviceAlreadyConnectedError(ConnectionError):
+    """Exception raised when the device is already connected."""
+
+    def __init__(
+        self,
+        message="This device is already connected. Try not calling `connect()` twice.",
+    ):
+        self.message = message
+        super().__init__(self.message)

lerobot/common/motors/__init__.py

@@ -0,0 +1 @@
+from .motors_bus import Motor, MotorCalibration, MotorNormMode, MotorsBus

lerobot/common/motors/configs.py

@@ -0,0 +1,41 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import abc
+from dataclasses import dataclass
+
+import draccus
+
+
+@dataclass
+class MotorsBusConfig(draccus.ChoiceRegistry, abc.ABC):
+    @property
+    def type(self) -> str:
+        return self.get_choice_name(self.__class__)
+
+
+@MotorsBusConfig.register_subclass("dynamixel")
+@dataclass
+class DynamixelMotorsBusConfig(MotorsBusConfig):
+    port: str
+    motors: dict[str, tuple[int, str]]
+    mock: bool = False
+
+
+@MotorsBusConfig.register_subclass("feetech")
+@dataclass
+class FeetechMotorsBusConfig(MotorsBusConfig):
+    port: str
+    motors: dict[str, tuple[int, str]]
+    mock: bool = False

lerobot/common/motors/dynamixel/__init__.py

@@ -0,0 +1,3 @@
+from .dynamixel import DriveMode, DynamixelMotorsBus, OperatingMode, TorqueMode
+from .dynamixel_calibration import run_arm_calibration
+from .tables import *

lerobot/common/motors/dynamixel/dynamixel.py

@@ -0,0 +1,206 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# TODO(aliberts): Should we implement FastSyncRead/Write?
+# https://github.com/ROBOTIS-GIT/DynamixelSDK/pull/643
+# https://github.com/ROBOTIS-GIT/DynamixelSDK/releases/tag/3.8.2
+# https://emanual.robotis.com/docs/en/dxl/protocol2/#fast-sync-read-0x8a
+# -> Need to check compatibility across models
+
+import logging
+from copy import deepcopy
+from enum import Enum
+
+from lerobot.common.utils.encoding_utils import decode_twos_complement, encode_twos_complement
+
+from ..motors_bus import Motor, MotorCalibration, MotorsBus, NameOrID, Value
+from .tables import (
+    AVAILABLE_BAUDRATES,
+    MODEL_BAUDRATE_TABLE,
+    MODEL_CONTROL_TABLE,
+    MODEL_ENCODING_TABLE,
+    MODEL_NUMBER_TABLE,
+    MODEL_RESOLUTION,
+)
+
+PROTOCOL_VERSION = 2.0
+BAUDRATE = 1_000_000
+DEFAULT_TIMEOUT_MS = 1000
+
+NORMALIZED_DATA = ["Goal_Position", "Present_Position"]
+CONVERT_UINT32_TO_INT32_REQUIRED = ["Goal_Position", "Present_Position"]
+
+logger = logging.getLogger(__name__)
+
+
+class OperatingMode(Enum):
+    # DYNAMIXEL only controls current(torque) regardless of speed and position. This mode is ideal for a
+    # gripper or a system that only uses current(torque) control or a system that has additional
+    # velocity/position controllers.
+    CURRENT = 0
+
+    # This mode controls velocity. This mode is identical to the Wheel Mode(endless) from existing DYNAMIXEL.
+    # This mode is ideal for wheel-type robots.
+    VELOCITY = 1
+
+    # This mode controls position. This mode is identical to the Joint Mode from existing DYNAMIXEL. Operating
+    # position range is limited by the Max Position Limit(48) and the Min Position Limit(52). This mode is
+    # ideal for articulated robots that each joint rotates less than 360 degrees.
+    POSITION = 3
+
+    # This mode controls position. This mode is identical to the Multi-turn Position Control from existing
+    # DYNAMIXEL. 512 turns are supported(-256[rev] ~ 256[rev]). This mode is ideal for multi-turn wrists or
+    # conveyer systems or a system that requires an additional reduction gear. Note that Max Position
+    # Limit(48), Min Position Limit(52) are not used on Extended Position Control Mode.
+    EXTENDED_POSITION = 4
+
+    # This mode controls both position and current(torque). Up to 512 turns are supported (-256[rev] ~
+    # 256[rev]). This mode is ideal for a system that requires both position and current control such as
+    # articulated robots or grippers.
+    CURRENT_POSITION = 5
+
+    # This mode directly controls PWM output. (Voltage Control Mode)
+    PWM = 16
+
+
+class DriveMode(Enum):
+    NON_INVERTED = 0
+    INVERTED = 1
+
+
+class TorqueMode(Enum):
+    ENABLED = 1
+    DISABLED = 0
+
+
+def _split_into_byte_chunks(value: int, length: int) -> list[int]:
+    import dynamixel_sdk as dxl
+
+    if length == 1:
+        data = [value]
+    elif length == 2:
+        data = [dxl.DXL_LOBYTE(value), dxl.DXL_HIBYTE(value)]
+    elif length == 4:
+        data = [
+            dxl.DXL_LOBYTE(dxl.DXL_LOWORD(value)),
+            dxl.DXL_HIBYTE(dxl.DXL_LOWORD(value)),
+            dxl.DXL_LOBYTE(dxl.DXL_HIWORD(value)),
+            dxl.DXL_HIBYTE(dxl.DXL_HIWORD(value)),
+        ]
+    return data
+
+
+class DynamixelMotorsBus(MotorsBus):
+    """
+    The Dynamixel implementation for a MotorsBus. It relies on the python dynamixel sdk to communicate with
+    the motors. For more info, see the Dynamixel SDK Documentation:
+    https://emanual.robotis.com/docs/en/software/dynamixel/dynamixel_sdk/sample_code/python_read_write_protocol_2_0/#python-read-write-protocol-20
+    """
+
+    available_baudrates = deepcopy(AVAILABLE_BAUDRATES)
+    default_timeout = DEFAULT_TIMEOUT_MS
+    model_baudrate_table = deepcopy(MODEL_BAUDRATE_TABLE)
+    model_ctrl_table = deepcopy(MODEL_CONTROL_TABLE)
+    model_encoding_table = deepcopy(MODEL_ENCODING_TABLE)
+    model_number_table = deepcopy(MODEL_NUMBER_TABLE)
+    model_resolution_table = deepcopy(MODEL_RESOLUTION)
+    normalized_data = deepcopy(NORMALIZED_DATA)
+
+    def __init__(
+        self,
+        port: str,
+        motors: dict[str, Motor],
+        calibration: dict[str, MotorCalibration] | None = None,
+    ):
+        super().__init__(port, motors, calibration)
+        import dynamixel_sdk as dxl
+
+        self.port_handler = dxl.PortHandler(self.port)
+        self.packet_handler = dxl.PacketHandler(PROTOCOL_VERSION)
+        self.sync_reader = dxl.GroupSyncRead(self.port_handler, self.packet_handler, 0, 0)
+        self.sync_writer = dxl.GroupSyncWrite(self.port_handler, self.packet_handler, 0, 0)
+        self._comm_success = dxl.COMM_SUCCESS
+        self._no_error = 0x00
+
+    def _assert_protocol_is_compatible(self, instruction_name: str) -> None:
+        pass
+
+    def _handshake(self) -> None:
+        self._assert_motors_exist()
+
+    def configure_motors(self) -> None:
+        # By default, Dynamixel motors have a 500µs delay response time (corresponding to a value of 250 on
+        # the 'Return_Delay_Time' address). We ensure this is reduced to the minimum of 2µs (value of 0).
+        for motor in self.motors:
+            self.write("Return_Delay_Time", motor, 0)
+
+    def disable_torque(self, motors: str | list[str] | None = None, num_retry: int = 0) -> None:
+        for name in self._get_motors_list(motors):
+            self.write("Torque_Enable", name, TorqueMode.DISABLED.value, num_retry=num_retry)
+
+    def enable_torque(self, motors: str | list[str] | None = None, num_retry: int = 0) -> None:
+        for name in self._get_motors_list(motors):
+            self.write("Torque_Enable", name, TorqueMode.ENABLED.value, num_retry=num_retry)
+
+    def _encode_sign(self, data_name: str, ids_values: dict[int, int]) -> dict[int, int]:
+        for id_ in ids_values:
+            model = self._id_to_model(id_)
+            encoding_table = self.model_encoding_table.get(model)
+            if encoding_table and data_name in encoding_table:
+                n_bytes = encoding_table[data_name]
+                ids_values[id_] = encode_twos_complement(ids_values[id_], n_bytes)
+
+        return ids_values
+
+    def _decode_sign(self, data_name: str, ids_values: dict[int, int]) -> dict[int, int]:
+        for id_ in ids_values:
+            model = self._id_to_model(id_)
+            encoding_table = self.model_encoding_table.get(model)
+            if encoding_table and data_name in encoding_table:
+                n_bytes = encoding_table[data_name]
+                ids_values[id_] = decode_twos_complement(ids_values[id_], n_bytes)
+
+        return ids_values
+
+    def _get_half_turn_homings(self, positions: dict[NameOrID, Value]) -> dict[NameOrID, Value]:
+        """
+        On Dynamixel Motors:
+        Present_Position = Actual_Position + Homing_Offset
+        """
+        half_turn_homings = {}
+        for motor, pos in positions.items():
+            model = self._get_motor_model(motor)
+            max_res = self.model_resolution_table[model] - 1
+            half_turn_homings[motor] = int(max_res / 2) - pos
+
+        return half_turn_homings
+
+    def _split_into_byte_chunks(self, value: int, length: int) -> list[int]:
+        return _split_into_byte_chunks(value, length)
+
+    def broadcast_ping(self, num_retry: int = 0, raise_on_error: bool = False) -> dict[int, int] | None:
+        for n_try in range(1 + num_retry):
+            data_list, comm = self.packet_handler.broadcastPing(self.port_handler)
+            if self._is_comm_success(comm):
+                break
+            logger.debug(f"Broadcast ping failed on port '{self.port}' ({n_try=})")
+            logger.debug(self.packet_handler.getTxRxResult(comm))
+
+        if not self._is_comm_success(comm):
+            if raise_on_error:
+                raise ConnectionError(self.packet_handler.getTxRxResult(comm))
+
+            return
+
+        return {id_: data[0] for id_, data in data_list.items()}

lerobot/common/motors/dynamixel/dynamixel_calibration.py

@@ -1,14 +1,25 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 """Logic to calibrate a robot arm built with dynamixel motors"""
 # TODO(rcadene, aliberts): move this logic into the robot code when refactoring
 
 import numpy as np
 
-from lerobot.common.robot_devices.motors.dynamixel import (
-    CalibrationMode,
-    TorqueMode,
-    convert_degrees_to_steps,
-)
-from lerobot.common.robot_devices.motors.utils import MotorsBus
+from ..motors_bus import MotorNormMode, MotorsBus
+from .dynamixel import TorqueMode
+from .tables import MODEL_RESOLUTION
 
 URL_TEMPLATE = (
     "https://raw.githubusercontent.com/huggingface/lerobot/main/media/{robot}/{arm}_{position}.webp"
@@ -35,6 +46,17 @@ def apply_drive_mode(position, drive_mode):
     return position
 
 
+def convert_degrees_to_steps(degrees: float | np.ndarray, models: str | list[str]) -> np.ndarray:
+    """This function converts the degree range to the step range for indicating motors rotation.
+    It assumes a motor achieves a full rotation by going from -180 degree position to +180.
+    The motor resolution (e.g. 4096) corresponds to the number of steps needed to achieve a full rotation.
+    """
+    resolutions = [MODEL_RESOLUTION[model] for model in models]
+    steps = degrees / 180 * np.array(resolutions) / 2
+    steps = steps.astype(int)
+    return steps
+
+
 def compute_nearest_rounded_position(position, models):
     delta_turn = convert_degrees_to_steps(ROTATED_POSITION_DEGREE, models)
     nearest_pos = np.round(position.astype(float) / delta_turn) * delta_turn
@@ -75,11 +97,11 @@ def run_arm_calibration(arm: MotorsBus, robot_type: str, arm_name: str, arm_type
     # We arbitrarily chose our zero target position to be a straight horizontal position with gripper upwards and closed.
     # It is easy to identify and all motors are in a "quarter turn" position. Once calibration is done, this position will
     # correspond to every motor angle being 0. If you set all 0 as Goal Position, the arm will move in this position.
-    zero_target_pos = convert_degrees_to_steps(ZERO_POSITION_DEGREE, arm.motor_models)
+    zero_target_pos = convert_degrees_to_steps(ZERO_POSITION_DEGREE, arm.models)
 
     # Compute homing offset so that `present_position + homing_offset ~= target_position`.
     zero_pos = arm.read("Present_Position")
-    zero_nearest_pos = compute_nearest_rounded_position(zero_pos, arm.motor_models)
+    zero_nearest_pos = compute_nearest_rounded_position(zero_pos, arm.models)
     homing_offset = zero_target_pos - zero_nearest_pos
 
     # The rotated target position corresponds to a rotation of a quarter turn from the zero position.
@@ -93,7 +115,7 @@ def run_arm_calibration(arm: MotorsBus, robot_type: str, arm_name: str, arm_type
     print("See: " + URL_TEMPLATE.format(robot=robot_type, arm=arm_type, position="rotated"))
     input("Press Enter to continue...")
 
-    rotated_target_pos = convert_degrees_to_steps(ROTATED_POSITION_DEGREE, arm.motor_models)
+    rotated_target_pos = convert_degrees_to_steps(ROTATED_POSITION_DEGREE, arm.models)
 
     # Find drive mode by rotating each motor by a quarter of a turn.
     # Drive mode indicates if the motor rotation direction should be inverted (=1) or not (=0).
@@ -102,7 +124,7 @@ def run_arm_calibration(arm: MotorsBus, robot_type: str, arm_name: str, arm_type
 
     # Re-compute homing offset to take into account drive mode
     rotated_drived_pos = apply_drive_mode(rotated_pos, drive_mode)
-    rotated_nearest_pos = compute_nearest_rounded_position(rotated_drived_pos, arm.motor_models)
+    rotated_nearest_pos = compute_nearest_rounded_position(rotated_drived_pos, arm.models)
     homing_offset = rotated_target_pos - rotated_nearest_pos
 
     print("\nMove arm to rest position")
@@ -111,13 +133,13 @@ def run_arm_calibration(arm: MotorsBus, robot_type: str, arm_name: str, arm_type
     print()
 
     # Joints with rotational motions are expressed in degrees in nominal range of [-180, 180]
-    calib_mode = [CalibrationMode.DEGREE.name] * len(arm.motor_names)
+    calib_mode = [MotorNormMode.DEGREE.name] * len(arm.names)
 
     # TODO(rcadene): make type of joints (DEGREE or LINEAR) configurable from yaml?
-    if robot_type in ["aloha"] and "gripper" in arm.motor_names:
+    if robot_type in ["aloha"] and "gripper" in arm.names:
         # Joints with linear motions (like gripper of Aloha) are expressed in nominal range of [0, 100]
-        calib_idx = arm.motor_names.index("gripper")
-        calib_mode[calib_idx] = CalibrationMode.LINEAR.name
+        calib_idx = arm.names.index("gripper")
+        calib_mode[calib_idx] = MotorNormMode.LINEAR.name
 
     calib_data = {
         "homing_offset": homing_offset.tolist(),
@@ -125,6 +147,6 @@ def run_arm_calibration(arm: MotorsBus, robot_type: str, arm_name: str, arm_type
         "start_pos": zero_pos.tolist(),
         "end_pos": rotated_pos.tolist(),
         "calib_mode": calib_mode,
-        "motor_names": arm.motor_names,
+        "motor_names": arm.names,
     }
     return calib_data

lerobot/common/motors/dynamixel/tables.py

@@ -0,0 +1,162 @@
+# {data_name: (address, size_byte)}
+# https://emanual.robotis.com/docs/en/dxl/x/{MODEL}/#control-table
+X_SERIES_CONTROL_TABLE = {
+    "Model_Number": (0, 2),
+    "Model_Information": (2, 4),
+    "Firmware_Version": (6, 1),
+    "ID": (7, 1),
+    "Baud_Rate": (8, 1),
+    "Return_Delay_Time": (9, 1),
+    "Drive_Mode": (10, 1),
+    "Operating_Mode": (11, 1),
+    "Secondary_ID": (12, 1),
+    "Protocol_Type": (13, 1),
+    "Homing_Offset": (20, 4),
+    "Moving_Threshold": (24, 4),
+    "Temperature_Limit": (31, 1),
+    "Max_Voltage_Limit": (32, 2),
+    "Min_Voltage_Limit": (34, 2),
+    "PWM_Limit": (36, 2),
+    "Current_Limit": (38, 2),
+    "Acceleration_Limit": (40, 4),
+    "Velocity_Limit": (44, 4),
+    "Max_Position_Limit": (48, 4),
+    "Min_Position_Limit": (52, 4),
+    "Shutdown": (63, 1),
+    "Torque_Enable": (64, 1),
+    "LED": (65, 1),
+    "Status_Return_Level": (68, 1),
+    "Registered_Instruction": (69, 1),
+    "Hardware_Error_Status": (70, 1),
+    "Velocity_I_Gain": (76, 2),
+    "Velocity_P_Gain": (78, 2),
+    "Position_D_Gain": (80, 2),
+    "Position_I_Gain": (82, 2),
+    "Position_P_Gain": (84, 2),
+    "Feedforward_2nd_Gain": (88, 2),
+    "Feedforward_1st_Gain": (90, 2),
+    "Bus_Watchdog": (98, 1),
+    "Goal_PWM": (100, 2),
+    "Goal_Current": (102, 2),
+    "Goal_Velocity": (104, 4),
+    "Profile_Acceleration": (108, 4),
+    "Profile_Velocity": (112, 4),
+    "Goal_Position": (116, 4),
+    "Realtime_Tick": (120, 2),
+    "Moving": (122, 1),
+    "Moving_Status": (123, 1),
+    "Present_PWM": (124, 2),
+    "Present_Current": (126, 2),
+    "Present_Velocity": (128, 4),
+    "Present_Position": (132, 4),
+    "Velocity_Trajectory": (136, 4),
+    "Position_Trajectory": (140, 4),
+    "Present_Input_Voltage": (144, 2),
+    "Present_Temperature": (146, 1),
+}
+
+# https://emanual.robotis.com/docs/en/dxl/x/{MODEL}/#baud-rate8
+X_SERIES_BAUDRATE_TABLE = {
+    0: 9_600,
+    1: 57_600,
+    2: 115_200,
+    3: 1_000_000,
+    4: 2_000_000,
+    5: 3_000_000,
+    6: 4_000_000,
+}
+
+# {data_name: size_byte}
+X_SERIES_ENCODINGS_TABLE = {
+    "Homing_Offset": X_SERIES_CONTROL_TABLE["Homing_Offset"][1],
+    "Goal_PWM": X_SERIES_CONTROL_TABLE["Goal_PWM"][1],
+    "Goal_Current": X_SERIES_CONTROL_TABLE["Goal_Current"][1],
+    "Goal_Velocity": X_SERIES_CONTROL_TABLE["Goal_Velocity"][1],
+    "Present_PWM": X_SERIES_CONTROL_TABLE["Present_PWM"][1],
+    "Present_Current": X_SERIES_CONTROL_TABLE["Present_Current"][1],
+    "Present_Velocity": X_SERIES_CONTROL_TABLE["Present_Velocity"][1],
+}
+
+MODEL_ENCODING_TABLE = {
+    "x_series": X_SERIES_ENCODINGS_TABLE,
+    "xl330-m077": X_SERIES_ENCODINGS_TABLE,
+    "xl330-m288": X_SERIES_ENCODINGS_TABLE,
+    "xl430-w250": X_SERIES_ENCODINGS_TABLE,
+    "xm430-w350": X_SERIES_ENCODINGS_TABLE,
+    "xm540-w270": X_SERIES_ENCODINGS_TABLE,
+    "xc430-w150": X_SERIES_ENCODINGS_TABLE,
+}
+
+# {model: model_resolution}
+# https://emanual.robotis.com/docs/en/dxl/x/{MODEL}/#specifications
+MODEL_RESOLUTION = {
+    "x_series": 4096,
+    "xl330-m077": 4096,
+    "xl330-m288": 4096,
+    "xl430-w250": 4096,
+    "xm430-w350": 4096,
+    "xm540-w270": 4096,
+    "xc430-w150": 4096,
+}
+
+# {model: model_number}
+# https://emanual.robotis.com/docs/en/dxl/x/{MODEL}/#control-table-of-eeprom-area
+MODEL_NUMBER_TABLE = {
+    "xl330-m077": 1190,
+    "xl330-m288": 1200,
+    "xl430-w250": 1060,
+    "xm430-w350": 1020,
+    "xm540-w270": 1120,
+    "xc430-w150": 1070,
+}
+
+# {model: available_operating_modes}
+# https://emanual.robotis.com/docs/en/dxl/x/{MODEL}/#operating-mode11
+MODEL_OPERATING_MODES = {
+    "xl330-m077": [0, 1, 3, 4, 5, 16],
+    "xl330-m288": [0, 1, 3, 4, 5, 16],
+    "xl430-w250": [1, 3, 4, 16],
+    "xm430-w350": [0, 1, 3, 4, 5, 16],
+    "xm540-w270": [0, 1, 3, 4, 5, 16],
+    "xc430-w150": [1, 3, 4, 16],
+}
+
+MODEL_CONTROL_TABLE = {
+    "x_series": X_SERIES_CONTROL_TABLE,
+    "xl330-m077": X_SERIES_CONTROL_TABLE,
+    "xl330-m288": X_SERIES_CONTROL_TABLE,
+    "xl430-w250": X_SERIES_CONTROL_TABLE,
+    "xm430-w350": X_SERIES_CONTROL_TABLE,
+    "xm540-w270": X_SERIES_CONTROL_TABLE,
+    "xc430-w150": X_SERIES_CONTROL_TABLE,
+}
+
+MODEL_BAUDRATE_TABLE = {
+    "x_series": X_SERIES_BAUDRATE_TABLE,
+    "xl330-m077": X_SERIES_BAUDRATE_TABLE,
+    "xl330-m288": X_SERIES_BAUDRATE_TABLE,
+    "xl430-w250": X_SERIES_BAUDRATE_TABLE,
+    "xm430-w350": X_SERIES_BAUDRATE_TABLE,
+    "xm540-w270": X_SERIES_BAUDRATE_TABLE,
+    "xc430-w150": X_SERIES_BAUDRATE_TABLE,
+}
+
+AVAILABLE_BAUDRATES = [
+    9_600,
+    19_200,
+    38_400,
+    57_600,
+    115_200,
+    230_400,
+    460_800,
+    500_000,
+    576_000,
+    921_600,
+    1_000_000,
+    1_152_000,
+    2_000_000,
+    2_500_000,
+    3_000_000,
+    3_500_000,
+    4_000_000,
+]

lerobot/common/motors/feetech/__init__.py

@@ -0,0 +1,2 @@
+from .feetech import DriveMode, FeetechMotorsBus, OperatingMode, TorqueMode
+from .tables import *

lerobot/common/motors/feetech/feetech.py

@@ -0,0 +1,367 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+from copy import deepcopy
+from enum import Enum
+from pprint import pformat
+
+from lerobot.common.utils.encoding_utils import decode_sign_magnitude, encode_sign_magnitude
+
+from ..motors_bus import Motor, MotorCalibration, MotorsBus, NameOrID, Value
+from .tables import (
+    FIRMWARE_MAJOR_VERSION,
+    FIRMWARE_MINOR_VERSION,
+    MODEL_BAUDRATE_TABLE,
+    MODEL_CONTROL_TABLE,
+    MODEL_ENCODING_TABLE,
+    MODEL_NUMBER,
+    MODEL_NUMBER_TABLE,
+    MODEL_PROTOCOL,
+    MODEL_RESOLUTION,
+    SCAN_BAUDRATES,
+)
+
+DEFAULT_PROTOCOL_VERSION = 0
+BAUDRATE = 1_000_000
+DEFAULT_TIMEOUT_MS = 1000
+
+NORMALIZED_DATA = ["Goal_Position", "Present_Position"]
+
+logger = logging.getLogger(__name__)
+
+
+class OperatingMode(Enum):
+    # position servo mode
+    POSITION = 0
+    # The motor is in constant speed mode, which is controlled by parameter 0x2e, and the highest bit 15 is
+    # the direction bit
+    VELOCITY = 1
+    # PWM open-loop speed regulation mode, with parameter 0x2c running time parameter control, bit11 as
+    # direction bit
+    PWM = 2
+    # In step servo mode, the number of step progress is represented by parameter 0x2a, and the highest bit 15
+    # is the direction bit
+    STEP = 3
+
+
+class DriveMode(Enum):
+    NON_INVERTED = 0
+    INVERTED = 1
+
+
+class TorqueMode(Enum):
+    ENABLED = 1
+    DISABLED = 0
+
+
+def _split_into_byte_chunks(value: int, length: int) -> list[int]:
+    import scservo_sdk as scs
+
+    if length == 1:
+        data = [value]
+    elif length == 2:
+        data = [scs.SCS_LOBYTE(value), scs.SCS_HIBYTE(value)]
+    elif length == 4:
+        data = [
+            scs.SCS_LOBYTE(scs.SCS_LOWORD(value)),
+            scs.SCS_HIBYTE(scs.SCS_LOWORD(value)),
+            scs.SCS_LOBYTE(scs.SCS_HIWORD(value)),
+            scs.SCS_HIBYTE(scs.SCS_HIWORD(value)),
+        ]
+    return data
+
+
+def patch_setPacketTimeout(self, packet_length):  # noqa: N802
+    """
+    HACK: This patches the PortHandler behavior to set the correct packet timeouts.
+
+    It fixes https://gitee.com/ftservo/SCServoSDK/issues/IBY2S6
+    The bug is fixed on the official Feetech SDK repo (https://gitee.com/ftservo/FTServo_Python)
+    but because that version is not published on PyPI, we rely on the (unofficial) on that is, which needs
+    patching.
+    """
+    self.packet_start_time = self.getCurrentTime()
+    self.packet_timeout = (self.tx_time_per_byte * packet_length) + (self.tx_time_per_byte * 3.0) + 50
+
+
+class FeetechMotorsBus(MotorsBus):
+    """
+    The FeetechMotorsBus class allows to efficiently read and write to the attached motors. It relies on the
+    python feetech sdk to communicate with the motors, which is itself based on the dynamixel sdk.
+    """
+
+    available_baudrates = deepcopy(SCAN_BAUDRATES)
+    default_timeout = DEFAULT_TIMEOUT_MS
+    model_baudrate_table = deepcopy(MODEL_BAUDRATE_TABLE)
+    model_ctrl_table = deepcopy(MODEL_CONTROL_TABLE)
+    model_encoding_table = deepcopy(MODEL_ENCODING_TABLE)
+    model_number_table = deepcopy(MODEL_NUMBER_TABLE)
+    model_resolution_table = deepcopy(MODEL_RESOLUTION)
+    normalized_data = deepcopy(NORMALIZED_DATA)
+
+    def __init__(
+        self,
+        port: str,
+        motors: dict[str, Motor],
+        calibration: dict[str, MotorCalibration] | None = None,
+        protocol_version: int = DEFAULT_PROTOCOL_VERSION,
+    ):
+        super().__init__(port, motors, calibration)
+        self.protocol_version = protocol_version
+        self._assert_same_protocol()
+        import scservo_sdk as scs
+
+        self.port_handler = scs.PortHandler(self.port)
+        # HACK: monkeypatch
+        self.port_handler.setPacketTimeout = patch_setPacketTimeout.__get__(
+            self.port_handler, scs.PortHandler
+        )
+        self.packet_handler = scs.PacketHandler(protocol_version)
+        self.sync_reader = scs.GroupSyncRead(self.port_handler, self.packet_handler, 0, 0)
+        self.sync_writer = scs.GroupSyncWrite(self.port_handler, self.packet_handler, 0, 0)
+        self._comm_success = scs.COMM_SUCCESS
+        self._no_error = 0x00
+
+        if any(MODEL_PROTOCOL[model] != self.protocol_version for model in self.models):
+            raise ValueError(f"Some motors are incompatible with protocol_version={self.protocol_version}")
+
+    def _assert_same_protocol(self) -> None:
+        if any(MODEL_PROTOCOL[model] != self.protocol_version for model in self.models):
+            raise RuntimeError("Some motors use an incompatible protocol.")
+
+    def _assert_protocol_is_compatible(self, instruction_name: str) -> None:
+        if instruction_name == "sync_read" and self.protocol_version == 1:
+            raise NotImplementedError(
+                "'Sync Read' is not available with Feetech motors using Protocol 1. Use 'Read' sequentially instead."
+            )
+        if instruction_name == "broadcast_ping" and self.protocol_version == 1:
+            raise NotImplementedError(
+                "'Broadcast Ping' is not available with Feetech motors using Protocol 1. Use 'Ping' sequentially instead."
+            )
+
+    def _assert_same_firmware(self) -> None:
+        firmware_versions = self._read_firmware_version(self.ids)
+        if len(set(firmware_versions.values())) != 1:
+            raise RuntimeError(
+                "Some Motors use different firmware versions. Update their firmware first using Feetech's software. "
+                "Visit https://www.feetechrc.com/software."
+            )
+
+    def _handshake(self) -> None:
+        self._assert_motors_exist()
+        self._assert_same_firmware()
+
+    def configure_motors(self) -> None:
+        for motor in self.motors:
+            # By default, Feetech motors have a 500µs delay response time (corresponding to a value of 250 on
+            # the 'Return_Delay_Time' address). We ensure this is reduced to the minimum of 2µs (value of 0).
+            self.write("Return_Delay_Time", motor, 0)
+            # Set 'Maximum_Acceleration' to 254 to speedup acceleration and deceleration of the motors.
+            # Note: this address is not in the official STS3215 Memory Table
+            self.write("Maximum_Acceleration", motor, 254)
+            self.write("Acceleration", motor, 254)
+
+    def _get_half_turn_homings(self, positions: dict[NameOrID, Value]) -> dict[NameOrID, Value]:
+        """
+        On Feetech Motors:
+        Present_Position = Actual_Position - Homing_Offset
+        """
+        half_turn_homings = {}
+        for motor, pos in positions.items():
+            model = self._get_motor_model(motor)
+            max_res = self.model_resolution_table[model] - 1
+            half_turn_homings[motor] = pos - int(max_res / 2)
+
+        return half_turn_homings
+
+    def disable_torque(self, motors: str | list[str] | None = None, num_retry: int = 0) -> None:
+        for name in self._get_motors_list(motors):
+            self.write("Torque_Enable", name, TorqueMode.DISABLED.value, num_retry=num_retry)
+            self.write("Lock", name, 0, num_retry=num_retry)
+
+    def enable_torque(self, motors: str | list[str] | None = None, num_retry: int = 0) -> None:
+        for name in self._get_motors_list(motors):
+            self.write("Torque_Enable", name, TorqueMode.ENABLED.value, num_retry=num_retry)
+            self.write("Lock", name, 1, num_retry=num_retry)
+
+    def _encode_sign(self, data_name: str, ids_values: dict[int, int]) -> dict[int, int]:
+        for id_ in ids_values:
+            model = self._id_to_model(id_)
+            encoding_table = self.model_encoding_table.get(model)
+            if encoding_table and data_name in encoding_table:
+                sign_bit = encoding_table[data_name]
+                ids_values[id_] = encode_sign_magnitude(ids_values[id_], sign_bit)
+
+        return ids_values
+
+    def _decode_sign(self, data_name: str, ids_values: dict[int, int]) -> dict[int, int]:
+        for id_ in ids_values:
+            model = self._id_to_model(id_)
+            encoding_table = self.model_encoding_table.get(model)
+            if encoding_table and data_name in encoding_table:
+                sign_bit = encoding_table[data_name]
+                ids_values[id_] = decode_sign_magnitude(ids_values[id_], sign_bit)
+
+        return ids_values
+
+    def _split_into_byte_chunks(self, value: int, length: int) -> list[int]:
+        return _split_into_byte_chunks(value, length)
+
+    def _broadcast_ping_p1(
+        self, known_motors_only: bool = True, n_motors: int | None = None, num_retry: int = 0
+    ) -> dict[int, int]:
+        if known_motors_only:
+            ids = self.ids
+        else:
+            import scservo_sdk as scs
+
+            ids = range(scs.MAX_ID + 1)
+
+        ids_models = {}
+        motors_found = 0
+        for id_ in ids:
+            model_number = self.ping(id_, num_retry)
+            if model_number is not None:
+                ids_models[id_] = model_number
+                motors_found += 1
+                if motors_found >= n_motors:
+                    break
+
+        return ids_models
+
+    def _broadcast_ping_p0(self) -> tuple[dict[int, int], int]:
+        import scservo_sdk as scs
+
+        data_list = {}
+
+        status_length = 6
+
+        rx_length = 0
+        wait_length = status_length * scs.MAX_ID
+
+        txpacket = [0] * 6
+
+        tx_time_per_byte = (1000.0 / self.port_handler.getBaudRate()) * 10.0
+
+        txpacket[scs.PKT_ID] = scs.BROADCAST_ID
+        txpacket[scs.PKT_LENGTH] = 2
+        txpacket[scs.PKT_INSTRUCTION] = scs.INST_PING
+
+        result = self.packet_handler.txPacket(self.port_handler, txpacket)
+        if result != scs.COMM_SUCCESS:
+            self.port_handler.is_using = False
+            return data_list, result
+
+        # set rx timeout
+        self.port_handler.setPacketTimeoutMillis((wait_length * tx_time_per_byte) + (3.0 * scs.MAX_ID) + 16.0)
+
+        rxpacket = []
+        while True:
+            rxpacket += self.port_handler.readPort(wait_length - rx_length)
+            rx_length = len(rxpacket)
+
+            if self.port_handler.isPacketTimeout():  # or rx_length >= wait_length
+                break
+
+        self.port_handler.is_using = False
+
+        if rx_length == 0:
+            return data_list, scs.COMM_RX_TIMEOUT
+
+        while True:
+            if rx_length < status_length:
+                return data_list, scs.COMM_RX_CORRUPT
+
+            # find packet header
+            for idx in range(0, (rx_length - 1)):
+                if (rxpacket[idx] == 0xFF) and (rxpacket[idx + 1] == 0xFF):
+                    break
+
+            if idx == 0:  # found at the beginning of the packet
+                # calculate checksum
+                checksum = 0
+                for idx in range(2, status_length - 1):  # except header & checksum
+                    checksum += rxpacket[idx]
+
+                checksum = ~checksum & 0xFF
+                if rxpacket[status_length - 1] == checksum:
+                    result = scs.COMM_SUCCESS
+                    data_list[rxpacket[scs.PKT_ID]] = rxpacket[scs.PKT_ERROR]
+
+                    del rxpacket[0:status_length]
+                    rx_length = rx_length - status_length
+
+                    if rx_length == 0:
+                        return data_list, result
+                else:
+                    result = scs.COMM_RX_CORRUPT
+                    # remove header (0xFF 0xFF)
+                    del rxpacket[0:2]
+                    rx_length = rx_length - 2
+            else:
+                # remove unnecessary packets
+                del rxpacket[0:idx]
+                rx_length = rx_length - idx
+
+    def broadcast_ping(self, num_retry: int = 0, raise_on_error: bool = False) -> dict[int, int] | None:
+        self._assert_protocol_is_compatible("broadcast_ping")
+        for n_try in range(1 + num_retry):
+            ids_status, comm = self._broadcast_ping_p0()
+            if self._is_comm_success(comm):
+                break
+            logger.debug(f"Broadcast ping failed on port '{self.port}' ({n_try=})")
+            logger.debug(self.packet_handler.getTxRxResult(comm))
+
+        if not self._is_comm_success(comm):
+            if raise_on_error:
+                raise ConnectionError(self.packet_handler.getTxRxResult(comm))
+            return
+
+        ids_errors = {id_: status for id_, status in ids_status.items() if self._is_error(status)}
+        if ids_errors:
+            display_dict = {id_: self.packet_handler.getRxPacketError(err) for id_, err in ids_errors.items()}
+            logger.error(f"Some motors found returned an error status:\n{pformat(display_dict, indent=4)}")
+
+        return self._read_model_number(list(ids_status), raise_on_error)
+
+    def _read_firmware_version(self, motor_ids: list[int], raise_on_error: bool = False) -> dict[int, str]:
+        firmware_versions = {}
+        for id_ in motor_ids:
+            firm_ver_major, comm, error = self._read(
+                *FIRMWARE_MAJOR_VERSION, id_, raise_on_error=raise_on_error
+            )
+            if not self._is_comm_success(comm) or self._is_error(error):
+                return
+
+            firm_ver_minor, comm, error = self._read(
+                *FIRMWARE_MINOR_VERSION, id_, raise_on_error=raise_on_error
+            )
+            if not self._is_comm_success(comm) or self._is_error(error):
+                return
+
+            firmware_versions[id_] = f"{firm_ver_major}.{firm_ver_minor}"
+
+        return firmware_versions
+
+    def _read_model_number(self, motor_ids: list[int], raise_on_error: bool = False) -> dict[int, int]:
+        model_numbers = {}
+        for id_ in motor_ids:
+            model_nb, comm, error = self._read(*MODEL_NUMBER, id_, raise_on_error=raise_on_error)
+            if not self._is_comm_success(comm) or self._is_error(error):
+                return
+
+            model_numbers[id_] = model_nb
+
+        return model_numbers

lerobot/common/motors/feetech/tables.py

@@ -0,0 +1,202 @@
+FIRMWARE_MAJOR_VERSION = (0, 1)
+FIRMWARE_MINOR_VERSION = (1, 1)
+MODEL_NUMBER = (3, 2)
+
+# See this link for STS3215 Memory Table:
+# https://docs.google.com/spreadsheets/d/1GVs7W1VS1PqdhA1nW-abeyAHhTUxKUdR/edit?usp=sharing&ouid=116566590112741600240&rtpof=true&sd=true
+# data_name: (address, size_byte)
+STS_SMS_SERIES_CONTROL_TABLE = {
+    # EPROM
+    "Firmware_Major_Version": FIRMWARE_MAJOR_VERSION,  # read-only
+    "Firmware_Minor_Version": FIRMWARE_MINOR_VERSION,  # read-only
+    "Model_Number": MODEL_NUMBER,  # read-only
+    "ID": (5, 1),
+    "Baud_Rate": (6, 1),
+    "Return_Delay_Time": (7, 1),
+    "Response_Status_Level": (8, 1),
+    "Min_Position_Limit": (9, 2),
+    "Max_Position_Limit": (11, 2),
+    "Max_Temperature_Limit": (13, 1),
+    "Max_Voltage_Limit": (14, 1),
+    "Min_Voltage_Limit": (15, 1),
+    "Max_Torque_Limit": (16, 2),
+    "Phase": (18, 1),
+    "Unloading_Condition": (19, 1),
+    "LED_Alarm_Condition": (20, 1),
+    "P_Coefficient": (21, 1),
+    "D_Coefficient": (22, 1),
+    "I_Coefficient": (23, 1),
+    "Minimum_Startup_Force": (24, 2),
+    "CW_Dead_Zone": (26, 1),
+    "CCW_Dead_Zone": (27, 1),
+    "Protection_Current": (28, 2),
+    "Angular_Resolution": (30, 1),
+    "Homing_Offset": (31, 2),
+    "Operating_Mode": (33, 1),
+    "Protective_Torque": (34, 1),
+    "Protection_Time": (35, 1),
+    "Overload_Torque": (36, 1),
+    "Speed_closed_loop_P_proportional_coefficient": (37, 1),
+    "Over_Current_Protection_Time": (38, 1),
+    "Velocity_closed_loop_I_integral_coefficient": (39, 1),
+    # SRAM
+    "Torque_Enable": (40, 1),
+    "Acceleration": (41, 1),
+    "Goal_Position": (42, 2),
+    "Goal_Time": (44, 2),
+    "Goal_Speed": (46, 2),
+    "Torque_Limit": (48, 2),
+    "Lock": (55, 1),
+    "Present_Position": (56, 2),  # read-only
+    "Present_Speed": (58, 2),  # read-only
+    "Present_Load": (60, 2),  # read-only
+    "Present_Voltage": (62, 1),  # read-only
+    "Present_Temperature": (63, 1),  # read-only
+    "Status": (65, 1),  # read-only
+    "Moving": (66, 1),  # read-only
+    "Present_Current": (69, 2),  # read-only
+    # Not in the Memory Table
+    "Maximum_Acceleration": (85, 2),
+}
+
+SCS_SERIES_CONTROL_TABLE = {
+    # EPROM
+    "Firmware_Major_Version": FIRMWARE_MAJOR_VERSION,  # read-only
+    "Firmware_Minor_Version": FIRMWARE_MINOR_VERSION,  # read-only
+    "Model_Number": MODEL_NUMBER,  # read-only
+    "ID": (5, 1),
+    "Baud_Rate": (6, 1),
+    "Return_Delay": (7, 1),
+    "Response_Status_Level": (8, 1),
+    "Min_Position_Limit": (9, 2),
+    "Max_Position_Limit": (11, 2),
+    "Max_Temperature_Limit": (13, 1),
+    "Max_Voltage_Limit": (14, 1),
+    "Min_Voltage_Limit": (15, 1),
+    "Max_Torque_Limit": (16, 2),
+    "Phase": (18, 1),
+    "Unloading_Condition": (19, 1),
+    "LED_Alarm_Condition": (20, 1),
+    "P_Coefficient": (21, 1),
+    "D_Coefficient": (22, 1),
+    "I_Coefficient": (23, 1),
+    "Minimum_Startup_Force": (24, 2),
+    "CW_Dead_Zone": (26, 1),
+    "CCW_Dead_Zone": (27, 1),
+    "Protective_Torque": (37, 1),
+    "Protection_Time": (38, 1),
+    # SRAM
+    "Torque_Enable": (40, 1),
+    "Acceleration": (41, 1),
+    "Goal_Position": (42, 2),
+    "Running_Time": (44, 2),
+    "Goal_Speed": (46, 2),
+    "Lock": (48, 1),
+    "Present_Position": (56, 2),  # read-only
+    "Present_Speed": (58, 2),  # read-only
+    "Present_Load": (60, 2),  # read-only
+    "Present_Voltage": (62, 1),  # read-only
+    "Present_Temperature": (63, 1),  # read-only
+    "Sync_Write_Flag": (64, 1),  # read-only
+    "Status": (65, 1),  # read-only
+    "Moving": (66, 1),  # read-only
+}
+
+STS_SMS_SERIES_BAUDRATE_TABLE = {
+    0: 1_000_000,
+    1: 500_000,
+    2: 250_000,
+    3: 128_000,
+    4: 115_200,
+    5: 57_600,
+    6: 38_400,
+    7: 19_200,
+}
+
+SCS_SERIES_BAUDRATE_TABLE = {
+    0: 1_000_000,
+    1: 500_000,
+    2: 250_000,
+    3: 128_000,
+    4: 115_200,
+    5: 57_600,
+    6: 38_400,
+    7: 19_200,
+}
+
+MODEL_CONTROL_TABLE = {
+    "sts_series": STS_SMS_SERIES_CONTROL_TABLE,
+    "scs_series": SCS_SERIES_CONTROL_TABLE,
+    "sms_series": STS_SMS_SERIES_CONTROL_TABLE,
+    "sts3215": STS_SMS_SERIES_CONTROL_TABLE,
+    "sts3250": STS_SMS_SERIES_CONTROL_TABLE,
+    "scs0009": SCS_SERIES_CONTROL_TABLE,
+    "sm8512bl": STS_SMS_SERIES_CONTROL_TABLE,
+}
+
+MODEL_RESOLUTION = {
+    "sts_series": 4096,
+    "sms_series": 4096,
+    "scs_series": 1024,
+    "sts3215": 4096,
+    "sts3250": 4096,
+    "sm8512bl": 65536,
+    "scs0009": 1024,
+}
+
+MODEL_BAUDRATE_TABLE = {
+    "sts_series": STS_SMS_SERIES_BAUDRATE_TABLE,
+    "sms_series": STS_SMS_SERIES_BAUDRATE_TABLE,
+    "scs_series": SCS_SERIES_BAUDRATE_TABLE,
+    "sm8512bl": STS_SMS_SERIES_BAUDRATE_TABLE,
+    "sts3215": STS_SMS_SERIES_BAUDRATE_TABLE,
+    "sts3250": STS_SMS_SERIES_BAUDRATE_TABLE,
+    "scs0009": SCS_SERIES_BAUDRATE_TABLE,
+}
+
+# Sign-Magnitude encoding bits
+STS_SMS_SERIES_ENCODINGS_TABLE = {
+    "Homing_Offset": 11,
+    "Goal_Speed": 15,
+}
+
+MODEL_ENCODING_TABLE = {
+    "sts_series": STS_SMS_SERIES_ENCODINGS_TABLE,
+    "sms_series": STS_SMS_SERIES_ENCODINGS_TABLE,
+    "scs_series": {},
+    "sts3215": STS_SMS_SERIES_ENCODINGS_TABLE,
+    "sts3250": STS_SMS_SERIES_ENCODINGS_TABLE,
+    "sm8512bl": STS_SMS_SERIES_ENCODINGS_TABLE,
+    "scs0009": {},
+}
+
+SCAN_BAUDRATES = [
+    4_800,
+    9_600,
+    14_400,
+    19_200,
+    38_400,
+    57_600,
+    115_200,
+    128_000,
+    250_000,
+    500_000,
+    1_000_000,
+]
+
+MODEL_NUMBER_TABLE = {
+    "sts3215": 777,
+    "sts3250": 2825,
+    "sm8512bl": 11272,
+    "scs0009": 1284,
+}
+
+MODEL_PROTOCOL = {
+    "sts_series": 0,
+    "sms_series": 0,
+    "scs_series": 1,
+    "sts3215": 0,
+    "sts3250": 0,
+    "sm8512bl": 0,
+    "scs0009": 1,
+}

lerobot/common/motors/motors_bus.py

@@ -0,0 +1,987 @@
+#!/usr/bin/env python
+
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# ruff: noqa: N802
+# This noqa is for the Protocols classes: PortHandler, PacketHandler GroupSyncRead/Write
+# TODO(aliberts): Add block noqa when feature below is available
+# https://github.com/astral-sh/ruff/issues/3711
+
+import abc
+import logging
+from contextlib import contextmanager
+from dataclasses import dataclass
+from enum import Enum
+from functools import cached_property
+from pprint import pformat
+from typing import Protocol, TypeAlias
+
+import serial
+from deepdiff import DeepDiff
+from tqdm import tqdm
+
+from lerobot.common.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
+from lerobot.common.utils.utils import enter_pressed, move_cursor_up
+
+NameOrID: TypeAlias = str | int
+Value: TypeAlias = int | float
+
+MAX_ID_RANGE = 252
+
+logger = logging.getLogger(__name__)
+
+
+def get_ctrl_table(model_ctrl_table: dict[str, dict], model: str) -> dict[str, tuple[int, int]]:
+    ctrl_table = model_ctrl_table.get(model)
+    if ctrl_table is None:
+        raise KeyError(f"Control table for {model=} not found.")
+    return ctrl_table
+
+
+def get_address(model_ctrl_table: dict[str, dict], model: str, data_name: str) -> tuple[int, int]:
+    ctrl_table = get_ctrl_table(model_ctrl_table, model)
+    addr_bytes = ctrl_table.get(data_name)
+    if addr_bytes is None:
+        raise KeyError(f"Address for '{data_name}' not found in {model} control table.")
+    return addr_bytes
+
+
+def assert_same_address(model_ctrl_table: dict[str, dict], motor_models: list[str], data_name: str) -> None:
+    all_addr = []
+    all_bytes = []
+    for model in motor_models:
+        addr, bytes = get_address(model_ctrl_table, model, data_name)
+        all_addr.append(addr)
+        all_bytes.append(bytes)
+
+    if len(set(all_addr)) != 1:
+        raise NotImplementedError(
+            f"At least two motor models use a different address for `data_name`='{data_name}'"
+            f"({list(zip(motor_models, all_addr, strict=False))})."
+        )
+
+    if len(set(all_bytes)) != 1:
+        raise NotImplementedError(
+            f"At least two motor models use a different bytes representation for `data_name`='{data_name}'"
+            f"({list(zip(motor_models, all_bytes, strict=False))})."
+        )
+
+
+class MotorNormMode(Enum):
+    DEGREE = 0
+    RANGE_0_100 = 1
+    RANGE_M100_100 = 2
+    VELOCITY = 3
+
+
+@dataclass
+class MotorCalibration:
+    id: int
+    drive_mode: int
+    homing_offset: int
+    range_min: int
+    range_max: int
+
+
+@dataclass
+class Motor:
+    id: int
+    model: str
+    norm_mode: MotorNormMode
+
+
+class JointOutOfRangeError(Exception):
+    def __init__(self, message="Joint is out of range"):
+        self.message = message
+        super().__init__(self.message)
+
+
+class PortHandler(Protocol):
+    def __init__(self, port_name):
+        self.is_open: bool
+        self.baudrate: int
+        self.packet_start_time: float
+        self.packet_timeout: float
+        self.tx_time_per_byte: float
+        self.is_using: bool
+        self.port_name: str
+        self.ser: serial.Serial
+
+    def openPort(self): ...
+    def closePort(self): ...
+    def clearPort(self): ...
+    def setPortName(self, port_name): ...
+    def getPortName(self): ...
+    def setBaudRate(self, baudrate): ...
+    def getBaudRate(self): ...
+    def getBytesAvailable(self): ...
+    def readPort(self, length): ...
+    def writePort(self, packet): ...
+    def setPacketTimeout(self, packet_length): ...
+    def setPacketTimeoutMillis(self, msec): ...
+    def isPacketTimeout(self): ...
+    def getCurrentTime(self): ...
+    def getTimeSinceStart(self): ...
+    def setupPort(self, cflag_baud): ...
+    def getCFlagBaud(self, baudrate): ...
+
+
+class PacketHandler(Protocol):
+    def getTxRxResult(self, result): ...
+    def getRxPacketError(self, error): ...
+    def txPacket(self, port, txpacket): ...
+    def rxPacket(self, port): ...
+    def txRxPacket(self, port, txpacket): ...
+    def ping(self, port, id): ...
+    def action(self, port, id): ...
+    def readTx(self, port, id, address, length): ...
+    def readRx(self, port, id, length): ...
+    def readTxRx(self, port, id, address, length): ...
+    def read1ByteTx(self, port, id, address): ...
+    def read1ByteRx(self, port, id): ...
+    def read1ByteTxRx(self, port, id, address): ...
+    def read2ByteTx(self, port, id, address): ...
+    def read2ByteRx(self, port, id): ...
+    def read2ByteTxRx(self, port, id, address): ...
+    def read4ByteTx(self, port, id, address): ...
+    def read4ByteRx(self, port, id): ...
+    def read4ByteTxRx(self, port, id, address): ...
+    def writeTxOnly(self, port, id, address, length, data): ...
+    def writeTxRx(self, port, id, address, length, data): ...
+    def write1ByteTxOnly(self, port, id, address, data): ...
+    def write1ByteTxRx(self, port, id, address, data): ...
+    def write2ByteTxOnly(self, port, id, address, data): ...
+    def write2ByteTxRx(self, port, id, address, data): ...
+    def write4ByteTxOnly(self, port, id, address, data): ...
+    def write4ByteTxRx(self, port, id, address, data): ...
+    def regWriteTxOnly(self, port, id, address, length, data): ...
+    def regWriteTxRx(self, port, id, address, length, data): ...
+    def syncReadTx(self, port, start_address, data_length, param, param_length): ...
+    def syncWriteTxOnly(self, port, start_address, data_length, param, param_length): ...
+
+
+class GroupSyncRead(Protocol):
+    def __init__(self, port, ph, start_address, data_length):
+        self.port: str
+        self.ph: PortHandler
+        self.start_address: int
+        self.data_length: int
+        self.last_result: bool
+        self.is_param_changed: bool
+        self.param: list
+        self.data_dict: dict
+
+    def makeParam(self): ...
+    def addParam(self, id): ...
+    def removeParam(self, id): ...
+    def clearParam(self): ...
+    def txPacket(self): ...
+    def rxPacket(self): ...
+    def txRxPacket(self): ...
+    def isAvailable(self, id, address, data_length): ...
+    def getData(self, id, address, data_length): ...
+
+
+class GroupSyncWrite(Protocol):
+    def __init__(self, port, ph, start_address, data_length):
+        self.port: str
+        self.ph: PortHandler
+        self.start_address: int
+        self.data_length: int
+        self.is_param_changed: bool
+        self.param: list
+        self.data_dict: dict
+
+    def makeParam(self): ...
+    def addParam(self, id, data): ...
+    def removeParam(self, id): ...
+    def changeParam(self, id, data): ...
+    def clearParam(self): ...
+    def txPacket(self): ...
+
+
+class MotorsBus(abc.ABC):
+    """The main LeRobot class for implementing motors buses.
+
+    There are currently two implementations of this abstract class:
+        - DynamixelMotorsBus
+        - FeetechMotorsBus
+
+    Note: This class may evolve in the future should we add support for other manufacturers SDKs.
+
+    A MotorsBus allows to efficiently read and write to the attached motors.
+    It represents several motors daisy-chained together and connected through a serial port.
+
+    A MotorsBus subclass instance requires a port (e.g. `FeetechMotorsBus(port="/dev/tty.usbmodem575E0031751"`)).
+    To find the port, you can run our utility script:
+    ```bash
+    python lerobot/scripts/find_motors_bus_port.py
+    >>> Finding all available ports for the MotorsBus.
+    >>> ['/dev/tty.usbmodem575E0032081', '/dev/tty.usbmodem575E0031751']
+    >>> Remove the usb cable from your MotorsBus and press Enter when done.
+    >>> The port of this MotorsBus is /dev/tty.usbmodem575E0031751.
+    >>> Reconnect the usb cable.
+    ```
+
+    Example of usage for 1 Feetech sts3215 motor connected to the bus:
+    ```python
+    motors_bus = FeetechMotorsBus(
+        port="/dev/tty.usbmodem575E0031751",
+        motors={"gripper": (6, "sts3215")},
+    )
+    motors_bus.connect()
+
+    position = motors_bus.read("Present_Position")
+
+    # Move from a few motor steps as an example
+    few_steps = 30
+    motors_bus.write("Goal_Position", position + few_steps)
+
+    # When done, properly disconnect the port using
+    motors_bus.disconnect()
+    ```
+    """
+
+    available_baudrates: list[int]
+    default_timeout: int
+    model_baudrate_table: dict[str, dict]
+    model_ctrl_table: dict[str, dict]
+    model_encoding_table: dict[str, dict]
+    model_number_table: dict[str, int]
+    model_resolution_table: dict[str, int]
+    normalized_data: list[str]
+
+    def __init__(
+        self,
+        port: str,
+        motors: dict[str, Motor],
+        calibration: dict[str, MotorCalibration] | None = None,
+    ):
+        self.port = port
+        self.motors = motors
+        self.calibration = calibration if calibration else {}
+
+        self.port_handler: PortHandler
+        self.packet_handler: PacketHandler
+        self.sync_reader: GroupSyncRead
+        self.sync_writer: GroupSyncWrite
+        self._comm_success: int
+        self._no_error: int
+
+        self._id_to_model_dict = {m.id: m.model for m in self.motors.values()}
+        self._id_to_name_dict = {m.id: name for name, m in self.motors.items()}
+        self._model_nb_to_model_dict = {v: k for k, v in self.model_number_table.items()}
+
+        self._validate_motors()
+
+    def __len__(self):
+        return len(self.motors)
+
+    def __repr__(self):
+        return (
+            f"{self.__class__.__name__}(\n"
+            f"    Port: '{self.port}',\n"
+            f"    Motors: \n{pformat(self.motors, indent=8, sort_dicts=False)},\n"
+            ")',\n"
+        )
+
+    @cached_property
+    def _has_different_ctrl_tables(self) -> bool:
+        if len(self.models) < 2:
+            return False
+
+        first_table = self.model_ctrl_table[self.models[0]]
+        return any(
+            DeepDiff(first_table, get_ctrl_table(self.model_ctrl_table, model)) for model in self.models[1:]
+        )
+
+    @cached_property
+    def names(self) -> list[str]:
+        return list(self.motors)
+
+    @cached_property
+    def models(self) -> list[str]:
+        return [m.model for m in self.motors.values()]
+
+    @cached_property
+    def ids(self) -> list[int]:
+        return [m.id for m in self.motors.values()]
+
+    def _model_nb_to_model(self, motor_nb: int) -> str:
+        return self._model_nb_to_model_dict[motor_nb]
+
+    def _id_to_model(self, motor_id: int) -> str:
+        return self._id_to_model_dict[motor_id]
+
+    def _id_to_name(self, motor_id: int) -> str:
+        return self._id_to_name_dict[motor_id]
+
+    def _get_motor_id(self, motor: NameOrID) -> int:
+        if isinstance(motor, str):
+            return self.motors[motor].id
+        elif isinstance(motor, int):
+            return motor
+        else:
+            raise TypeError(f"'{motor}' should be int, str.")
+
+    def _get_motor_model(self, motor: NameOrID) -> int:
+        if isinstance(motor, str):
+            return self.motors[motor].model
+        elif isinstance(motor, int):
+            return self._id_to_model_dict[motor]
+        else:
+            raise TypeError(f"'{motor}' should be int, str.")
+
+    def _get_motors_list(self, motors: str | list[str] | None) -> list[str]:
+        if motors is None:
+            return self.names
+        elif isinstance(motors, str):
+            return [motors]
+        elif isinstance(motors, list):
+            return motors.copy()
+        else:
+            raise TypeError(motors)
+
+    def _get_ids_values_dict(self, values: Value | dict[str, Value] | None) -> list[str]:
+        if isinstance(values, (int, float)):
+            return dict.fromkeys(self.ids, values)
+        elif isinstance(values, dict):
+            return {self.motors[motor].id: val for motor, val in values.items()}
+        else:
+            raise TypeError(f"'values' is expected to be a single value or a dict. Got {values}")
+
+    def _validate_motors(self) -> None:
+        if len(self.ids) != len(set(self.ids)):
+            raise ValueError(f"Some motors have the same id!\n{self}")
+
+        # Ensure ctrl table available for all models
+        for model in self.models:
+            get_ctrl_table(self.model_ctrl_table, model)
+
+    def _is_comm_success(self, comm: int) -> bool:
+        return comm == self._comm_success
+
+    def _is_error(self, error: int) -> bool:
+        return error != self._no_error
+
+    def _assert_motors_exist(self) -> None:
+        # TODO(aliberts): collect all wrong ids/models and display them at once
+        found_models = {}
+        for id_ in self.ids:
+            model_nb = self.ping(id_)
+            if model_nb is not None:
+                found_models[id_] = model_nb
+        expected_models = {m.id: self.model_number_table[m.model] for m in self.motors.values()}
+        if set(found_models) != set(self.ids):
+            raise RuntimeError(
+                f"{self.__class__.__name__} is supposed to have these motors: ({{id: model_nb}})"
+                f"\n{pformat(expected_models, indent=4, sort_dicts=False)}\n"
+                f"But it found these motors on port '{self.port}':"
+                f"\n{pformat(found_models, indent=4, sort_dicts=False)}\n"
+            )
+
+        for id_, model in expected_models.items():
+            if found_models[id_] != model:
+                raise RuntimeError(
+                    f"Motor '{self._id_to_name(id_)}' (id={id_}) is supposed to be of model_number={model} "
+                    f"('{self._id_to_model(id_)}') but a model_number={found_models[id_]} "
+                    "was found instead for that id."
+                )
+
+    @abc.abstractmethod
+    def _assert_protocol_is_compatible(self, instruction_name: str) -> None:
+        pass
+
+    @property
+    def is_connected(self) -> bool:
+        return self.port_handler.is_open
+
+    def connect(self, handshake: bool = True) -> None:
+        if self.is_connected:
+            raise DeviceAlreadyConnectedError(
+                f"{self.__class__.__name__}('{self.port}') is already connected. Do not call `{self.__class__.__name__}.connect()` twice."
+            )
+
+        try:
+            if not self.port_handler.openPort():
+                raise OSError(f"Failed to open port '{self.port}'.")
+            elif handshake:
+                self._handshake()
+        except (FileNotFoundError, OSError, serial.SerialException) as e:
+            raise ConnectionError(
+                f"\nCould not connect on port '{self.port}'. Make sure you are using the correct port."
+                "\nTry running `python lerobot/scripts/find_motors_bus_port.py`\n"
+            ) from e
+
+        self.set_timeout()
+        logger.debug(f"{self.__class__.__name__} connected.")
+
+    @abc.abstractmethod
+    def _handshake(self) -> None:
+        pass
+
+    @classmethod
+    def scan_port(cls, port: str, *args, **kwargs) -> dict[int, list[int]]:
+        bus = cls(port, {}, *args, **kwargs)
+        try:
+            bus.port_handler.openPort()
+        except (FileNotFoundError, OSError, serial.SerialException) as e:
+            raise ConnectionError(
+                f"Could not connect to port '{port}'. Make sure you are using the correct port."
+                "\nTry running `python lerobot/scripts/find_motors_bus_port.py`\n"
+            ) from e
+        baudrate_ids = {}
+        for baudrate in tqdm(bus.available_baudrates, desc="Scanning port"):
+            bus.set_baudrate(baudrate)
+            ids_models = bus.broadcast_ping()
+            if ids_models:
+                tqdm.write(f"Motors found for {baudrate=}: {pformat(ids_models, indent=4)}")
+                baudrate_ids[baudrate] = list(ids_models)
+
+        return baudrate_ids
+
+    @abc.abstractmethod
+    def configure_motors(self) -> None:
+        pass
+
+    @abc.abstractmethod
+    def disable_torque(self, motors: str | list[str] | None = None, num_retry: int = 0) -> None:
+        pass
+
+    @abc.abstractmethod
+    def enable_torque(self, motors: str | list[str] | None = None, num_retry: int = 0) -> None:
+        pass
+
+    @contextmanager
+    def torque_disabled(self):
+        self.disable_torque()
+        try:
+            yield
+        finally:
+            self.enable_torque()
+
+    def set_timeout(self, timeout_ms: int | None = None):
+        timeout_ms = timeout_ms if timeout_ms is not None else self.default_timeout
+        self.port_handler.setPacketTimeoutMillis(timeout_ms)
+
+    def get_baudrate(self) -> int:
+        return self.port_handler.getBaudRate()
+
+    def set_baudrate(self, baudrate: int) -> None:
+        present_bus_baudrate = self.port_handler.getBaudRate()
+        if present_bus_baudrate != baudrate:
+            logger.info(f"Setting bus baud rate to {baudrate}. Previously {present_bus_baudrate}.")
+            self.port_handler.setBaudRate(baudrate)
+
+            if self.port_handler.getBaudRate() != baudrate:
+                raise OSError("Failed to write bus baud rate.")
+
+    @property
+    def is_calibrated(self) -> bool:
+        return self.calibration == self.read_calibration()
+
+    def read_calibration(self) -> dict[str, MotorCalibration]:
+        offsets = self.sync_read("Homing_Offset", normalize=False)
+        mins = self.sync_read("Min_Position_Limit", normalize=False)
+        maxes = self.sync_read("Max_Position_Limit", normalize=False)
+
+        try:
+            drive_modes = self.sync_read("Drive_Mode", normalize=False)
+        except KeyError:
+            drive_modes = dict.fromkeys(self.names, 0)
+
+        calibration = {}
+        for name, motor in self.motors.items():
+            calibration[name] = MotorCalibration(
+                id=motor.id,
+                drive_mode=drive_modes[name],
+                homing_offset=offsets[name],
+                range_min=mins[name],
+                range_max=maxes[name],
+            )
+
+        return calibration
+
+    def write_calibration(self, calibration_dict: dict[str, MotorCalibration]) -> None:
+        for motor, calibration in calibration_dict.items():
+            self.write("Homing_Offset", motor, calibration.homing_offset)
+            self.write("Min_Position_Limit", motor, calibration.range_min)
+            self.write("Max_Position_Limit", motor, calibration.range_max)
+
+        self.calibration = calibration_dict
+
+    def reset_calibration(self, motors: NameOrID | list[NameOrID] | None = None) -> None:
+        if motors is None:
+            motors = self.names
+        elif isinstance(motors, (str, int)):
+            motors = [motors]
+        elif not isinstance(motors, list):
+            raise TypeError(motors)
+
+        for motor in motors:
+            model = self._get_motor_model(motor)
+            max_res = self.model_resolution_table[model] - 1
+            self.write("Homing_Offset", motor, 0, normalize=False)
+            self.write("Min_Position_Limit", motor, 0, normalize=False)
+            self.write("Max_Position_Limit", motor, max_res, normalize=False)
+
+        self.calibration = {}
+
+    def set_half_turn_homings(self, motors: NameOrID | list[NameOrID] | None = None) -> dict[NameOrID, Value]:
+        """
+        This assumes motors present positions are roughly in the middle of their desired range
+
+        Step 1: Set homing and min max to 0
+
+        Step 2: Read Present_Position which will be Actual_Position since
+        Present_Position = Actual_Position ± Homing_Offset (1)
+        and Homing_Offset = 0 from step 1
+
+        Step 3: We want to set the Homing_Offset such that the current Present_Position to be half range of 1
+        revolution. For instance, if 1 revolution corresponds to 4095 (4096 steps), this means we want the
+        current Present_Position to be 2047.
+
+        In that example:
+        Present_Position = 2047 (2)
+        Actual_Position = X (read in step 2)
+        from (1) and (2):
+        => Homing_Offset = ±(X - 2048)
+        """
+        if motors is None:
+            motors = self.names
+        elif isinstance(motors, (str, int)):
+            motors = [motors]
+        else:
+            raise TypeError(motors)
+
+        self.reset_calibration(motors)
+        actual_positions = self.sync_read("Present_Position", motors, normalize=False)
+        homing_offsets = self._get_half_turn_homings(actual_positions)
+        for motor, offset in homing_offsets.items():
+            self.write("Homing_Offset", motor, offset)
+
+        return homing_offsets
+
+    @abc.abstractmethod
+    def _get_half_turn_homings(self, positions: dict[NameOrID, Value]) -> dict[NameOrID, Value]:
+        pass
+
+    def record_ranges_of_motion(
+        self, motors: NameOrID | list[NameOrID] | None = None, display_values: bool = True
+    ) -> tuple[dict[NameOrID, Value], dict[NameOrID, Value]]:
+        """
+        This assumes that the homing offsets have been set such that all possible values in the range of
+        motion are positive and that the zero is not crossed. To that end, `set_half_turn_homings` should
+        typically be called prior to this.
+        """
+        if motors is None:
+            motors = self.names
+        elif isinstance(motors, (str, int)):
+            motors = [motors]
+        elif not isinstance(motors, list):
+            raise TypeError(motors)
+
+        start_positions = self.sync_read("Present_Position", motors, normalize=False)
+        mins = start_positions.copy()
+        maxes = start_positions.copy()
+        while True:
+            positions = self.sync_read("Present_Position", motors, normalize=False)
+            mins = {motor: min(positions[motor], min_) for motor, min_ in mins.items()}
+            maxes = {motor: max(positions[motor], max_) for motor, max_ in maxes.items()}
+
+            if display_values:
+                print("\n-------------------------------------------")
+                print(f"{'NAME':<15} | {'MIN':>6} | {'POS':>6} | {'MAX':>6}")
+                for name in motors:
+                    print(f"{name:<15} | {mins[name]:>6} | {positions[name]:>6} | {maxes[name]:>6}")
+
+            if enter_pressed():
+                break
+
+            if display_values:
+                # Move cursor up to overwrite the previous output
+                move_cursor_up(len(motors) + 3)
+
+        return mins, maxes
+
+    def _normalize(self, data_name: str, ids_values: dict[int, int]) -> dict[int, float]:
+        if not self.calibration:
+            raise RuntimeError(f"{self} has no calibration registered.")
+
+        normalized_values = {}
+        for id_, val in ids_values.items():
+            name = self._id_to_name(id_)
+            min_ = self.calibration[name].range_min
+            max_ = self.calibration[name].range_max
+            bounded_val = min(max_, max(min_, val))
+            if self.motors[name].norm_mode is MotorNormMode.RANGE_M100_100:
+                normalized_values[id_] = (((bounded_val - min_) / (max_ - min_)) * 200) - 100
+            elif self.motors[name].norm_mode is MotorNormMode.RANGE_0_100:
+                normalized_values[id_] = ((bounded_val - min_) / (max_ - min_)) * 100
+            else:
+                # TODO(alibers): velocity and degree modes
+                raise NotImplementedError
+
+        return normalized_values
+
+    def _unnormalize(self, data_name: str, ids_values: dict[int, float]) -> dict[int, int]:
+        if not self.calibration:
+            raise RuntimeError(f"{self} has no calibration registered.")
+
+        unnormalized_values = {}
+        for id_, val in ids_values.items():
+            name = self._id_to_name(id_)
+            min_ = self.calibration[name].range_min
+            max_ = self.calibration[name].range_max
+            if self.motors[name].norm_mode is MotorNormMode.RANGE_M100_100:
+                bounded_val = min(100.0, max(-100.0, val))
+                unnormalized_values[id_] = int(((bounded_val + 100) / 200) * (max_ - min_) + min_)
+            elif self.motors[name].norm_mode is MotorNormMode.RANGE_0_100:
+                bounded_val = min(100.0, max(0.0, val))
+                unnormalized_values[id_] = int((bounded_val / 100) * (max_ - min_) + min_)
+            else:
+                # TODO(alibers): velocity and degree modes
+                raise NotImplementedError
+
+        return unnormalized_values
+
+    @abc.abstractmethod
+    def _encode_sign(self, data_name: str, ids_values: dict[int, int]) -> dict[int, int]:
+        pass
+
+    @abc.abstractmethod
+    def _decode_sign(self, data_name: str, ids_values: dict[int, int]) -> dict[int, int]:
+        pass
+
+    def _serialize_data(self, value: int, length: int) -> list[int]:
+        """
+        Converts an unsigned integer value into a list of byte-sized integers to be sent via a communication
+        protocol. Depending on the protocol, split values can be in big-endian or little-endian order.
+
+        Supported data length for both Feetech and Dynamixel:
+            - 1 (for values 0 to 255)
+            - 2 (for values 0 to 65,535)
+            - 4 (for values 0 to 4,294,967,295)
+        """
+        if value < 0:
+            raise ValueError(f"Negative values are not allowed: {value}")
+
+        max_value = {1: 0xFF, 2: 0xFFFF, 4: 0xFFFFFFFF}.get(length)
+        if max_value is None:
+            raise NotImplementedError(f"Unsupported byte size: {length}. Expected [1, 2, 4].")
+
+        if value > max_value:
+            raise ValueError(f"Value {value} exceeds the maximum for {length} bytes ({max_value}).")
+
+        return self._split_into_byte_chunks(value, length)
+
+    @abc.abstractmethod
+    def _split_into_byte_chunks(self, value: int, length: int) -> list[int]:
+        """Convert an integer into a list of byte-sized integers."""
+        pass
+
+    def ping(self, motor: NameOrID, num_retry: int = 0, raise_on_error: bool = False) -> int | None:
+        id_ = self._get_motor_id(motor)
+        for n_try in range(1 + num_retry):
+            model_number, comm, error = self.packet_handler.ping(self.port_handler, id_)
+            if self._is_comm_success(comm):
+                break
+            logger.debug(f"ping failed for {id_=}: {n_try=} got {comm=} {error=}")
+
+        if not self._is_comm_success(comm):
+            if raise_on_error:
+                raise ConnectionError(self.packet_handler.getTxRxResult(comm))
+            else:
+                return
+        if self._is_error(error):
+            if raise_on_error:
+                raise RuntimeError(self.packet_handler.getRxPacketError(error))
+            else:
+                return
+
+        return model_number
+
+    @abc.abstractmethod
+    def broadcast_ping(self, num_retry: int = 0, raise_on_error: bool = False) -> dict[int, int] | None:
+        pass
+
+    def read(
+        self,
+        data_name: str,
+        motor: str,
+        *,
+        normalize: bool = True,
+        num_retry: int = 0,
+    ) -> Value:
+        if not self.is_connected:
+            raise DeviceNotConnectedError(
+                f"{self.__class__.__name__}('{self.port}') is not connected. You need to run `{self.__class__.__name__}.connect()`."
+            )
+
+        id_ = self.motors[motor].id
+        model = self.motors[motor].model
+        addr, length = get_address(self.model_ctrl_table, model, data_name)
+
+        err_msg = f"Failed to read '{data_name}' on {id_=} after {num_retry + 1} tries."
+        value, _, _ = self._read(addr, length, id_, num_retry=num_retry, raise_on_error=True, err_msg=err_msg)
+
+        id_value = self._decode_sign(data_name, {id_: value})
+
+        if normalize and data_name in self.normalized_data:
+            id_value = self._normalize(data_name, id_value)
+
+        return id_value[id_]
+
+    def _read(
+        self,
+        address: int,
+        length: int,
+        motor_id: int,
+        *,
+        num_retry: int = 0,
+        raise_on_error: bool = True,
+        err_msg: str = "",
+    ) -> tuple[int, int]:
+        if length == 1:
+            read_fn = self.packet_handler.read1ByteTxRx
+        elif length == 2:
+            read_fn = self.packet_handler.read2ByteTxRx
+        elif length == 4:
+            read_fn = self.packet_handler.read4ByteTxRx
+        else:
+            raise ValueError(length)
+
+        for n_try in range(1 + num_retry):
+            value, comm, error = read_fn(self.port_handler, motor_id, address)
+            if self._is_comm_success(comm):
+                break
+            logger.debug(
+                f"Failed to read @{address=} ({length=}) on {motor_id=} ({n_try=}): "
+                + self.packet_handler.getTxRxResult(comm)
+            )
+
+        if not self._is_comm_success(comm) and raise_on_error:
+            raise ConnectionError(f"{err_msg} {self.packet_handler.getTxRxResult(comm)}")
+        elif self._is_error(error) and raise_on_error:
+            raise RuntimeError(f"{err_msg} {self.packet_handler.getRxPacketError(error)}")
+
+        return value, comm, error
+
+    def write(
+        self, data_name: str, motor: str, value: Value, *, normalize: bool = True, num_retry: int = 0
+    ) -> None:
+        if not self.is_connected:
+            raise DeviceNotConnectedError(
+                f"{self.__class__.__name__}('{self.port}') is not connected. You need to run `{self.__class__.__name__}.connect()`."
+            )
+
+        id_ = self.motors[motor].id
+        model = self.motors[motor].model
+        addr, length = get_address(self.model_ctrl_table, model, data_name)
+
+        if normalize and data_name in self.normalized_data:
+            value = self._unnormalize(data_name, {id_: value})[id_]
+
+        value = self._encode_sign(data_name, {id_: value})[id_]
+
+        err_msg = f"Failed to write '{data_name}' on {id_=} with '{value}' after {num_retry + 1} tries."
+        self._write(addr, length, id_, value, num_retry=num_retry, raise_on_error=True, err_msg=err_msg)
+
+    def _write(
+        self,
+        addr: int,
+        length: int,
+        motor_id: int,
+        value: int,
+        *,
+        num_retry: int = 0,
+        raise_on_error: bool = True,
+        err_msg: str = "",
+    ) -> tuple[int, int]:
+        data = self._serialize_data(value, length)
+        for n_try in range(1 + num_retry):
+            comm, error = self.packet_handler.writeTxRx(self.port_handler, motor_id, addr, length, data)
+            if self._is_comm_success(comm):
+                break
+            logger.debug(
+                f"Failed to sync write @{addr=} ({length=}) on id={motor_id} with {value=} ({n_try=}): "
+                + self.packet_handler.getTxRxResult(comm)
+            )
+
+        if not self._is_comm_success(comm) and raise_on_error:
+            raise ConnectionError(f"{err_msg} {self.packet_handler.getTxRxResult(comm)}")
+        elif self._is_error(error) and raise_on_error:
+            raise RuntimeError(f"{err_msg} {self.packet_handler.getRxPacketError(error)}")
+
+        return comm, error
+
+    def sync_read(
+        self,
+        data_name: str,
+        motors: str | list[str] | None = None,
+        *,
+        normalize: bool = True,
+        num_retry: int = 0,
+    ) -> dict[str, Value]:
+        if not self.is_connected:
+            raise DeviceNotConnectedError(
+                f"{self.__class__.__name__}('{self.port}') is not connected. You need to run `{self.__class__.__name__}.connect()`."
+            )
+
+        self._assert_protocol_is_compatible("sync_read")
+
+        names = self._get_motors_list(motors)
+        ids = [self.motors[name].id for name in names]
+        models = [self.motors[name].model for name in names]
+
+        if self._has_different_ctrl_tables:
+            assert_same_address(self.model_ctrl_table, models, data_name)
+
+        model = next(iter(models))
+        addr, length = get_address(self.model_ctrl_table, model, data_name)
+
+        err_msg = f"Failed to sync read '{data_name}' on {ids=} after {num_retry + 1} tries."
+        ids_values, _ = self._sync_read(
+            addr, length, ids, num_retry=num_retry, raise_on_error=True, err_msg=err_msg
+        )
+
+        ids_values = self._decode_sign(data_name, ids_values)
+
+        if normalize and data_name in self.normalized_data:
+            ids_values = self._normalize(data_name, ids_values)
+
+        return {self._id_to_name(id_): value for id_, value in ids_values.items()}
+
+    def _sync_read(
+        self,
+        addr: int,
+        length: int,
+        motor_ids: list[int],
+        *,
+        num_retry: int = 0,
+        raise_on_error: bool = True,
+        err_msg: str = "",
+    ) -> tuple[dict[int, int], int]:
+        self._setup_sync_reader(motor_ids, addr, length)
+        for n_try in range(1 + num_retry):
+            comm = self.sync_reader.txRxPacket()
+            if self._is_comm_success(comm):
+                break
+            logger.debug(
+                f"Failed to sync read @{addr=} ({length=}) on {motor_ids=} ({n_try=}): "
+                + self.packet_handler.getTxRxResult(comm)
+            )
+
+        if not self._is_comm_success(comm) and raise_on_error:
+            raise ConnectionError(f"{err_msg} {self.packet_handler.getTxRxResult(comm)}")
+
+        values = {id_: self.sync_reader.getData(id_, addr, length) for id_ in motor_ids}
+        return values, comm
+
+    def _setup_sync_reader(self, motor_ids: list[int], addr: int, length: int) -> None:
+        self.sync_reader.clearParam()
+        self.sync_reader.start_address = addr
+        self.sync_reader.data_length = length
+        for id_ in motor_ids:
+            self.sync_reader.addParam(id_)
+
+    # TODO(aliberts, pkooij): Implementing something like this could get even much faster read times if need be.
+    # Would have to handle the logic of checking if a packet has been sent previously though but doable.
+    # This could be at the cost of increase latency between the moment the data is produced by the motors and
+    # the moment it is used by a policy.
+    # def _async_read(self, motor_ids: list[int], address: int, length: int):
+    #     if self.sync_reader.start_address != address or self.sync_reader.data_length != length or ...:
+    #         self._setup_sync_reader(motor_ids, address, length)
+    #     else:
+    #         self.sync_reader.rxPacket()
+    #         self.sync_reader.txPacket()
+
+    #     for id_ in motor_ids:
+    #         value = self.sync_reader.getData(id_, address, length)
+
+    def sync_write(
+        self,
+        data_name: str,
+        values: Value | dict[str, Value],
+        *,
+        normalize: bool = True,
+        num_retry: int = 0,
+    ) -> None:
+        if not self.is_connected:
+            raise DeviceNotConnectedError(
+                f"{self.__class__.__name__}('{self.port}') is not connected. You need to run `{self.__class__.__name__}.connect()`."
+            )
+
+        ids_values = self._get_ids_values_dict(values)
+        models = [self._id_to_model(id_) for id_ in ids_values]
+        if self._has_different_ctrl_tables:
+            assert_same_address(self.model_ctrl_table, models, data_name)
+
+        model = next(iter(models))
+        addr, length = get_address(self.model_ctrl_table, model, data_name)
+
+        if normalize and data_name in self.normalized_data:
+            ids_values = self._unnormalize(data_name, ids_values)
+
+        ids_values = self._encode_sign(data_name, ids_values)
+
+        err_msg = f"Failed to sync write '{data_name}' with {ids_values=} after {num_retry + 1} tries."
+        self._sync_write(addr, length, ids_values, num_retry=num_retry, raise_on_error=True, err_msg=err_msg)
+
+    def _sync_write(
+        self,
+        addr: int,
+        length: int,
+        ids_values: dict[int, int],
+        num_retry: int = 0,
+        raise_on_error: bool = True,
+        err_msg: str = "",
+    ) -> int:
+        self._setup_sync_writer(ids_values, addr, length)
+        for n_try in range(1 + num_retry):
+            comm = self.sync_writer.txPacket()
+            if self._is_comm_success(comm):
+                break
+            logger.debug(
+                f"Failed to sync write @{addr=} ({length=}) with {ids_values=} ({n_try=}): "
+                + self.packet_handler.getTxRxResult(comm)
+            )
+
+        if not self._is_comm_success(comm) and raise_on_error:
+            raise ConnectionError(f"{err_msg} {self.packet_handler.getTxRxResult(comm)}")
+
+        return comm
+
+    def _setup_sync_writer(self, ids_values: dict[int, int], addr: int, length: int) -> None:
+        self.sync_writer.clearParam()
+        self.sync_writer.start_address = addr
+        self.sync_writer.data_length = length
+        for id_, value in ids_values.items():
+            data = self._serialize_data(value, length)
+            self.sync_writer.addParam(id_, data)
+
+    def disconnect(self, disable_torque: bool = True) -> None:
+        if not self.is_connected:
+            raise DeviceNotConnectedError(
+                f"{self.__class__.__name__}('{self.port}') is not connected. Try running `{self.__class__.__name__}.connect()` first."
+            )
+
+        if disable_torque:
+            self.port_handler.clearPort()
+            self.port_handler.is_using = False
+            self.disable_torque(num_retry=5)
+
+        self.port_handler.closePort()
+        logger.debug(f"{self.__class__.__name__} disconnected.")

lerobot/common/motors/utils.py

@@ -0,0 +1,56 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .configs import MotorsBusConfig
+from .motors_bus import MotorsBus
+
+
+def make_motors_buses_from_configs(motors_bus_configs: dict[str, MotorsBusConfig]) -> list[MotorsBus]:
+    motors_buses = {}
+
+    for key, cfg in motors_bus_configs.items():
+        if cfg.type == "dynamixel":
+            from .dynamixel import DynamixelMotorsBus
+
+            motors_buses[key] = DynamixelMotorsBus(cfg)
+
+        elif cfg.type == "feetech":
+            from lerobot.common.motors.feetech.feetech import FeetechMotorsBus
+
+            motors_buses[key] = FeetechMotorsBus(cfg)
+
+        else:
+            raise ValueError(f"The motor type '{cfg.type}' is not valid.")
+
+    return motors_buses
+
+
+def make_motors_bus(motor_type: str, **kwargs) -> MotorsBus:
+    if motor_type == "dynamixel":
+        from .configs import DynamixelMotorsBusConfig
+        from .dynamixel import DynamixelMotorsBus
+
+        config = DynamixelMotorsBusConfig(**kwargs)
+        return DynamixelMotorsBus(config)
+
+    elif motor_type == "feetech":
+        from feetech import FeetechMotorsBus
+
+        from .configs import FeetechMotorsBusConfig
+
+        config = FeetechMotorsBusConfig(**kwargs)
+        return FeetechMotorsBus(config)
+
+    else:
+        raise ValueError(f"The motor type '{motor_type}' is not valid.")

lerobot/common/optim/__init__.py

@@ -1 +1,15 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from .optimizers import OptimizerConfig as OptimizerConfig

lerobot/common/policies/__init__.py

@@ -1,3 +1,17 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from .act.configuration_act import ACTConfig as ACTConfig
 from .diffusion.configuration_diffusion import DiffusionConfig as DiffusionConfig
 from .pi0.configuration_pi0 import PI0Config as PI0Config