User/pepijn/2025 03 17 act different image shapes (#870)

Co-authored-by: Simon Alibert <75076266+aliberts@users.noreply.github.com>

.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.
+
 # 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
+    rev: v1.30.2
     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.9.10
     hooks:
       - id: ruff
         args: [--fix]
       - id: ruff-format
+
+
+  ##### Security #####
   - repo: https://github.com/gitleaks/gitleaks
-    rev: v8.23.3
+    rev: v8.24.0
     hooks:
       - id: gitleaks
+
   - repo: https://github.com/woodruffw/zizmor-pre-commit
-    rev: v1.3.1
+    rev: v1.4.1
     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/>
@@ -223,8 +232,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 +384,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/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/10_use_so100.md

@@ -4,8 +4,8 @@
 
   - [A. Source the parts](#a-source-the-parts)
   - [B. Install LeRobot](#b-install-lerobot)
-  - [C. Configure the motors](#c-configure-the-motors)
-  - [D. Assemble the arms](#d-assemble-the-arms)
+  - [C. Configure the Motors](#c-configure-the-motors)
+  - [D. Step-by-Step Assembly Instructions](#d-step-by-step-assembly-instructions)
   - [E. Calibrate](#e-calibrate)
   - [F. Teleoperate](#f-teleoperate)
   - [G. Record a dataset](#g-record-a-dataset)
@@ -70,6 +70,7 @@ conda install -y -c conda-forge "opencv>=4.10.0"
 ```
 Great :hugs:! You are now done installing LeRobot and we can begin assembling the SO100 arms :robot:.
 Every time you now want to use LeRobot you can go to the `~/lerobot` folder where we installed LeRobot and run one of the commands.
+
 ## C. Configure the motors
 
 > [!NOTE]
@@ -98,22 +99,22 @@ Example output when identifying the leader arm's port (e.g., `/dev/tty.usbmodem5
 ```
 Finding all available ports for the MotorBus.
 ['/dev/tty.usbmodem575E0032081', '/dev/tty.usbmodem575E0031751']
-Remove the usb cable from your DynamixelMotorsBus and press Enter when done.
+Remove the usb cable from your MotorsBus and press Enter when done.
 
 [...Disconnect leader arm and press Enter...]
 
-The port of this DynamixelMotorsBus is /dev/tty.usbmodem575E0031751
+The port of this MotorsBus is /dev/tty.usbmodem575E0031751
 Reconnect the usb cable.
 ```
 Example output when identifying the follower arm's port (e.g., `/dev/tty.usbmodem575E0032081`, or possibly `/dev/ttyACM1` on Linux):
 ```
 Finding all available ports for the MotorBus.
 ['/dev/tty.usbmodem575E0032081', '/dev/tty.usbmodem575E0031751']
-Remove the usb cable from your DynamixelMotorsBus and press Enter when done.
+Remove the usb cable from your MotorsBus and press Enter when done.
 
 [...Disconnect follower arm and press Enter...]
 
-The port of this DynamixelMotorsBus is /dev/tty.usbmodem575E0032081
+The port of this MotorsBus is /dev/tty.usbmodem575E0032081
 Reconnect the usb cable.
 ```
 
@@ -221,19 +222,13 @@ Redo the process for all your motors until ID 6. Do the same for the 6 motors of
 
 Follow the video for removing gears. You need to remove the gear for the motors of the leader arm. As a result, you will only use the position encoding of the motor and reduce friction to more easily operate the leader arm.
 
-#### c. Add motor horn to all 12 motors
+## D. Step-by-Step Assembly Instructions
 
-<details>
-<summary><strong>Video adding motor horn</strong></summary>
+**Step 1: Clean Parts**
+- Remove all support material from the 3D-printed parts.
+---
 
-<video src="https://github.com/user-attachments/assets/ef3391a4-ad05-4100-b2bd-1699bf86c969"></video>
-
-</details>
-
-Follow the video for adding the motor horn. For SO-100, you need to align the holes on the motor horn to the motor spline to be approximately 1:30, 4:30, 7:30 and 10:30.
-Try to avoid rotating the motor while doing so to keep position 2048 set during configuration. It is especially tricky for the leader motors as it is more sensible without the gears, but it's ok if it's a bit rotated.
-
-## D. Assemble the arms
+### Additional Guidance
 
 <details>
 <summary><strong>Video assembling arms</strong></summary>
@@ -242,7 +237,211 @@ Try to avoid rotating the motor while doing so to keep position 2048 set during
 
 </details>
 
-Follow the video for assembling the arms. It is important to insert the cables into the motor that is being assembled before you assemble the motor into the arm! Inserting the cables beforehand is much easier than doing this afterward. The first arm should take a bit more than 1 hour to assemble, but once you get used to it, you can do it under 1 hour for the second arm.
+**Note:**
+This video provides visual guidance for assembling the arms, but it doesn't specify when or how to do the wiring. Inserting the cables beforehand is much easier than doing it afterward. The first arm may take a bit more than 1 hour to assemble, but once you get used to it, you can assemble the second arm in under 1 hour.
+
+---
+
+### First Motor
+
+**Step 2: Insert Wires**
+- Insert two wires into the first motor.
+
+  <img src="../media/tutorial/img1.jpg" style="height:300px;">
+
+**Step 3: Install in Base**
+- Place the first motor into the base.
+
+  <img src="../media/tutorial/img2.jpg" style="height:300px;">
+
+**Step 4: Secure Motor**
+- Fasten the motor with 4 screws. Two from the bottom and two from top.
+
+**Step 5: Attach Motor Holder**
+- Slide over the first motor holder and fasten it using two screws (one on each side).
+
+  <img src="../media/tutorial/img4.jpg" style="height:300px;">
+
+**Step 6: Attach Motor Horns**
+- Install both motor horns, securing the top horn with a screw. Try not to move the motor position when attaching the motor horn, especially for the leader arms, where we removed the gears.
+
+  <img src="../media/tutorial/img5.jpg" style="height:300px;">
+<details>
+  <summary><strong>Video adding motor horn</strong></summary>
+  <video src="https://github.com/user-attachments/assets/ef3391a4-ad05-4100-b2bd-1699bf86c969"></video>
+</details>
+
+**Step 7: Attach Shoulder Part**
+- Route one wire to the back of the robot and the other to the left or in photo towards you (see photo).
+- Attach the shoulder part.
+
+  <img src="../media/tutorial/img6.jpg" style="height:300px;">
+
+**Step 8: Secure Shoulder**
+- Tighten the shoulder part with 4 screws on top and 4 on the bottom
+*(access bottom holes by turning the shoulder).*
+
+---
+
+### Second Motor Assembly
+
+**Step 9: Install Motor 2**
+- Slide the second motor in from the top and link the wire from motor 1 to motor 2.
+
+  <img src="../media/tutorial/img8.jpg" style="height:300px;">
+
+**Step 10: Attach Shoulder Holder**
+- Add the shoulder motor holder.
+- Ensure the wire from motor 1 to motor 2 goes behind the holder while the other wire is routed upward (see photo).
+- This part can be tight to assemble, you can use a workbench like the image or a similar setup to push the part around the motor.
+
+  <div style="display: flex;">
+    <img src="../media/tutorial/img9.jpg" style="height:250px;">
+    <img src="../media/tutorial/img10.jpg" style="height:250px;">
+    <img src="../media/tutorial/img12.jpg" style="height:250px;">
+  </div>
+
+**Step 11: Secure Motor 2**
+- Fasten the second motor with 4 screws.
+
+**Step 12: Attach Motor Horn**
+- Attach both motor horns to motor 2, again use the horn screw.
+
+**Step 13: Attach Base**
+- Install the base attachment using 2 screws.
+
+  <img src="../media/tutorial/img11.jpg" style="height:300px;">
+
+**Step 14: Attach Upper Arm**
+- Attach the upper arm with 4 screws on each side.
+
+  <img src="../media/tutorial/img13.jpg" style="height:300px;">
+
+---
+
+### Third Motor Assembly
+
+**Step 15: Install Motor 3**
+- Route the motor cable from motor 2 through the cable holder to motor 3, then secure motor 3 with 4 screws.
+
+**Step 16: Attach Motor Horn**
+- Attach both motor horns to motor 3 and secure one again with a horn screw.
+
+  <img src="../media/tutorial/img14.jpg" style="height:300px;">
+
+**Step 17: Attach Forearm**
+- Connect the forearm to motor 3 using 4 screws on each side.
+
+  <img src="../media/tutorial/img15.jpg" style="height:300px;">
+
+---
+
+### Fourth Motor Assembly
+
+**Step 18: Install Motor 4**
+- Slide in motor 4, attach the cable from motor 3, and secure the cable in its holder with a screw.
+
+  <div style="display: flex;">
+    <img src="../media/tutorial/img16.jpg" style="height:300px;">
+    <img src="../media/tutorial/img19.jpg" style="height:300px;">
+  </div>
+
+**Step 19: Attach Motor Holder 4**
+- Install the fourth motor holder (a tight fit). Ensure one wire is routed upward and the wire from motor 3 is routed downward (see photo).
+
+  <img src="../media/tutorial/img17.jpg" style="height:300px;">
+
+**Step 20: Secure Motor 4 & Attach Horn**
+- Fasten motor 4 with 4 screws and attach its motor horns, use for one a horn screw.
+
+  <img src="../media/tutorial/img18.jpg" style="height:300px;">
+
+---
+
+### Wrist Assembly
+
+**Step 21: Install Motor 5**
+- Insert motor 5 into the wrist holder and secure it with 2 front screws.
+
+  <img src="../media/tutorial/img20.jpg" style="height:300px;">
+
+**Step 22: Attach Wrist**
+- Connect the wire from motor 4 to motor 5. And already insert the other wire for the gripper.
+- Secure the wrist to motor 4 using 4 screws on both sides.
+
+  <img src="../media/tutorial/img22.jpg" style="height:300px;">
+
+**Step 23: Attach Wrist Horn**
+- Install only one motor horn on the wrist motor and secure it with a horn screw.
+
+  <img src="../media/tutorial/img23.jpg" style="height:300px;">
+
+---
+
+### Follower Configuration
+
+**Step 24: Attach Gripper**
+- Attach the gripper to motor 5.
+
+  <img src="../media/tutorial/img24.jpg" style="height:300px;">
+
+**Step 25: Install Gripper Motor**
+- Insert the gripper motor, connect the motor wire from motor 5 to motor 6, and secure it with 3 screws on each side.
+
+  <img src="../media/tutorial/img25.jpg" style="height:300px;">
+
+**Step 26: Attach Gripper Horn & Claw**
+- Attach the motor horns and again use a horn screw.
+- Install the gripper claw and secure it with 4 screws on both sides.
+
+  <img src="../media/tutorial/img26.jpg" style="height:300px;">
+
+**Step 27: Mount Controller**
+- Attach the motor controller on the back.
+
+  <div style="display: flex;">
+    <img src="../media/tutorial/img27.jpg" style="height:300px;">
+    <img src="../media/tutorial/img28.jpg" style="height:300px;">
+  </div>
+
+*Assembly complete – proceed to Leader arm assembly.*
+
+---
+
+### Leader Configuration
+
+For the leader configuration, perform **Steps 1–23**. Make sure that you removed the motor gears from the motors.
+
+**Step 24: Attach Leader Holder**
+- Mount the leader holder onto the wrist and secure it with a screw.
+
+  <img src="../media/tutorial/img29.jpg" style="height:300px;">
+
+**Step 25: Attach Handle**
+- Attach the handle to motor 5 using 4 screws.
+
+  <img src="../media/tutorial/img30.jpg" style="height:300px;">
+
+**Step 26: Install Gripper Motor**
+- Insert the gripper motor, secure it with 3 screws on each side, attach a motor horn using a horn screw, and connect the motor wire.
+
+  <img src="../media/tutorial/img31.jpg" style="height:300px;">
+
+**Step 27: Attach Trigger**
+- Attach the follower trigger with 4 screws.
+
+  <img src="../media/tutorial/img32.jpg" style="height:300px;">
+
+**Step 28: Mount Controller**
+- Attach the motor controller on the back.
+
+  <div style="display: flex;">
+    <img src="../media/tutorial/img27.jpg" style="height:300px;">
+    <img src="../media/tutorial/img28.jpg" style="height:300px;">
+  </div>
+
+*Assembly complete – proceed to calibration.*
+
 
 ## E. Calibrate
 
@@ -255,8 +454,8 @@ Next, you'll need to calibrate your SO-100 robot to ensure that the leader and f
 
 You will need to move the follower arm to these positions sequentially:
 
-| 1. Zero position | 2. Rotated position | 3. Rest position |
-|---|---|---|
+| 1. Zero position                                                                                                                                             | 2. Rotated position                                                                                                                                                   | 3. Rest position                                                                                                                                             |
+| ------------------------------------------------------------------------------------------------------------------------------------------------------------ | --------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------ |
 | <img src="../media/so100/follower_zero.webp?raw=true" alt="SO-100 follower arm zero position" title="SO-100 follower arm zero position" style="width:100%;"> | <img src="../media/so100/follower_rotated.webp?raw=true" alt="SO-100 follower arm rotated position" title="SO-100 follower arm rotated position" style="width:100%;"> | <img src="../media/so100/follower_rest.webp?raw=true" alt="SO-100 follower arm rest position" title="SO-100 follower arm rest position" style="width:100%;"> |
 
 Make sure both arms are connected and run this script to launch manual calibration:
@@ -271,8 +470,8 @@ python lerobot/scripts/control_robot.py \
 #### b. Manual calibration of leader arm
 Follow step 6 of the [assembly video](https://youtu.be/FioA2oeFZ5I?t=724) which illustrates the manual calibration. You will need to move the leader arm to these positions sequentially:
 
-| 1. Zero position | 2. Rotated position | 3. Rest position |
-|---|---|---|
+| 1. Zero position                                                                                                                                       | 2. Rotated position                                                                                                                                             | 3. Rest position                                                                                                                                       |
+| ------------------------------------------------------------------------------------------------------------------------------------------------------ | --------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------ |
 | <img src="../media/so100/leader_zero.webp?raw=true" alt="SO-100 leader arm zero position" title="SO-100 leader arm zero position" style="width:100%;"> | <img src="../media/so100/leader_rotated.webp?raw=true" alt="SO-100 leader arm rotated position" title="SO-100 leader arm rotated position" style="width:100%;"> | <img src="../media/so100/leader_rest.webp?raw=true" alt="SO-100 leader arm rest position" title="SO-100 leader arm rest position" style="width:100%;"> |
 
 Run this script to launch manual calibration:
@@ -372,18 +571,25 @@ python lerobot/scripts/train.py \
   --policy.type=act \
   --output_dir=outputs/train/act_so100_test \
   --job_name=act_so100_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}/so100_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`.
 
 Training should take several hours. You will find checkpoints in `outputs/train/act_so100_test/checkpoints`.
 
+To resume training from a checkpoint, below is an example command to resume from `last` checkpoint of the `act_so100_test` policy:
+```bash
+python lerobot/scripts/train.py \
+  --config_path=outputs/train/act_so100_test/checkpoints/last/pretrained_model/train_config.json \
+  --resume=true
+```
+
 ## K. Evaluate your policy
 
 You can use the `record` function from [`lerobot/scripts/control_robot.py`](../lerobot/scripts/control_robot.py) but with a policy checkpoint as input. For instance, run this command to record 10 evaluation episodes:

examples/11_use_lekiwi.md

@@ -23,6 +23,9 @@ Follow this [README](https://github.com/SIGRobotics-UIUC/LeKiwi). It contains th
 
 Before assembling, you will first need to configure your motors. To this end, we provide a nice script, so let's first install LeRobot. After configuration, we will also guide you through assembly.
 
+### Wired version
+If you have the **wired** LeKiwi version you can skip the installation of the Raspberry Pi and setting up SSH. You can also run all commands directly on your PC for both the LeKiwi scripts and the leader arm scripts for teleoperating.
+
 ## B. Install software on Pi
 Now we have to setup the remote PC that will run on the LeKiwi Robot. This is normally a Raspberry Pi, but can be any PC that can run on 5V and has enough usb ports (2 or more) for the cameras and motor control board.
 
@@ -246,6 +249,110 @@ class LeKiwiRobotConfig(RobotConfig):
         }
     )
 
+    teleop_keys: dict[str, str] = field(
+        default_factory=lambda: {
+            # Movement
+            "forward": "w",
+            "backward": "s",
+            "left": "a",
+            "right": "d",
+            "rotate_left": "z",
+            "rotate_right": "x",
+            # Speed control
+            "speed_up": "r",
+            "speed_down": "f",
+            # quit teleop
+            "quit": "q",
+        }
+    )
+
+    mock: bool = False
+```
+
+## Wired version
+
+For the wired LeKiwi version your configured IP address should refer to your own laptop (127.0.0.1), because leader arm and LeKiwi are in this case connected to own laptop. Below and example configuration for this wired setup:
+```python
+@RobotConfig.register_subclass("lekiwi")
+@dataclass
+class LeKiwiRobotConfig(RobotConfig):
+    # `max_relative_target` limits the magnitude of the relative positional target vector for safety purposes.
+    # Set this to a positive scalar to have the same value for all motors, or a list that is the same length as
+    # the number of motors in your follower arms.
+    max_relative_target: int | None = None
+
+    # Network Configuration
+    ip: str = "127.0.0.1"
+    port: int = 5555
+    video_port: int = 5556
+
+    cameras: dict[str, CameraConfig] = field(
+        default_factory=lambda: {
+            "front": OpenCVCameraConfig(
+                camera_index=0, fps=30, width=640, height=480, rotation=90
+            ),
+            "wrist": OpenCVCameraConfig(
+                camera_index=1, fps=30, width=640, height=480, rotation=180
+            ),
+        }
+    )
+
+    calibration_dir: str = ".cache/calibration/lekiwi"
+
+    leader_arms: dict[str, MotorsBusConfig] = field(
+        default_factory=lambda: {
+            "main": FeetechMotorsBusConfig(
+                port="/dev/tty.usbmodem585A0077581",
+                motors={
+                    # name: (index, model)
+                    "shoulder_pan": [1, "sts3215"],
+                    "shoulder_lift": [2, "sts3215"],
+                    "elbow_flex": [3, "sts3215"],
+                    "wrist_flex": [4, "sts3215"],
+                    "wrist_roll": [5, "sts3215"],
+                    "gripper": [6, "sts3215"],
+                },
+            ),
+        }
+    )
+
+    follower_arms: dict[str, MotorsBusConfig] = field(
+        default_factory=lambda: {
+            "main": FeetechMotorsBusConfig(
+                port="/dev/tty.usbmodem58760431061",
+                motors={
+                    # name: (index, model)
+                    "shoulder_pan": [1, "sts3215"],
+                    "shoulder_lift": [2, "sts3215"],
+                    "elbow_flex": [3, "sts3215"],
+                    "wrist_flex": [4, "sts3215"],
+                    "wrist_roll": [5, "sts3215"],
+                    "gripper": [6, "sts3215"],
+                    "left_wheel": (7, "sts3215"),
+                    "back_wheel": (8, "sts3215"),
+                    "right_wheel": (9, "sts3215"),
+                },
+            ),
+        }
+    )
+
+    teleop_keys: dict[str, str] = field(
+        default_factory=lambda: {
+            # Movement
+            "forward": "w",
+            "backward": "s",
+            "left": "a",
+            "right": "d",
+            "rotate_left": "z",
+            "rotate_right": "x",
+            # Speed control
+            "speed_up": "r",
+            "speed_down": "f",
+            # quit teleop
+            "quit": "q",
+        }
+    )
+
     mock: bool = False
 ```
 
@@ -259,8 +366,8 @@ Now we have to calibrate the leader arm and the follower arm. The wheel motors d
 
 You will need to move the follower arm to these positions sequentially:
 
-| 1. Zero position | 2. Rotated position | 3. Rest position |
-|---|---|---|
+| 1. Zero position                                                                                                                                                  | 2. Rotated position                                                                                                                                                        | 3. Rest position                                                                                                                                                  |
+| ----------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------- |
 | <img src="../media/lekiwi/mobile_calib_zero.webp?raw=true" alt="SO-100 follower arm zero position" title="SO-100 follower arm zero position" style="width:100%;"> | <img src="../media/lekiwi/mobile_calib_rotated.webp?raw=true" alt="SO-100 follower arm rotated position" title="SO-100 follower arm rotated position" style="width:100%;"> | <img src="../media/lekiwi/mobile_calib_rest.webp?raw=true" alt="SO-100 follower arm rest position" title="SO-100 follower arm rest position" style="width:100%;"> |
 
 Make sure the arm is connected to the Raspberry Pi and run this script (on the Raspberry Pi) to launch manual calibration:
@@ -272,11 +379,14 @@ python lerobot/scripts/control_robot.py \
   --control.arms='["main_follower"]'
 ```
 
+### Wired version
+If you have the **wired** LeKiwi version please run all commands including this calibration command on your laptop.
+
 ### Calibrate leader arm
 Then to calibrate the leader arm (which is attached to the laptop/pc). You will need to move the leader arm to these positions sequentially:
 
-| 1. Zero position | 2. Rotated position | 3. Rest position |
-|---|---|---|
+| 1. Zero position                                                                                                                                       | 2. Rotated position                                                                                                                                             | 3. Rest position                                                                                                                                       |
+| ------------------------------------------------------------------------------------------------------------------------------------------------------ | --------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------ |
 | <img src="../media/so100/leader_zero.webp?raw=true" alt="SO-100 leader arm zero position" title="SO-100 leader arm zero position" style="width:100%;"> | <img src="../media/so100/leader_rotated.webp?raw=true" alt="SO-100 leader arm rotated position" title="SO-100 leader arm rotated position" style="width:100%;"> | <img src="../media/so100/leader_rest.webp?raw=true" alt="SO-100 leader arm rest position" title="SO-100 leader arm rest position" style="width:100%;"> |
 
 Run this script (on your laptop/pc) to launch manual calibration:
@@ -306,26 +416,29 @@ python lerobot/scripts/control_robot.py \
 
 You should see on your laptop something like this: ```[INFO] Connected to remote robot at tcp://172.17.133.91:5555 and video stream at tcp://172.17.133.91:5556.``` Now you can move the leader arm and use the keyboard (w,a,s,d) to drive forward, left, backwards, right. And use (z,x) to turn left or turn right. You can use (r,f) to increase and decrease the speed of the mobile robot. There are three speed modes, see the table below:
 | Speed Mode | Linear Speed (m/s) | Rotation Speed (deg/s) |
-|------------|-------------------|-----------------------|
-| Fast      | 0.4               | 90                    |
-| Medium    | 0.25              | 60                    |
-| Slow      | 0.1               | 30                    |
+| ---------- | ------------------ | ---------------------- |
+| Fast       | 0.4                | 90                     |
+| Medium     | 0.25               | 60                     |
+| Slow       | 0.1                | 30                     |
 
 
-| Key  | Action                         |
-|------|--------------------------------|
-| W    | Move forward                   |
-| A    | Move left                       |
-| S    | Move backward                   |
-| D    | Move right                      |
-| Z    | Turn left                       |
-| X    | Turn right                      |
-| R    | Increase speed                  |
-| F    | Decrease speed                  |
+| Key | Action         |
+| --- | -------------- |
+| W   | Move forward   |
+| A   | Move left      |
+| S   | Move backward  |
+| D   | Move right     |
+| Z   | Turn left      |
+| X   | Turn right     |
+| R   | Increase speed |
+| F   | Decrease speed |
 
 > [!TIP]
 >  If you use a different keyboard you can change the keys for each command in the [`LeKiwiRobotConfig`](../lerobot/common/robot_devices/robots/configs.py).
 
+### Wired version
+If you have the **wired** LeKiwi version please run all commands including both these teleoperation commands on your laptop.
+
 ## Troubleshoot communication
 
 If you are having trouble connecting to the Mobile SO100, follow these steps to diagnose and resolve the issue.
@@ -364,6 +477,13 @@ Make sure the configuration file on both your laptop/pc and the Raspberry Pi is
 # G. Record a dataset
 Once you're familiar with teleoperation, you can record your first dataset with LeKiwi.
 
+To start the program on LeKiwi, SSH into your Raspberry Pi, and run `conda activate lerobot` and this script:
+```bash
+python lerobot/scripts/control_robot.py \
+  --robot.type=lekiwi \
+  --control.type=remote_robot
+```
+
 If you want to use the Hugging Face hub features for uploading your dataset and you haven't previously done it, make sure you've logged in using a write-access token, which can be generated from the [Hugging Face settings](https://huggingface.co/settings/tokens):
 ```bash
 huggingface-cli login --token ${HUGGINGFACE_TOKEN} --add-to-git-credential
@@ -374,8 +494,7 @@ Store your Hugging Face repository name in a variable to run these commands:
 HF_USER=$(huggingface-cli whoami | head -n 1)
 echo $HF_USER
 ```
-
-Record 2 episodes and upload your dataset to the hub:
+On your laptop then run this command to record 2 episodes and upload your dataset to the hub:
 ```bash
 python lerobot/scripts/control_robot.py \
   --robot.type=lekiwi \
@@ -393,6 +512,9 @@ python lerobot/scripts/control_robot.py \
 
 Note: You can resume recording by adding `--control.resume=true`.
 
+### Wired version
+If you have the **wired** LeKiwi version please run all commands including both these record dataset commands on your laptop.
+
 # H. Visualize a dataset
 
 If you uploaded your dataset to the hub with `--control.push_to_hub=true`, you can [visualize your dataset online](https://huggingface.co/spaces/lerobot/visualize_dataset) by copy pasting your repo id given by:
@@ -427,14 +549,14 @@ python lerobot/scripts/train.py \
   --policy.type=act \
   --output_dir=outputs/train/act_lekiwi_test \
   --job_name=act_lekiwi_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}/lekiwi_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`.
 
 Training should take several hours. You will find checkpoints in `outputs/train/act_lekiwi_test/checkpoints`.

examples/11_use_moss.md

@@ -176,8 +176,8 @@ Next, you'll need to calibrate your Moss v1 robot to ensure that the leader and
 
 You will need to move the follower arm to these positions sequentially:
 
-| 1. Zero position | 2. Rotated position | 3. Rest position |
-|---|---|---|
+| 1. Zero position                                                                                                                                              | 2. Rotated position                                                                                                                                                    | 3. Rest position                                                                                                                                              |
+| ------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------- |
 | <img src="../media/moss/follower_zero.webp?raw=true" alt="Moss v1 follower arm zero position" title="Moss v1 follower arm zero position" style="width:100%;"> | <img src="../media/moss/follower_rotated.webp?raw=true" alt="Moss v1 follower arm rotated position" title="Moss v1 follower arm rotated position" style="width:100%;"> | <img src="../media/moss/follower_rest.webp?raw=true" alt="Moss v1 follower arm rest position" title="Moss v1 follower arm rest position" style="width:100%;"> |
 
 Make sure both arms are connected and run this script to launch manual calibration:
@@ -192,8 +192,8 @@ python lerobot/scripts/control_robot.py \
 **Manual calibration of leader arm**
 Follow step 6 of the [assembly video](https://www.youtube.com/watch?v=DA91NJOtMic) which illustrates the manual calibration. You will need to move the leader arm to these positions sequentially:
 
-| 1. Zero position | 2. Rotated position | 3. Rest position |
-|---|---|---|
+| 1. Zero position                                                                                                                                        | 2. Rotated position                                                                                                                                              | 3. Rest position                                                                                                                                        |
+| ------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------- |
 | <img src="../media/moss/leader_zero.webp?raw=true" alt="Moss v1 leader arm zero position" title="Moss v1 leader arm zero position" style="width:100%;"> | <img src="../media/moss/leader_rotated.webp?raw=true" alt="Moss v1 leader arm rotated position" title="Moss v1 leader arm rotated position" style="width:100%;"> | <img src="../media/moss/leader_rest.webp?raw=true" alt="Moss v1 leader arm rest position" title="Moss v1 leader arm rest position" style="width:100%;"> |
 
 Run this script to launch manual calibration:
@@ -293,14 +293,14 @@ python lerobot/scripts/train.py \
   --policy.type=act \
   --output_dir=outputs/train/act_moss_test \
   --job_name=act_moss_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}/moss_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`.
 
 Training should take several hours. You will find checkpoints in `outputs/train/act_moss_test/checkpoints`.

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.

examples/2_evaluate_pretrained_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 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.
@@ -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

@@ -386,14 +386,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.
@@ -898,14 +898,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/9_use_aloha.md

@@ -135,14 +135,14 @@ python lerobot/scripts/train.py \
   --policy.type=act \
   --output_dir=outputs/train/act_aloha_test \
   --job_name=act_aloha_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}/aloha_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,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 shutil
 from pathlib import Path
 

lerobot/common/constants.py

@@ -1,3 +1,16 @@
+# 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

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,8 +67,9 @@ 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
@@ -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
 
 

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
 
@@ -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/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,

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,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 abc
 from dataclasses import dataclass, field
 

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

lerobot/common/policies/act/modeling_act.py

@@ -119,9 +119,7 @@ class ACTPolicy(PreTrainedPolicy):
         batch = self.normalize_inputs(batch)
         if self.config.image_features:
             batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
-            batch["observation.images"] = torch.stack(
-                [batch[key] for key in self.config.image_features], dim=-4
-            )
+            batch["observation.images"] = [batch[key] for key in self.config.image_features]
 
         # If we are doing temporal ensembling, do online updates where we keep track of the number of actions
         # we are ensembling over.
@@ -149,9 +147,8 @@ class ACTPolicy(PreTrainedPolicy):
         batch = self.normalize_inputs(batch)
         if self.config.image_features:
             batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
-            batch["observation.images"] = torch.stack(
-                [batch[key] for key in self.config.image_features], dim=-4
-            )
+            batch["observation.images"] = [batch[key] for key in self.config.image_features]
+
         batch = self.normalize_targets(batch)
         actions_hat, (mu_hat, log_sigma_x2_hat) = self.model(batch)
 
@@ -413,11 +410,10 @@ class ACT(nn.Module):
                 "actions must be provided when using the variational objective in training mode."
             )
 
-        batch_size = (
-            batch["observation.images"]
-            if "observation.images" in batch
-            else batch["observation.environment_state"]
-        ).shape[0]
+        if "observation.images" in batch:
+            batch_size = batch["observation.images"][0].shape[0]
+        else:
+            batch_size = batch["observation.environment_state"].shape[0]
 
         # Prepare the latent for input to the transformer encoder.
         if self.config.use_vae and "action" in batch:
@@ -490,20 +486,21 @@ class ACT(nn.Module):
             all_cam_features = []
             all_cam_pos_embeds = []
 
-            for cam_index in range(batch["observation.images"].shape[-4]):
-                cam_features = self.backbone(batch["observation.images"][:, cam_index])["feature_map"]
-                # TODO(rcadene, alexander-soare): remove call to `.to` to speedup forward ; precompute and use
-                # buffer
+            # For a list of images, the H and W may vary but H*W is constant.
+            for img in batch["observation.images"]:
+                cam_features = self.backbone(img)["feature_map"]
                 cam_pos_embed = self.encoder_cam_feat_pos_embed(cam_features).to(dtype=cam_features.dtype)
-                cam_features = self.encoder_img_feat_input_proj(cam_features)  # (B, C, h, w)
+                cam_features = self.encoder_img_feat_input_proj(cam_features)
+
+                # Rearrange features to (sequence, batch, dim).
+                cam_features = einops.rearrange(cam_features, "b c h w -> (h w) b c")
+                cam_pos_embed = einops.rearrange(cam_pos_embed, "b c h w -> (h w) b c")
+
                 all_cam_features.append(cam_features)
                 all_cam_pos_embeds.append(cam_pos_embed)
-            # Concatenate camera observation feature maps and positional embeddings along the width dimension,
-            # and move to (sequence, batch, dim).
-            all_cam_features = torch.cat(all_cam_features, axis=-1)
-            encoder_in_tokens.extend(einops.rearrange(all_cam_features, "b c h w -> (h w) b c"))
-            all_cam_pos_embeds = torch.cat(all_cam_pos_embeds, axis=-1)
-            encoder_in_pos_embed.extend(einops.rearrange(all_cam_pos_embeds, "b c h w -> (h w) b c"))
+
+            encoder_in_tokens.extend(torch.cat(all_cam_features, axis=0))
+            encoder_in_pos_embed.extend(torch.cat(all_cam_pos_embeds, axis=0))
 
         # Stack all tokens along the sequence dimension.
         encoder_in_tokens = torch.stack(encoder_in_tokens, axis=0)

lerobot/common/policies/factory.py

@@ -16,7 +16,6 @@
 
 import logging
 
-import torch
 from torch import nn
 
 from lerobot.common.datasets.lerobot_dataset import LeRobotDatasetMetadata
@@ -76,7 +75,6 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
 
 def make_policy(
     cfg: PreTrainedConfig,
-    device: str | torch.device,
     ds_meta: LeRobotDatasetMetadata | None = None,
     env_cfg: EnvConfig | None = None,
 ) -> PreTrainedPolicy:
@@ -88,7 +86,6 @@ def make_policy(
     Args:
         cfg (PreTrainedConfig): The config of the policy to make. If `pretrained_path` is set, the policy will
             be loaded with the weights from that path.
-        device (str): the device to load the policy onto.
         ds_meta (LeRobotDatasetMetadata | None, optional): Dataset metadata to take input/output shapes and
             statistics to use for (un)normalization of inputs/outputs in the policy. Defaults to None.
         env_cfg (EnvConfig | None, optional): The config of a gym environment to parse features from. Must be
@@ -96,7 +93,7 @@ def make_policy(
 
     Raises:
         ValueError: Either ds_meta or env and env_cfg must be provided.
-        NotImplementedError: if the policy.type is 'vqbet' and the device 'mps' (due to an incompatibility)
+        NotImplementedError: if the policy.type is 'vqbet' and the policy device 'mps' (due to an incompatibility)
 
     Returns:
         PreTrainedPolicy: _description_
@@ -111,7 +108,7 @@ def make_policy(
     # https://github.com/pytorch/pytorch/issues/77764. As a temporary fix, you can set the environment
     # variable `PYTORCH_ENABLE_MPS_FALLBACK=1` to use the CPU as a fallback for this op. WARNING: this will be
     # slower than running natively on MPS.
-    if cfg.type == "vqbet" and str(device) == "mps":
+    if cfg.type == "vqbet" and cfg.device == "mps":
         raise NotImplementedError(
             "Current implementation of VQBeT does not support `mps` backend. "
             "Please use `cpu` or `cuda` backend."
@@ -145,7 +142,7 @@ def make_policy(
         # Make a fresh policy.
         policy = policy_cls(**kwargs)
 
-    policy.to(device)
+    policy.to(cfg.device)
     assert isinstance(policy, nn.Module)
 
     # policy = torch.compile(policy, mode="reduce-overhead")

lerobot/common/policies/pi0/configuration_pi0.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 dataclasses import dataclass, field
 
 from lerobot.common.optim.optimizers import AdamWConfig
@@ -76,6 +90,7 @@ class PI0Config(PreTrainedConfig):
     def __post_init__(self):
         super().__post_init__()
 
+        # TODO(Steven): Validate device and amp? in all policy configs?
         """Input validation (not exhaustive)."""
         if self.n_action_steps > self.chunk_size:
             raise ValueError(

lerobot/common/policies/pi0/conversion_scripts/benchmark.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 torch
 
 from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
@@ -31,7 +45,7 @@ def main():
 
     cfg = PreTrainedConfig.from_pretrained(ckpt_torch_dir)
     cfg.pretrained_path = ckpt_torch_dir
-    policy = make_policy(cfg, device, ds_meta=dataset.meta)
+    policy = make_policy(cfg, ds_meta=dataset.meta)
 
     # policy = torch.compile(policy, mode="reduce-overhead")
 

lerobot/common/policies/pi0/conversion_scripts/compare_with_jax.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 json
 import pickle
 from pathlib import Path
@@ -87,7 +101,7 @@ def main():
 
     cfg = PreTrainedConfig.from_pretrained(ckpt_torch_dir)
     cfg.pretrained_path = ckpt_torch_dir
-    policy = make_policy(cfg, device, dataset_meta)
+    policy = make_policy(cfg, dataset_meta)
 
     # loss_dict = policy.forward(batch, noise=noise, time=time_beta)
     # loss_dict["loss"].backward()

lerobot/common/policies/pi0/conversion_scripts/conversion_utils.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 transformers import GemmaConfig, PaliGemmaConfig
 
 

lerobot/common/policies/pi0/conversion_scripts/convert_pi0_to_hf_lerobot.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.
+
 """
 Convert pi0 parameters from Jax to Pytorch
 

lerobot/common/policies/pi0/flex_attention.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 torch
 import torch.nn.functional as F  # noqa: N812
 from packaging.version import Version

lerobot/common/policies/pi0/modeling_pi0.py

@@ -313,7 +313,7 @@ class PI0Policy(PreTrainedPolicy):
         state = self.prepare_state(batch)
         lang_tokens, lang_masks = self.prepare_language(batch)
         actions = self.prepare_action(batch)
-        actions_is_pad = batch.get("actions_id_pad")
+        actions_is_pad = batch.get("actions_is_pad")
 
         loss_dict = {}
         losses = self.model.forward(images, img_masks, lang_tokens, lang_masks, state, actions, noise, time)

lerobot/common/policies/pi0/paligemma_with_expert.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 typing import List, Optional, Union
 
 import torch

lerobot/common/policies/pretrained.py

@@ -1,3 +1,16 @@
+# 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
 import logging
 import os
@@ -73,7 +86,6 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
         cache_dir: str | Path | None = None,
         local_files_only: bool = False,
         revision: str | None = None,
-        map_location: str = "cpu",
         strict: bool = False,
         **kwargs,
     ) -> T:
@@ -98,7 +110,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
         if os.path.isdir(model_id):
             print("Loading weights from local directory")
             model_file = os.path.join(model_id, SAFETENSORS_SINGLE_FILE)
-            policy = cls._load_as_safetensor(instance, model_file, map_location, strict)
+            policy = cls._load_as_safetensor(instance, model_file, config.device, strict)
         else:
             try:
                 model_file = hf_hub_download(
@@ -112,13 +124,13 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
                     token=token,
                     local_files_only=local_files_only,
                 )
-                policy = cls._load_as_safetensor(instance, model_file, map_location, strict)
+                policy = cls._load_as_safetensor(instance, model_file, config.device, strict)
             except HfHubHTTPError as e:
                 raise FileNotFoundError(
                     f"{SAFETENSORS_SINGLE_FILE} not found on the HuggingFace Hub in {model_id}"
                 ) from e
 
-        policy.to(map_location)
+        policy.to(config.device)
         policy.eval()
         return policy
 

lerobot/common/robot_devices/cameras/configs.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 abc
 from dataclasses import dataclass
 

lerobot/common/robot_devices/cameras/intelrealsense.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.
 """
@@ -34,7 +48,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 +100,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
 
@@ -100,7 +114,7 @@ def save_images_from_cameras(
         camera = IntelRealSenseCamera(config)
         camera.connect()
         print(
-            f"IntelRealSenseCamera({camera.serial_number}, fps={camera.fps}, width={camera.width}, height={camera.height}, color_mode={camera.color_mode})"
+            f"IntelRealSenseCamera({camera.serial_number}, fps={camera.fps}, width={camera.capture_width}, height={camera.capture_height}, color_mode={camera.color_mode})"
         )
         cameras.append(camera)
 
@@ -210,9 +224,20 @@ class IntelRealSenseCamera:
             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 +253,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
@@ -263,22 +287,26 @@ class IntelRealSenseCamera:
             )
 
         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)
 
@@ -316,18 +344,18 @@ class IntelRealSenseCamera:
             raise OSError(
                 f"Can't set {self.fps=} for IntelRealSenseCamera({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 IntelRealSenseCamera({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 IntelRealSenseCamera({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
 
@@ -347,7 +375,7 @@ class IntelRealSenseCamera:
             )
 
         if self.mock:
-            import tests.mock_cv2 as cv2
+            import tests.cameras.mock_cv2 as cv2
         else:
             import cv2
 
@@ -373,7 +401,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."
             )
@@ -395,7 +423,7 @@ class IntelRealSenseCamera:
             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."
                 )

lerobot/common/robot_devices/cameras/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.
 """
@@ -66,7 +80,7 @@ def _find_cameras(
     possible_camera_ids: list[int | str], raise_when_empty=False, mock=False
 ) -> list[int | str]:
     if mock:
-        import tests.mock_cv2 as cv2
+        import tests.cameras.mock_cv2 as cv2
     else:
         import cv2
 
@@ -130,8 +144,8 @@ def save_images_from_cameras(
         camera = OpenCVCamera(config)
         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)
 
@@ -230,9 +244,19 @@ 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
@@ -245,11 +269,10 @@ class OpenCVCamera:
         self.logs = {}
 
         if self.mock:
-            import tests.mock_cv2 as cv2
+            import tests.cameras.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
@@ -263,7 +286,7 @@ class OpenCVCamera:
             raise RobotDeviceAlreadyConnectedError(f"OpenCVCamera({self.camera_index}) is already connected.")
 
         if self.mock:
-            import tests.mock_cv2 as cv2
+            import tests.cameras.mock_cv2 as cv2
         else:
             import cv2
 
@@ -271,10 +294,20 @@ class OpenCVCamera:
             # when other threads are used to save the images.
             cv2.setNumThreads(1)
 
+        backend = (
+            cv2.CAP_V4L2
+            if platform.system() == "Linux"
+            else cv2.CAP_DSHOW
+            if platform.system() == "Windows"
+            else cv2.CAP_AVFOUNDATION
+            if platform.system() == "Darwin"
+            else cv2.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.VideoCapture(camera_idx, backend)
         is_camera_open = tmp_camera.isOpened()
         # Release camera to make it accessible for `find_camera_indices`
         tmp_camera.release()
@@ -297,14 +330,14 @@ 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.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)
+        if self.capture_width is not None:
+            self.camera.set(cv2.CAP_PROP_FRAME_WIDTH, self.capture_width)
+        if self.capture_height is not None:
+            self.camera.set(cv2.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)
@@ -316,19 +349,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:
@@ -362,14 +398,14 @@ class OpenCVCamera:
         # so we convert the image color from BGR to RGB.
         if requested_color_mode == "rgb":
             if self.mock:
-                import tests.mock_cv2 as cv2
+                import tests.cameras.mock_cv2 as cv2
             else:
                 import cv2
 
             color_image = cv2.cvtColor(color_image, cv2.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."
             )

lerobot/common/robot_devices/cameras/utils.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 typing import Protocol
 
 import numpy as np
@@ -31,7 +45,7 @@ def make_cameras_from_configs(camera_configs: dict[str, CameraConfig]) -> list[C
 
             cameras[key] = IntelRealSenseCamera(cfg)
         else:
-            raise ValueError(f"The motor type '{cfg.type}' is not valid.")
+            raise ValueError(f"The camera type '{cfg.type}' is not valid.")
 
     return cameras
 

lerobot/common/robot_devices/control_configs.py

@@ -1,14 +1,25 @@
-import logging
+# 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
 from pathlib import Path
 
 import draccus
 
 from lerobot.common.robot_devices.robots.configs import RobotConfig
-from lerobot.common.utils.utils import auto_select_torch_device, is_amp_available, is_torch_device_available
 from lerobot.configs import parser
 from lerobot.configs.policies import PreTrainedConfig
-from lerobot.configs.train import TrainPipelineConfig
 
 
 @dataclass
@@ -43,11 +54,6 @@ class RecordControlConfig(ControlConfig):
     # Root directory where the dataset will be stored (e.g. 'dataset/path').
     root: str | Path | None = None
     policy: PreTrainedConfig | None = None
-    # TODO(rcadene, aliberts): By default, use device and use_amp values from policy checkpoint.
-    device: str | None = None  # cuda | cpu | mps
-    # `use_amp` determines whether to use Automatic Mixed Precision (AMP) for training and evaluation. With AMP,
-    # automatic gradient scaling is used.
-    use_amp: bool | None = None
     # Limit the frames per second. By default, uses the policy fps.
     fps: int | None = None
     # Number of seconds before starting data collection. It allows the robot devices to warmup and synchronize.
@@ -90,27 +96,6 @@ class RecordControlConfig(ControlConfig):
             self.policy = PreTrainedConfig.from_pretrained(policy_path, cli_overrides=cli_overrides)
             self.policy.pretrained_path = policy_path
 
-            # When no device or use_amp are given, use the one from training config.
-            if self.device is None or self.use_amp is None:
-                train_cfg = TrainPipelineConfig.from_pretrained(policy_path)
-                if self.device is None:
-                    self.device = train_cfg.device
-                if self.use_amp is None:
-                    self.use_amp = train_cfg.use_amp
-
-            # Automatically switch to available device if necessary
-            if not is_torch_device_available(self.device):
-                auto_device = auto_select_torch_device()
-                logging.warning(f"Device '{self.device}' is not available. Switching to '{auto_device}'.")
-                self.device = auto_device
-
-            # Automatically deactivate AMP if necessary
-            if self.use_amp and not is_amp_available(self.device):
-                logging.warning(
-                    f"Automatic Mixed Precision (amp) is not available on device '{self.device}'. Deactivating AMP."
-                )
-                self.use_amp = False
-
 
 @ControlConfig.register_subclass("replay")
 @dataclass

lerobot/common/robot_devices/control_utils.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.
+
 ########################################################################################
 # Utilities
 ########################################################################################
@@ -18,6 +32,7 @@ from termcolor import colored
 from lerobot.common.datasets.image_writer import safe_stop_image_writer
 from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.common.datasets.utils import get_features_from_robot
+from lerobot.common.policies.pretrained import PreTrainedPolicy
 from lerobot.common.robot_devices.robots.utils import Robot
 from lerobot.common.robot_devices.utils import busy_wait
 from lerobot.common.utils.utils import get_safe_torch_device, has_method
@@ -179,8 +194,6 @@ def record_episode(
     episode_time_s,
     display_cameras,
     policy,
-    device,
-    use_amp,
     fps,
     single_task,
 ):
@@ -191,8 +204,6 @@ def record_episode(
         dataset=dataset,
         events=events,
         policy=policy,
-        device=device,
-        use_amp=use_amp,
         fps=fps,
         teleoperate=policy is None,
         single_task=single_task,
@@ -207,9 +218,7 @@ def control_loop(
     display_cameras=False,
     dataset: LeRobotDataset | None = None,
     events=None,
-    policy=None,
-    device: torch.device | str | None = None,
-    use_amp: bool | None = None,
+    policy: PreTrainedPolicy = None,
     fps: int | None = None,
     single_task: str | None = None,
 ):
@@ -232,9 +241,6 @@ def control_loop(
     if dataset is not None and fps is not None and dataset.fps != fps:
         raise ValueError(f"The dataset fps should be equal to requested fps ({dataset['fps']} != {fps}).")
 
-    if isinstance(device, str):
-        device = get_safe_torch_device(device)
-
     timestamp = 0
     start_episode_t = time.perf_counter()
     while timestamp < control_time_s:
@@ -246,7 +252,9 @@ def control_loop(
             observation = robot.capture_observation()
 
             if policy is not None:
-                pred_action = predict_action(observation, policy, device, use_amp)
+                pred_action = predict_action(
+                    observation, policy, get_safe_torch_device(policy.config.device), policy.config.use_amp
+                )
                 # Action can eventually be clipped using `max_relative_target`,
                 # so action actually sent is saved in the dataset.
                 action = robot.send_action(pred_action)

lerobot/common/robot_devices/motors/configs.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 abc
 from dataclasses import dataclass
 

lerobot/common/robot_devices/motors/dynamixel.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 enum
 import logging
 import math
@@ -318,7 +332,7 @@ class DynamixelMotorsBus:
             )
 
         if self.mock:
-            import tests.mock_dynamixel_sdk as dxl
+            import tests.motors.mock_dynamixel_sdk as dxl
         else:
             import dynamixel_sdk as dxl
 
@@ -342,7 +356,7 @@ class DynamixelMotorsBus:
 
     def reconnect(self):
         if self.mock:
-            import tests.mock_dynamixel_sdk as dxl
+            import tests.motors.mock_dynamixel_sdk as dxl
         else:
             import dynamixel_sdk as dxl
 
@@ -632,7 +646,7 @@ class DynamixelMotorsBus:
 
     def read_with_motor_ids(self, motor_models, motor_ids, data_name, num_retry=NUM_READ_RETRY):
         if self.mock:
-            import tests.mock_dynamixel_sdk as dxl
+            import tests.motors.mock_dynamixel_sdk as dxl
         else:
             import dynamixel_sdk as dxl
 
@@ -677,7 +691,7 @@ class DynamixelMotorsBus:
         start_time = time.perf_counter()
 
         if self.mock:
-            import tests.mock_dynamixel_sdk as dxl
+            import tests.motors.mock_dynamixel_sdk as dxl
         else:
             import dynamixel_sdk as dxl
 
@@ -743,7 +757,7 @@ class DynamixelMotorsBus:
 
     def write_with_motor_ids(self, motor_models, motor_ids, data_name, values, num_retry=NUM_WRITE_RETRY):
         if self.mock:
-            import tests.mock_dynamixel_sdk as dxl
+            import tests.motors.mock_dynamixel_sdk as dxl
         else:
             import dynamixel_sdk as dxl
 
@@ -779,7 +793,7 @@ class DynamixelMotorsBus:
         start_time = time.perf_counter()
 
         if self.mock:
-            import tests.mock_dynamixel_sdk as dxl
+            import tests.motors.mock_dynamixel_sdk as dxl
         else:
             import dynamixel_sdk as dxl
 

lerobot/common/robot_devices/motors/feetech.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 enum
 import logging
 import math
@@ -299,7 +313,7 @@ class FeetechMotorsBus:
             )
 
         if self.mock:
-            import tests.mock_scservo_sdk as scs
+            import tests.motors.mock_scservo_sdk as scs
         else:
             import scservo_sdk as scs
 
@@ -323,7 +337,7 @@ class FeetechMotorsBus:
 
     def reconnect(self):
         if self.mock:
-            import tests.mock_scservo_sdk as scs
+            import tests.motors.mock_scservo_sdk as scs
         else:
             import scservo_sdk as scs
 
@@ -650,7 +664,7 @@ class FeetechMotorsBus:
 
     def read_with_motor_ids(self, motor_models, motor_ids, data_name, num_retry=NUM_READ_RETRY):
         if self.mock:
-            import tests.mock_scservo_sdk as scs
+            import tests.motors.mock_scservo_sdk as scs
         else:
             import scservo_sdk as scs
 
@@ -688,7 +702,7 @@ class FeetechMotorsBus:
 
     def read(self, data_name, motor_names: str | list[str] | None = None):
         if self.mock:
-            import tests.mock_scservo_sdk as scs
+            import tests.motors.mock_scservo_sdk as scs
         else:
             import scservo_sdk as scs
 
@@ -768,7 +782,7 @@ class FeetechMotorsBus:
 
     def write_with_motor_ids(self, motor_models, motor_ids, data_name, values, num_retry=NUM_WRITE_RETRY):
         if self.mock:
-            import tests.mock_scservo_sdk as scs
+            import tests.motors.mock_scservo_sdk as scs
         else:
             import scservo_sdk as scs
 
@@ -804,7 +818,7 @@ class FeetechMotorsBus:
         start_time = time.perf_counter()
 
         if self.mock:
-            import tests.mock_scservo_sdk as scs
+            import tests.motors.mock_scservo_sdk as scs
         else:
             import scservo_sdk as scs
 

lerobot/common/robot_devices/motors/utils.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 typing import Protocol
 
 from lerobot.common.robot_devices.motors.configs import (

lerobot/common/robot_devices/robots/configs.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 abc
 from dataclasses import dataclass, field
 from typing import Sequence

lerobot/common/robot_devices/robots/dynamixel_calibration.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.
+
 """Logic to calibrate a robot arm built with dynamixel motors"""
 # TODO(rcadene, aliberts): move this logic into the robot code when refactoring
 

lerobot/common/robot_devices/robots/feetech_calibration.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.
+
 """Logic to calibrate a robot arm built with feetech motors"""
 # TODO(rcadene, aliberts): move this logic into the robot code when refactoring
 

lerobot/common/robot_devices/robots/lekiwi_remote.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 base64
 import json
 import threading

lerobot/common/robot_devices/robots/manipulator.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.
+
 """Contains logic to instantiate a robot, read information from its motors and cameras,
 and send orders to its motors.
 """

lerobot/common/robot_devices/robots/mobile_manipulator.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 base64
 import json
 import os
@@ -392,21 +406,19 @@ class MobileManipulator:
         for name in self.leader_arms:
             pos = self.leader_arms[name].read("Present_Position")
             pos_tensor = torch.from_numpy(pos).float()
-            # Instead of pos_tensor.item(), use tolist() to convert the entire tensor to a list
             arm_positions.extend(pos_tensor.tolist())
 
-        # (The rest of your code for generating wheel commands remains unchanged)
-        x_cmd = 0.0  # m/s forward/backward
-        y_cmd = 0.0  # m/s lateral
+        y_cmd = 0.0  # m/s forward/backward
+        x_cmd = 0.0  # m/s lateral
         theta_cmd = 0.0  # deg/s rotation
         if self.pressed_keys["forward"]:
-            x_cmd += xy_speed
-        if self.pressed_keys["backward"]:
-            x_cmd -= xy_speed
-        if self.pressed_keys["left"]:
             y_cmd += xy_speed
-        if self.pressed_keys["right"]:
+        if self.pressed_keys["backward"]:
             y_cmd -= xy_speed
+        if self.pressed_keys["left"]:
+            x_cmd += xy_speed
+        if self.pressed_keys["right"]:
+            x_cmd -= xy_speed
         if self.pressed_keys["rotate_left"]:
             theta_cmd += theta_speed
         if self.pressed_keys["rotate_right"]:
@@ -584,8 +596,8 @@ class MobileManipulator:
         # Create the body velocity vector [x, y, theta_rad].
         velocity_vector = np.array([x_cmd, y_cmd, theta_rad])
 
-        # Define the wheel mounting angles with a -90° offset.
-        angles = np.radians(np.array([240, 120, 0]) - 90)
+        # Define the wheel mounting angles (defined from y axis cw)
+        angles = np.radians(np.array([300, 180, 60]))
         # Build the kinematic matrix: each row maps body velocities to a wheel’s linear speed.
         # The third column (base_radius) accounts for the effect of rotation.
         m = np.array([[np.cos(a), np.sin(a), base_radius] for a in angles])
@@ -641,8 +653,8 @@ class MobileManipulator:
         # Compute each wheel’s linear speed (m/s) from its angular speed.
         wheel_linear_speeds = wheel_radps * wheel_radius
 
-        # Define the wheel mounting angles with a -90° offset.
-        angles = np.radians(np.array([240, 120, 0]) - 90)
+        # Define the wheel mounting angles (defined from y axis cw)
+        angles = np.radians(np.array([300, 180, 60]))
         m = np.array([[np.cos(a), np.sin(a), base_radius] for a in angles])
 
         # Solve the inverse kinematics: body_velocity = M⁻¹ · wheel_linear_speeds.

lerobot/common/robot_devices/robots/utils.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 typing import Protocol
 
 from lerobot.common.robot_devices.robots.configs import (

lerobot/common/robot_devices/utils.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 platform
 import time
 

lerobot/common/utils/hub.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 pathlib import Path
 from tempfile import TemporaryDirectory
 from typing import Any, Type, TypeVar

lerobot/common/utils/utils.py

@@ -17,6 +17,7 @@ import logging
 import os
 import os.path as osp
 import platform
+import subprocess
 from copy import copy
 from datetime import datetime, timezone
 from pathlib import Path
@@ -50,8 +51,10 @@ def auto_select_torch_device() -> torch.device:
         return torch.device("cpu")
 
 
+# TODO(Steven): Remove log. log shouldn't be an argument, this should be handled by the logger level
 def get_safe_torch_device(try_device: str, log: bool = False) -> torch.device:
     """Given a string, return a torch.device with checks on whether the device is available."""
+    try_device = str(try_device)
     match try_device:
         case "cuda":
             assert torch.cuda.is_available()
@@ -84,6 +87,7 @@ def get_safe_dtype(dtype: torch.dtype, device: str | torch.device):
 
 
 def is_torch_device_available(try_device: str) -> bool:
+    try_device = str(try_device)  # Ensure try_device is a string
     if try_device == "cuda":
         return torch.cuda.is_available()
     elif try_device == "mps":
@@ -91,7 +95,7 @@ def is_torch_device_available(try_device: str) -> bool:
     elif try_device == "cpu":
         return True
     else:
-        raise ValueError(f"Unknown device '{try_device}.")
+        raise ValueError(f"Unknown device {try_device}. Supported devices are: cuda, mps or cpu.")
 
 
 def is_amp_available(device: str):
@@ -165,23 +169,31 @@ def capture_timestamp_utc():
 
 
 def say(text, blocking=False):
-    # Check if mac, linux, or windows.
-    if platform.system() == "Darwin":
-        cmd = f'say "{text}"'
-        if not blocking:
-            cmd += " &"
-    elif platform.system() == "Linux":
-        cmd = f'spd-say "{text}"'
-        if blocking:
-            cmd += "  --wait"
-    elif platform.system() == "Windows":
-        # TODO(rcadene): Make blocking option work for Windows
-        cmd = (
-            'PowerShell -Command "Add-Type -AssemblyName System.Speech; '
-            f"(New-Object System.Speech.Synthesis.SpeechSynthesizer).Speak('{text}')\""
-        )
+    system = platform.system()
 
-    os.system(cmd)
+    if system == "Darwin":
+        cmd = ["say", text]
+
+    elif system == "Linux":
+        cmd = ["spd-say", text]
+        if blocking:
+            cmd.append("--wait")
+
+    elif system == "Windows":
+        cmd = [
+            "PowerShell",
+            "-Command",
+            "Add-Type -AssemblyName System.Speech; "
+            f"(New-Object System.Speech.Synthesis.SpeechSynthesizer).Speak('{text}')",
+        ]
+
+    else:
+        raise RuntimeError("Unsupported operating system for text-to-speech.")
+
+    if blocking:
+        subprocess.run(cmd, check=True)
+    else:
+        subprocess.Popen(cmd, creationflags=subprocess.CREATE_NO_WINDOW if system == "Windows" else 0)
 
 
 def log_say(text, play_sounds, blocking=False):

lerobot/common/utils/wandb_utils.py

@@ -69,7 +69,13 @@ class WandBLogger:
         os.environ["WANDB_SILENT"] = "True"
         import wandb
 
-        wandb_run_id = get_wandb_run_id_from_filesystem(self.log_dir) if cfg.resume else None
+        wandb_run_id = (
+            cfg.wandb.run_id
+            if cfg.wandb.run_id
+            else get_wandb_run_id_from_filesystem(self.log_dir)
+            if cfg.resume
+            else None
+        )
         wandb.init(
             id=wandb_run_id,
             project=self.cfg.project,

lerobot/configs/default.py

@@ -20,6 +20,7 @@ from lerobot.common import (
     policies,  # noqa: F401
 )
 from lerobot.common.datasets.transforms import ImageTransformsConfig
+from lerobot.common.datasets.video_utils import get_safe_default_codec
 
 
 @dataclass
@@ -35,7 +36,7 @@ class DatasetConfig:
     image_transforms: ImageTransformsConfig = field(default_factory=ImageTransformsConfig)
     revision: str | None = None
     use_imagenet_stats: bool = True
-    video_backend: str = "pyav"
+    video_backend: str = field(default_factory=get_safe_default_codec)
 
 
 @dataclass
@@ -46,6 +47,7 @@ class WandBConfig:
     project: str = "lerobot"
     entity: str | None = None
     notes: str | None = None
+    run_id: str | None = None
 
 
 @dataclass

lerobot/configs/eval.py

@@ -1,14 +1,26 @@
+# 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 datetime as dt
 import logging
 from dataclasses import dataclass, field
 from pathlib import Path
 
 from lerobot.common import envs, policies  # noqa: F401
-from lerobot.common.utils.utils import auto_select_torch_device, is_amp_available, is_torch_device_available
 from lerobot.configs import parser
 from lerobot.configs.default import EvalConfig
 from lerobot.configs.policies import PreTrainedConfig
-from lerobot.configs.train import TrainPipelineConfig
 
 
 @dataclass
@@ -21,11 +33,6 @@ class EvalPipelineConfig:
     policy: PreTrainedConfig | None = None
     output_dir: Path | None = None
     job_name: str | None = None
-    # TODO(rcadene, aliberts): By default, use device and use_amp values from policy checkpoint.
-    device: str | None = None  # cuda | cpu | mps
-    # `use_amp` determines whether to use Automatic Mixed Precision (AMP) for training and evaluation. With AMP,
-    # automatic gradient scaling is used.
-    use_amp: bool = False
     seed: int | None = 1000
 
     def __post_init__(self):
@@ -36,27 +43,6 @@ class EvalPipelineConfig:
             self.policy = PreTrainedConfig.from_pretrained(policy_path, cli_overrides=cli_overrides)
             self.policy.pretrained_path = policy_path
 
-            # When no device or use_amp are given, use the one from training config.
-            if self.device is None or self.use_amp is None:
-                train_cfg = TrainPipelineConfig.from_pretrained(policy_path)
-                if self.device is None:
-                    self.device = train_cfg.device
-                if self.use_amp is None:
-                    self.use_amp = train_cfg.use_amp
-
-            # Automatically switch to available device if necessary
-            if not is_torch_device_available(self.device):
-                auto_device = auto_select_torch_device()
-                logging.warning(f"Device '{self.device}' is not available. Switching to '{auto_device}'.")
-                self.device = auto_device
-
-            # Automatically deactivate AMP if necessary
-            if self.use_amp and not is_amp_available(self.device):
-                logging.warning(
-                    f"Automatic Mixed Precision (amp) is not available on device '{self.device}'. Deactivating AMP."
-                )
-                self.use_amp = False
-
         else:
             logging.warning(
                 "No pretrained path was provided, evaluated policy will be built from scratch (random weights)."
@@ -73,11 +59,6 @@ class EvalPipelineConfig:
             eval_dir = f"{now:%Y-%m-%d}/{now:%H-%M-%S}_{self.job_name}"
             self.output_dir = Path("outputs/eval") / eval_dir
 
-        if self.device is None:
-            raise ValueError("Set one of the following device: cuda, cpu or mps")
-        elif self.device == "cuda" and self.use_amp is None:
-            raise ValueError("Set 'use_amp' to True or False.")
-
     @classmethod
     def __get_path_fields__(cls) -> list[str]:
         """This enables the parser to load config from the policy using `--policy.path=local/dir`"""

lerobot/configs/parser.py

@@ -1,4 +1,19 @@
+# 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 importlib
 import inspect
+import pkgutil
 import sys
 from argparse import ArgumentError
 from functools import wraps
@@ -10,6 +25,7 @@ import draccus
 from lerobot.common.utils.utils import has_method
 
 PATH_KEY = "path"
+PLUGIN_DISCOVERY_SUFFIX = "discover_packages_path"
 draccus.set_config_type("json")
 
 
@@ -45,6 +61,86 @@ def parse_arg(arg_name: str, args: Sequence[str] | None = None) -> str | None:
     return None
 
 
+def parse_plugin_args(plugin_arg_suffix: str, args: Sequence[str]) -> dict:
+    """Parse plugin-related arguments from command-line arguments.
+
+    This function extracts arguments from command-line arguments that match a specified suffix pattern.
+    It processes arguments in the format '--key=value' and returns them as a dictionary.
+
+    Args:
+        plugin_arg_suffix (str): The suffix to identify plugin-related arguments.
+        cli_args (Sequence[str]): A sequence of command-line arguments to parse.
+
+    Returns:
+        dict: A dictionary containing the parsed plugin arguments where:
+            - Keys are the argument names (with '--' prefix removed if present)
+            - Values are the corresponding argument values
+
+    Example:
+        >>> args = ['--env.discover_packages_path=my_package',
+        ...         '--other_arg=value']
+        >>> parse_plugin_args('discover_packages_path', args)
+        {'env.discover_packages_path': 'my_package'}
+    """
+    plugin_args = {}
+    for arg in args:
+        if "=" in arg and plugin_arg_suffix in arg:
+            key, value = arg.split("=", 1)
+            # Remove leading '--' if present
+            if key.startswith("--"):
+                key = key[2:]
+            plugin_args[key] = value
+    return plugin_args
+
+
+class PluginLoadError(Exception):
+    """Raised when a plugin fails to load."""
+
+
+def load_plugin(plugin_path: str) -> None:
+    """Load and initialize a plugin from a given Python package path.
+
+    This function attempts to load a plugin by importing its package and any submodules.
+    Plugin registration is expected to happen during package initialization, i.e. when
+    the package is imported the gym environment should be registered and the config classes
+    registered with their parents using the `register_subclass` decorator.
+
+    Args:
+        plugin_path (str): The Python package path to the plugin (e.g. "mypackage.plugins.myplugin")
+
+    Raises:
+        PluginLoadError: If the plugin cannot be loaded due to import errors or if the package path is invalid.
+
+    Examples:
+        >>> load_plugin("external_plugin.core")       # Loads plugin from external package
+
+    Notes:
+        - The plugin package should handle its own registration during import
+        - All submodules in the plugin package will be imported
+        - Implementation follows the plugin discovery pattern from Python packaging guidelines
+
+    See Also:
+        https://packaging.python.org/en/latest/guides/creating-and-discovering-plugins/
+    """
+    try:
+        package_module = importlib.import_module(plugin_path, __package__)
+    except (ImportError, ModuleNotFoundError) as e:
+        raise PluginLoadError(
+            f"Failed to load plugin '{plugin_path}'. Verify the path and installation: {str(e)}"
+        ) from e
+
+    def iter_namespace(ns_pkg):
+        return pkgutil.iter_modules(ns_pkg.__path__, ns_pkg.__name__ + ".")
+
+    try:
+        for _finder, pkg_name, _ispkg in iter_namespace(package_module):
+            importlib.import_module(pkg_name)
+    except ImportError as e:
+        raise PluginLoadError(
+            f"Failed to load plugin '{plugin_path}'. Verify the path and installation: {str(e)}"
+        ) from e
+
+
 def get_path_arg(field_name: str, args: Sequence[str] | None = None) -> str | None:
     return parse_arg(f"{field_name}.{PATH_KEY}", args)
 
@@ -92,10 +188,13 @@ def filter_path_args(fields_to_filter: str | list[str], args: Sequence[str] | No
 
 def wrap(config_path: Path | None = None):
     """
-    HACK: Similar to draccus.wrap but does two additional things:
+    HACK: Similar to draccus.wrap but does three additional things:
         - Will remove '.path' arguments from CLI in order to process them later on.
         - If a 'config_path' is passed and the main config class has a 'from_pretrained' method, will
           initialize it from there to allow to fetch configs from the hub directly
+        - Will load plugins specified in the CLI arguments. These plugins will typically register
+            their own subclasses of config classes, so that draccus can find the right class to instantiate
+            from the CLI '.type' arguments
     """
 
     def wrapper_outer(fn):
@@ -108,6 +207,14 @@ def wrap(config_path: Path | None = None):
                 args = args[1:]
             else:
                 cli_args = sys.argv[1:]
+                plugin_args = parse_plugin_args(PLUGIN_DISCOVERY_SUFFIX, cli_args)
+                for plugin_cli_arg, plugin_path in plugin_args.items():
+                    try:
+                        load_plugin(plugin_path)
+                    except PluginLoadError as e:
+                        # add the relevant CLI arg to the error message
+                        raise PluginLoadError(f"{e}\nFailed plugin CLI Arg: {plugin_cli_arg}") from e
+                    cli_args = filter_arg(plugin_cli_arg, cli_args)
                 config_path_cli = parse_arg("config_path", cli_args)
                 if has_method(argtype, "__get_path_fields__"):
                     path_fields = argtype.__get_path_fields__()

lerobot/configs/policies.py

@@ -1,4 +1,18 @@
+# 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
+import logging
 import os
 from dataclasses import dataclass, field
 from pathlib import Path
@@ -12,6 +26,7 @@ from huggingface_hub.errors import HfHubHTTPError
 from lerobot.common.optim.optimizers import OptimizerConfig
 from lerobot.common.optim.schedulers import LRSchedulerConfig
 from lerobot.common.utils.hub import HubMixin
+from lerobot.common.utils.utils import auto_select_torch_device, is_amp_available, is_torch_device_available
 from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
 
 # Generic variable that is either PreTrainedConfig or a subclass thereof
@@ -40,8 +55,24 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
     input_features: dict[str, PolicyFeature] = field(default_factory=dict)
     output_features: dict[str, PolicyFeature] = field(default_factory=dict)
 
+    device: str | None = None  # cuda | cpu | mp
+    # `use_amp` determines whether to use Automatic Mixed Precision (AMP) for training and evaluation. With AMP,
+    # automatic gradient scaling is used.
+    use_amp: bool = False
+
     def __post_init__(self):
         self.pretrained_path = None
+        if not self.device or not is_torch_device_available(self.device):
+            auto_device = auto_select_torch_device()
+            logging.warning(f"Device '{self.device}' is not available. Switching to '{auto_device}'.")
+            self.device = auto_device.type
+
+        # Automatically deactivate AMP if necessary
+        if self.use_amp and not is_amp_available(self.device):
+            logging.warning(
+                f"Automatic Mixed Precision (amp) is not available on device '{self.device}'. Deactivating AMP."
+            )
+            self.use_amp = False
 
     @property
     def type(self) -> str:

lerobot/configs/train.py

@@ -1,5 +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 datetime as dt
-import logging
 import os
 from dataclasses import dataclass, field
 from pathlib import Path
@@ -13,7 +25,6 @@ from lerobot.common import envs
 from lerobot.common.optim import OptimizerConfig
 from lerobot.common.optim.schedulers import LRSchedulerConfig
 from lerobot.common.utils.hub import HubMixin
-from lerobot.common.utils.utils import auto_select_torch_device, is_amp_available
 from lerobot.configs import parser
 from lerobot.configs.default import DatasetConfig, EvalConfig, WandBConfig
 from lerobot.configs.policies import PreTrainedConfig
@@ -35,10 +46,6 @@ class TrainPipelineConfig(HubMixin):
     # Note that when resuming a run, the default behavior is to use the configuration from the checkpoint,
     # regardless of what's provided with the training command at the time of resumption.
     resume: bool = False
-    device: str | None = None  # cuda | cpu | mp
-    # `use_amp` determines whether to use Automatic Mixed Precision (AMP) for training and evaluation. With AMP,
-    # automatic gradient scaling is used.
-    use_amp: bool = False
     # `seed` is used for training (eg: model initialization, dataset shuffling)
     # AND for the evaluation environments.
     seed: int | None = 1000
@@ -61,18 +68,6 @@ class TrainPipelineConfig(HubMixin):
         self.checkpoint_path = None
 
     def validate(self):
-        if not self.device:
-            logging.warning("No device specified, trying to infer device automatically")
-            device = auto_select_torch_device()
-            self.device = device.type
-
-        # Automatically deactivate AMP if necessary
-        if self.use_amp and not is_amp_available(self.device):
-            logging.warning(
-                f"Automatic Mixed Precision (amp) is not available on device '{self.device}'. Deactivating AMP."
-            )
-            self.use_amp = False
-
         # HACK: We parse again the cli args here to get the pretrained paths if there was some.
         policy_path = parser.get_path_arg("policy")
         if policy_path:
@@ -84,7 +79,9 @@ class TrainPipelineConfig(HubMixin):
             # The entire train config is already loaded, we just need to get the checkpoint dir
             config_path = parser.parse_arg("config_path")
             if not config_path:
-                raise ValueError("A config_path is expected when resuming a run.")
+                raise ValueError(
+                    f"A config_path is expected when resuming a run. Please specify path to {TRAIN_CONFIG_NAME}"
+                )
             if not Path(config_path).resolve().exists():
                 raise NotADirectoryError(
                     f"{config_path=} is expected to be a local path. "

lerobot/configs/types.py

@@ -1,3 +1,16 @@
+# 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.
 # Note: We subclass str so that serialization is straightforward
 # https://stackoverflow.com/questions/24481852/serialising-an-enum-member-to-json
 from dataclasses import dataclass

lerobot/scripts/configure_motor.py

@@ -1,3 +1,16 @@
+# 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 configure a single motor at a time to a given ID and baudrate.
 

lerobot/scripts/control_robot.py

@@ -1,3 +1,16 @@
+# 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.
 """
 Utilities to control a robot.
 
@@ -254,7 +267,7 @@ def record(
         )
 
     # Load pretrained policy
-    policy = None if cfg.policy is None else make_policy(cfg.policy, cfg.device, ds_meta=dataset.meta)
+    policy = None if cfg.policy is None else make_policy(cfg.policy, ds_meta=dataset.meta)
 
     if not robot.is_connected:
         robot.connect()
@@ -285,8 +298,6 @@ def record(
             episode_time_s=cfg.episode_time_s,
             display_cameras=cfg.display_cameras,
             policy=policy,
-            device=cfg.device,
-            use_amp=cfg.use_amp,
             fps=cfg.fps,
             single_task=cfg.single_task,
         )

lerobot/scripts/control_sim_robot.py

@@ -1,3 +1,16 @@
+# 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.
 """
 Utilities to control a robot in simulation.
 

lerobot/scripts/eval.py

@@ -454,11 +454,11 @@ def _compile_episode_data(
 
 
 @parser.wrap()
-def eval(cfg: EvalPipelineConfig):
+def eval_main(cfg: EvalPipelineConfig):
     logging.info(pformat(asdict(cfg)))
 
     # Check device is available
-    device = get_safe_torch_device(cfg.device, log=True)
+    device = get_safe_torch_device(cfg.policy.device, log=True)
 
     torch.backends.cudnn.benchmark = True
     torch.backends.cuda.matmul.allow_tf32 = True
@@ -470,14 +470,14 @@ def eval(cfg: EvalPipelineConfig):
     env = make_env(cfg.env, n_envs=cfg.eval.batch_size, use_async_envs=cfg.eval.use_async_envs)
 
     logging.info("Making policy.")
+
     policy = make_policy(
         cfg=cfg.policy,
-        device=device,
         env_cfg=cfg.env,
     )
     policy.eval()
 
-    with torch.no_grad(), torch.autocast(device_type=device.type) if cfg.use_amp else nullcontext():
+    with torch.no_grad(), torch.autocast(device_type=device.type) if cfg.policy.use_amp else nullcontext():
         info = eval_policy(
             env,
             policy,
@@ -499,4 +499,4 @@ def eval(cfg: EvalPipelineConfig):
 
 if __name__ == "__main__":
     init_logging()
-    eval()
+    eval_main()

lerobot/scripts/find_motors_bus_port.py

@@ -1,3 +1,16 @@
+# 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 os
 import time
 from pathlib import Path

lerobot/scripts/train.py

@@ -120,7 +120,7 @@ def train(cfg: TrainPipelineConfig):
         set_seed(cfg.seed)
 
     # Check device is available
-    device = get_safe_torch_device(cfg.device, log=True)
+    device = get_safe_torch_device(cfg.policy.device, log=True)
     torch.backends.cudnn.benchmark = True
     torch.backends.cuda.matmul.allow_tf32 = True
 
@@ -138,13 +138,12 @@ def train(cfg: TrainPipelineConfig):
     logging.info("Creating policy")
     policy = make_policy(
         cfg=cfg.policy,
-        device=device,
         ds_meta=dataset.meta,
     )
 
     logging.info("Creating optimizer and scheduler")
     optimizer, lr_scheduler = make_optimizer_and_scheduler(cfg, policy)
-    grad_scaler = GradScaler(device, enabled=cfg.use_amp)
+    grad_scaler = GradScaler(device.type, enabled=cfg.policy.use_amp)
 
     step = 0  # number of policy updates (forward + backward + optim)
 
@@ -218,7 +217,7 @@ def train(cfg: TrainPipelineConfig):
             cfg.optimizer.grad_clip_norm,
             grad_scaler=grad_scaler,
             lr_scheduler=lr_scheduler,
-            use_amp=cfg.use_amp,
+            use_amp=cfg.policy.use_amp,
         )
 
         # Note: eval and checkpoint happens *after* the `step`th training update has completed, so we
@@ -249,7 +248,10 @@ def train(cfg: TrainPipelineConfig):
         if cfg.env and is_eval_step:
             step_id = get_step_identifier(step, cfg.steps)
             logging.info(f"Eval policy at step {step}")
-            with torch.no_grad(), torch.autocast(device_type=device.type) if cfg.use_amp else nullcontext():
+            with (
+                torch.no_grad(),
+                torch.autocast(device_type=device.type) if cfg.policy.use_amp else nullcontext(),
+            ):
                 eval_info = eval_policy(
                     eval_env,
                     policy,

lerobot/scripts/visualize_dataset.py

@@ -265,13 +265,25 @@ def main():
         ),
     )
 
+    parser.add_argument(
+        "--tolerance-s",
+        type=float,
+        default=1e-4,
+        help=(
+            "Tolerance in seconds used to ensure data timestamps respect the dataset fps value"
+            "This is argument passed to the constructor of LeRobotDataset and maps to its tolerance_s constructor argument"
+            "If not given, defaults to 1e-4."
+        ),
+    )
+
     args = parser.parse_args()
     kwargs = vars(args)
     repo_id = kwargs.pop("repo_id")
     root = kwargs.pop("root")
+    tolerance_s = kwargs.pop("tolerance_s")
 
     logging.info("Loading dataset")
-    dataset = LeRobotDataset(repo_id, root=root)
+    dataset = LeRobotDataset(repo_id, root=root, tolerance_s=tolerance_s)
 
     visualize_dataset(dataset, **vars(args))
 

lerobot/scripts/visualize_dataset_html.py

@@ -158,7 +158,7 @@ def run_server(
             if major_version < 2:
                 return "Make sure to convert your LeRobotDataset to v2 & above."
 
-        episode_data_csv_str, columns = get_episode_data(dataset, episode_id)
+        episode_data_csv_str, columns, ignored_columns = get_episode_data(dataset, episode_id)
         dataset_info = {
             "repo_id": f"{dataset_namespace}/{dataset_name}",
             "num_samples": dataset.num_frames
@@ -194,7 +194,7 @@ def run_server(
             ]
 
             response = requests.get(
-                f"https://huggingface.co/datasets/{repo_id}/resolve/main/meta/episodes.jsonl"
+                f"https://huggingface.co/datasets/{repo_id}/resolve/main/meta/episodes.jsonl", timeout=5
             )
             response.raise_for_status()
             # Split into lines and parse each line as JSON
@@ -218,6 +218,7 @@ def run_server(
             videos_info=videos_info,
             episode_data_csv_str=episode_data_csv_str,
             columns=columns,
+            ignored_columns=ignored_columns,
         )
 
     app.run(host=host, port=port)
@@ -233,9 +234,17 @@ def get_episode_data(dataset: LeRobotDataset | IterableNamespace, episode_index)
     This file will be loaded by Dygraph javascript to plot data in real time."""
     columns = []
 
-    selected_columns = [col for col, ft in dataset.features.items() if ft["dtype"] == "float32"]
+    selected_columns = [col for col, ft in dataset.features.items() if ft["dtype"] in ["float32", "int32"]]
     selected_columns.remove("timestamp")
 
+    ignored_columns = []
+    for column_name in selected_columns:
+        shape = dataset.features[column_name]["shape"]
+        shape_dim = len(shape)
+        if shape_dim > 1:
+            selected_columns.remove(column_name)
+            ignored_columns.append(column_name)
+
     # init header of csv with state and action names
     header = ["timestamp"]
 
@@ -291,7 +300,7 @@ def get_episode_data(dataset: LeRobotDataset | IterableNamespace, episode_index)
     csv_writer.writerows(rows)
     csv_string = csv_buffer.getvalue()
 
-    return csv_string, columns
+    return csv_string, columns, ignored_columns
 
 
 def get_episode_video_paths(dataset: LeRobotDataset, ep_index: int) -> list[str]:
@@ -318,7 +327,9 @@ def get_episode_language_instruction(dataset: LeRobotDataset, ep_index: int) ->
 
 
 def get_dataset_info(repo_id: str) -> IterableNamespace:
-    response = requests.get(f"https://huggingface.co/datasets/{repo_id}/resolve/main/meta/info.json")
+    response = requests.get(
+        f"https://huggingface.co/datasets/{repo_id}/resolve/main/meta/info.json", timeout=5
+    )
     response.raise_for_status()  # Raises an HTTPError for bad responses
     dataset_info = response.json()
     dataset_info["repo_id"] = repo_id
@@ -435,15 +446,31 @@ def main():
         help="Delete the output directory if it exists already.",
     )
 
+    parser.add_argument(
+        "--tolerance-s",
+        type=float,
+        default=1e-4,
+        help=(
+            "Tolerance in seconds used to ensure data timestamps respect the dataset fps value"
+            "This is argument passed to the constructor of LeRobotDataset and maps to its tolerance_s constructor argument"
+            "If not given, defaults to 1e-4."
+        ),
+    )
+
     args = parser.parse_args()
     kwargs = vars(args)
     repo_id = kwargs.pop("repo_id")
     load_from_hf_hub = kwargs.pop("load_from_hf_hub")
     root = kwargs.pop("root")
+    tolerance_s = kwargs.pop("tolerance_s")
 
     dataset = None
     if repo_id:
-        dataset = LeRobotDataset(repo_id, root=root) if not load_from_hf_hub else get_dataset_info(repo_id)
+        dataset = (
+            LeRobotDataset(repo_id, root=root, tolerance_s=tolerance_s)
+            if not load_from_hf_hub
+            else get_dataset_info(repo_id)
+        )
 
     visualize_dataset_html(dataset, **vars(args))
 

lerobot/templates/visualize_dataset_template.html

@@ -42,22 +42,22 @@
             <ul>
                 <template x-for="episode in paginatedEpisodes" :key="episode">
                     <li class="font-mono text-sm mt-0.5">
-                        <a :href="'episode_' + episode" 
+                        <a :href="'episode_' + episode"
                            :class="{'underline': true, 'font-bold -ml-1': episode == {{ episode_id }}}"
                            x-text="'Episode ' + episode"></a>
                     </li>
                 </template>
             </ul>
-        
+
             <div class="flex items-center mt-3 text-xs" x-show="totalPages > 1">
-                <button @click="prevPage()" 
+                <button @click="prevPage()"
                         class="px-2 py-1 bg-slate-800 rounded mr-2"
                         :class="{'opacity-50 cursor-not-allowed': page === 1}"
                         :disabled="page === 1">
                     &laquo; Prev
                 </button>
                 <span class="font-mono mr-2" x-text="` ${page} / ${totalPages}`"></span>
-                <button @click="nextPage()" 
+                <button @click="nextPage()"
                         class="px-2 py-1 bg-slate-800 rounded"
                         :class="{'opacity-50 cursor-not-allowed': page === totalPages}"
                         :disabled="page === totalPages">
@@ -65,10 +65,10 @@
                 </button>
             </div>
         </div>
-        
+
         <!-- episodes menu for small screens -->
         <div class="flex overflow-x-auto md:hidden" x-data="episodePagination">
-            <button @click="prevPage()" 
+            <button @click="prevPage()"
                     class="px-2 bg-slate-800 rounded mr-2"
                     :class="{'opacity-50 cursor-not-allowed': page === 1}"
                     :disabled="page === 1">&laquo;</button>
@@ -83,7 +83,7 @@
                     </p>
                 </template>
             </div>
-            <button @click="nextPage()" 
+            <button @click="nextPage()"
                     class="px-2 bg-slate-800 rounded ml-2"
                     :class="{'opacity-50 cursor-not-allowed': page === totalPages}"
                     :disabled="page === totalPages">&raquo; </button>
@@ -224,49 +224,58 @@
                 </p>
             </div>
 
-            <table class="text-sm border-collapse border border-slate-700" x-show="currentFrameData">
-                <thead>
-                    <tr>
-                        <th></th>
-                        <template x-for="(_, colIndex) in Array.from({length: columns.length}, (_, index) => index)">
-                            <th class="border border-slate-700">
-                                <div class="flex gap-x-2 justify-between px-2">
-                                    <input type="checkbox" :checked="isColumnChecked(colIndex)"
-                                        @change="toggleColumn(colIndex)">
-                                    <p x-text="`${columns[colIndex].key}`"></p>
-                                </div>
-                            </th>
-                        </template>
-                    </tr>
-                </thead>
-                <tbody>
-                    <template x-for="(row, rowIndex) in rows">
-                        <tr class="odd:bg-gray-800 even:bg-gray-900">
-                            <td class="border border-slate-700">
-                                <div class="flex gap-x-2 max-w-64 font-semibold px-1 break-all">
-                                    <input type="checkbox" :checked="isRowChecked(rowIndex)"
-                                        @change="toggleRow(rowIndex)">
-                                </div>
-                            </td>
-                            <template x-for="(cell, colIndex) in row">
-                                <td x-show="cell" class="border border-slate-700">
-                                    <div class="flex gap-x-2 justify-between px-2" :class="{ 'hidden': cell.isNull }">
-                                        <div class="flex gap-x-2">
-                                            <input type="checkbox" x-model="cell.checked" @change="updateTableValues()">
-                                            <span x-text="`${!cell.isNull ? cell.label : null}`"></span>
-                                        </div>
-                                        <span class="w-14 text-right" x-text="`${!cell.isNull ? (typeof cell.value === 'number' ? cell.value.toFixed(2) : cell.value) : null}`"
-                                            :style="`color: ${cell.color}`"></span>
+            <div>
+                <table class="text-sm border-collapse border border-slate-700" x-show="currentFrameData">
+                    <thead>
+                        <tr>
+                            <th></th>
+                            <template x-for="(_, colIndex) in Array.from({length: columns.length}, (_, index) => index)">
+                                <th class="border border-slate-700">
+                                    <div class="flex gap-x-2 justify-between px-2">
+                                        <input type="checkbox" :checked="isColumnChecked(colIndex)"
+                                            @change="toggleColumn(colIndex)">
+                                        <p x-text="`${columns[colIndex].key}`"></p>
                                     </div>
-                                </td>
+                                </th>
                             </template>
                         </tr>
-                    </template>
-                </tbody>
-            </table>
+                    </thead>
+                    <tbody>
+                        <template x-for="(row, rowIndex) in rows">
+                            <tr class="odd:bg-gray-800 even:bg-gray-900">
+                                <td class="border border-slate-700">
+                                    <div class="flex gap-x-2 max-w-64 font-semibold px-1 break-all">
+                                        <input type="checkbox" :checked="isRowChecked(rowIndex)"
+                                            @change="toggleRow(rowIndex)">
+                                    </div>
+                                </td>
+                                <template x-for="(cell, colIndex) in row">
+                                    <td x-show="cell" class="border border-slate-700">
+                                        <div class="flex gap-x-2 justify-between px-2" :class="{ 'hidden': cell.isNull }">
+                                            <div class="flex gap-x-2">
+                                                <input type="checkbox" x-model="cell.checked" @change="updateTableValues()">
+                                                <span x-text="`${!cell.isNull ? cell.label : null}`"></span>
+                                            </div>
+                                            <span class="w-14 text-right" x-text="`${!cell.isNull ? (typeof cell.value === 'number' ? cell.value.toFixed(2) : cell.value) : null}`"
+                                                :style="`color: ${cell.color}`"></span>
+                                        </div>
+                                    </td>
+                                </template>
+                            </tr>
+                        </template>
+                    </tbody>
+                </table>
 
-            <div id="labels" class="hidden">
+                <div id="labels" class="hidden">
+                </div>
+
+                {% if ignored_columns|length > 0 %}
+                <div class="m-2 text-orange-700 max-w-96">
+                    Columns {{ ignored_columns }} are NOT shown since the visualizer currently does not support 2D or 3D data.
+                </div>
+                {% endif %}
             </div>
+
         </div>
     </div>
 
@@ -476,7 +485,7 @@
                 episodes: {{ episodes }},
                 pageSize: 100,
                 page: 1,
-                
+
                 init() {
                     // Find which page contains the current episode_id
                     const currentEpisodeId = {{ episode_id }};
@@ -485,23 +494,23 @@
                         this.page = Math.floor(episodeIndex / this.pageSize) + 1;
                     }
                 },
-                
+
                 get totalPages() {
                     return Math.ceil(this.episodes.length / this.pageSize);
                 },
-                
+
                 get paginatedEpisodes() {
                     const start = (this.page - 1) * this.pageSize;
                     const end = start + this.pageSize;
                     return this.episodes.slice(start, end);
                 },
-                
+
                 nextPage() {
                     if (this.page < this.totalPages) {
                         this.page++;
                     }
                 },
-                
+
                 prevPage() {
                     if (this.page > 1) {
                         this.page--;
@@ -515,7 +524,7 @@
         window.addEventListener('keydown', (e) => {
             // Use the space bar to play and pause, instead of default action (e.g. scrolling)
             const { keyCode, key } = e;
-            
+
             if (keyCode === 32 || key === ' ') {
                 e.preventDefault();
                 const btnPause = document.querySelector('[x-ref="btnPause"]');