TOREMOVE: isolate aloha on __init__ to see if it creates the bug

Co-authored-by: Remi <remi.cadene@huggingface.co>
Co-authored-by: Remi Cadene <re.cadene@gmail.com>

.cache/calibration/aloha_default/left_follower.json

@@ -0,0 +1,68 @@
+{
+    "homing_offset": [
+        2048,
+        3072,
+        3072,
+        -1024,
+        -1024,
+        2048,
+        -2048,
+        2048,
+        -2048
+    ],
+    "drive_mode": [
+        1,
+        1,
+        1,
+        0,
+        0,
+        1,
+        0,
+        1,
+        0
+    ],
+    "start_pos": [
+        2015,
+        3058,
+        3061,
+        1071,
+        1071,
+        2035,
+        2152,
+        2029,
+        2499
+    ],
+    "end_pos": [
+        -1008,
+        -1963,
+        -1966,
+        2141,
+        2143,
+        -971,
+        3043,
+        -1077,
+        3144
+    ],
+    "calib_mode": [
+        "DEGREE",
+        "DEGREE",
+        "DEGREE",
+        "DEGREE",
+        "DEGREE",
+        "DEGREE",
+        "DEGREE",
+        "DEGREE",
+        "LINEAR"
+    ],
+    "motor_names": [
+        "waist",
+        "shoulder",
+        "shoulder_shadow",
+        "elbow",
+        "elbow_shadow",
+        "forearm_roll",
+        "wrist_angle",
+        "wrist_rotate",
+        "gripper"
+    ]
+}

.cache/calibration/aloha_default/left_leader.json

@@ -0,0 +1,68 @@
+{
+    "homing_offset": [
+        2048,
+        3072,
+        3072,
+        -1024,
+        -1024,
+        2048,
+        -2048,
+        2048,
+        -1024
+    ],
+    "drive_mode": [
+        1,
+        1,
+        1,
+        0,
+        0,
+        1,
+        0,
+        1,
+        0
+    ],
+    "start_pos": [
+        2035,
+        3024,
+        3019,
+        979,
+        981,
+        1982,
+        2166,
+        2124,
+        1968
+    ],
+    "end_pos": [
+        -990,
+        -2017,
+        -2015,
+        2078,
+        2076,
+        -1030,
+        3117,
+        -1016,
+        2556
+    ],
+    "calib_mode": [
+        "DEGREE",
+        "DEGREE",
+        "DEGREE",
+        "DEGREE",
+        "DEGREE",
+        "DEGREE",
+        "DEGREE",
+        "DEGREE",
+        "LINEAR"
+    ],
+    "motor_names": [
+        "waist",
+        "shoulder",
+        "shoulder_shadow",
+        "elbow",
+        "elbow_shadow",
+        "forearm_roll",
+        "wrist_angle",
+        "wrist_rotate",
+        "gripper"
+    ]
+}

.cache/calibration/aloha_default/right_follower.json

@@ -0,0 +1,68 @@
+{
+    "homing_offset": [
+        2048,
+        3072,
+        3072,
+        -1024,
+        -1024,
+        2048,
+        -2048,
+        2048,
+        -2048
+    ],
+    "drive_mode": [
+        1,
+        1,
+        1,
+        0,
+        0,
+        1,
+        0,
+        1,
+        0
+    ],
+    "start_pos": [
+        2056,
+        2895,
+        2896,
+        1191,
+        1190,
+        2018,
+        2051,
+        2056,
+        2509
+    ],
+    "end_pos": [
+        -1040,
+        -2004,
+        -2006,
+        2126,
+        2127,
+        -1010,
+        3050,
+        -1117,
+        3143
+    ],
+    "calib_mode": [
+        "DEGREE",
+        "DEGREE",
+        "DEGREE",
+        "DEGREE",
+        "DEGREE",
+        "DEGREE",
+        "DEGREE",
+        "DEGREE",
+        "LINEAR"
+    ],
+    "motor_names": [
+        "waist",
+        "shoulder",
+        "shoulder_shadow",
+        "elbow",
+        "elbow_shadow",
+        "forearm_roll",
+        "wrist_angle",
+        "wrist_rotate",
+        "gripper"
+    ]
+}

.cache/calibration/aloha_default/right_leader.json

@@ -0,0 +1,68 @@
+{
+    "homing_offset": [
+        2048,
+        3072,
+        3072,
+        -1024,
+        -1024,
+        2048,
+        -2048,
+        2048,
+        -2048
+    ],
+    "drive_mode": [
+        1,
+        1,
+        1,
+        0,
+        0,
+        1,
+        0,
+        1,
+        0
+    ],
+    "start_pos": [
+        2068,
+        3034,
+        3030,
+        1038,
+        1041,
+        1991,
+        1948,
+        2090,
+        1985
+    ],
+    "end_pos": [
+        -1025,
+        -2014,
+        -2015,
+        2058,
+        2060,
+        -955,
+        3091,
+        -940,
+        2576
+    ],
+    "calib_mode": [
+        "DEGREE",
+        "DEGREE",
+        "DEGREE",
+        "DEGREE",
+        "DEGREE",
+        "DEGREE",
+        "DEGREE",
+        "DEGREE",
+        "LINEAR"
+    ],
+    "motor_names": [
+        "waist",
+        "shoulder",
+        "shoulder_shadow",
+        "elbow",
+        "elbow_shadow",
+        "forearm_roll",
+        "wrist_angle",
+        "wrist_rotate",
+        "gripper"
+    ]
+}

.gitattributes

@@ -3,4 +3,4 @@
 *.safetensors filter=lfs diff=lfs merge=lfs -text
 *.mp4 filter=lfs diff=lfs merge=lfs -text
 *.arrow filter=lfs diff=lfs merge=lfs -text
-*.json filter=lfs diff=lfs merge=lfs -text
+*.json !text !filter !merge !diff

.github/workflows/test.yml

@@ -11,6 +11,7 @@ on:
       - ".github/**"
       - "poetry.lock"
       - "Makefile"
+      - ".cache/**"
   push:
     branches:
       - main
@@ -21,6 +22,7 @@ on:
       - ".github/**"
       - "poetry.lock"
       - "Makefile"
+      - ".cache/**"
 
 jobs:
   pytest:
@@ -35,7 +37,10 @@ jobs:
           lfs: true  # Ensure LFS files are pulled
 
       - name: Install apt dependencies
-        run: sudo apt-get update && sudo apt-get install -y libegl1-mesa-dev ffmpeg
+      # portaudio19-dev is needed to install pyaudio
+        run: |
+          sudo apt-get update && \
+          sudo apt-get install -y libegl1-mesa-dev ffmpeg portaudio19-dev
 
       - name: Install poetry
         run: |
@@ -60,7 +65,6 @@ jobs:
             -W ignore::UserWarning:gymnasium.utils.env_checker:247 \
             && rm -rf tests/outputs outputs
 
-
   pytest-minimal:
     name: Pytest (minimal install)
     runs-on: ubuntu-latest
@@ -110,7 +114,10 @@ jobs:
           lfs: true  # Ensure LFS files are pulled
 
       - name: Install apt dependencies
-        run: sudo apt-get update && sudo apt-get install -y libegl1-mesa-dev
+      # portaudio19-dev is needed to install pyaudio
+        run: |
+          sudo apt-get update && \
+          sudo apt-get install -y libegl1-mesa-dev portaudio19-dev
 
       - name: Install poetry
         run: |

.gitignore

@@ -66,7 +66,6 @@ htmlcov/
 .nox/
 .coverage
 .coverage.*
-.cache
 nosetests.xml
 coverage.xml
 *.cover
@@ -74,6 +73,11 @@ coverage.xml
 .hypothesis/
 .pytest_cache/
 
+# Ignore .cache except calibration
+.cache/*
+!.cache/calibration/
+!.cache/calibration/**
+
 # Translations
 *.mo
 *.pot

CONTRIBUTING.md

@@ -20,7 +20,7 @@ Some of the ways you can contribute to 🤗 LeRobot:
 * Contributing to the examples or to the documentation.
 * Submitting issues related to bugs or desired new features.
 
-Following the guides below, feel free to open issues and PRs and to coordinate your efforts with the community on our [Discord Channel](https://discord.gg/VjFz58wn3R). For specific inquiries, reach out to [Remi Cadene](remi.cadene@huggingface.co).
+Following the guides below, feel free to open issues and PRs and to coordinate your efforts with the community on our [Discord Channel](https://discord.gg/VjFz58wn3R). For specific inquiries, reach out to [Remi Cadene](mailto:remi.cadene@huggingface.co).
 
 If you are not sure how to contribute or want to know the next features we working on, look on this project page: [LeRobot TODO](https://github.com/orgs/huggingface/projects/46)
 

docker/lerobot-cpu/Dockerfile

@@ -22,7 +22,7 @@ RUN echo "source /opt/venv/bin/activate" >> /root/.bashrc
 COPY . /lerobot
 WORKDIR /lerobot
 RUN pip install --upgrade --no-cache-dir pip
-RUN pip install --no-cache-dir ".[test, aloha, xarm, pusht, koch]" \
+RUN pip install --no-cache-dir ".[test, aloha, xarm, pusht, dynamixel]" \
     --extra-index-url https://download.pytorch.org/whl/cpu
 
 # Set EGL as the rendering backend for MuJoCo

docker/lerobot-gpu/Dockerfile

@@ -24,7 +24,7 @@ RUN echo "source /opt/venv/bin/activate" >> /root/.bashrc
 COPY . /lerobot
 WORKDIR /lerobot
 RUN pip install --upgrade --no-cache-dir pip
-RUN pip install --no-cache-dir ".[test, aloha, xarm, pusht, koch]"
+RUN pip install --no-cache-dir ".[test, aloha, xarm, pusht, dynamixel]"
 
 # Set EGL as the rendering backend for MuJoCo
 ENV MUJOCO_GL="egl"

examples/4_train_policy_with_script.md

@@ -196,7 +196,7 @@ These logs will also be saved in wandb if `wandb.enable` is set to `true`. Here
 - `success`: average success rate of eval episodes. Reward and success are usually different except for the sparsing reward setting, where reward=1 only when the task is completed successfully.
 - `eval_s`: time to evaluate the policy in the environment, in second.
 - `updt_s`: time to update the network parameters, in second.
-- `data_s`: time to load a batch of data, in second. 
+- `data_s`: time to load a batch of data, in second.
 
 Some metrics are useful for initial performance profiling. For example, if you find the current GPU utilization is low via the `nvidia-smi` command and `data_s` sometimes is too high, you may need to modify batch size or number of dataloading workers to accelerate dataloading. We also recommend [pytorch profiler](https://github.com/huggingface/lerobot?tab=readme-ov-file#improve-your-code-with-profiling) for detailed performance probing.
 

examples/7_get_started_with_real_robot.md

@@ -11,7 +11,7 @@ This tutorial will guide you through the process of setting up and training a ne
 
 By following these steps, you'll be able to replicate tasks like picking up a Lego block and placing it in a bin with a high success rate, as demonstrated in [this video](https://x.com/RemiCadene/status/1814680760592572934).
 
-Although this tutorial is general and can be easily adapted to various types of robots by changing the configuration, it is specifically based on the [Koch v1.1](https://github.com/jess-moss/koch-v1-1), an affordable robot. The Koch v1.1 consists of a leader arm and a follower arm, each with 6 motors. It can work with one or several cameras to record the scene, which serve as visual sensors for the robot.
+This tutorial is specifically made for the affordable [Koch v1.1](https://github.com/jess-moss/koch-v1-1) robot, but it contains additional information to be easily adapted to various types of robots like [Aloha bimanual robot](aloha-2.github.io) by changing some configurations. The Koch v1.1 consists of a leader arm and a follower arm, each with 6 motors. It can work with one or several cameras to record the scene, which serve as visual sensors for the robot.
 
 During the data collection phase, you will control the follower arm by moving the leader arm. This process is known as "teleoperation." This technique is used to collect robot trajectories. Afterward, you'll train a neural network to imitate these trajectories and deploy the network to enable your robot to operate autonomously.
 
@@ -29,16 +29,23 @@ For a visual walkthrough of the assembly process, you can refer to [this video t
 
 ## 2. Configure motors, calibrate arms, teleoperate your Koch v1.1
 
-First, install the additional dependencies required for Koch v1.1 by running one of the following commands.
+First, install the additional dependencies required for robots built with dynamixel motors like Koch v1.1 by running one of the following commands.
 
 Using `pip`:
 ```bash
-pip install -e ".[koch]"
+pip install -e ".[dynamixel]"
 ```
 
 Or using `poetry`:
 ```bash
-poetry install --sync --extras "koch"
+poetry install --sync --extras "dynamixel"
+```
+
+/!\ For Linux only, ffmpeg and opencv requires conda install for now. Run this exact sequence of commands:
+```bash
+conda install -c conda-forge ffmpeg
+pip uninstall opencv-python
+conda install -c conda-forge "opencv>=4.10.0"
 ```
 
 You are now ready to plug the 5V power supply to the motor bus of the leader arm (the smaller one) since all its motors only require 5V.
@@ -147,6 +154,7 @@ follower_arm = DynamixelMotorsBus(
 Next, update the port values in the YAML configuration file for the Koch robot at [`lerobot/configs/robot/koch.yaml`](../lerobot/configs/robot/koch.yaml) with the ports you've identified:
 ```yaml
 [...]
+robot_type: koch
 leader_arms:
   main:
     _target_: lerobot.common.robot_devices.motors.dynamixel.DynamixelMotorsBus
@@ -174,6 +182,8 @@ follower_arms:
 [...]
 ```
 
+Don't forget to set `robot_type: aloha` if you follow this tutorial with [Aloha bimanual robot](aloha-2.github.io) instead of Koch v1.1
+
 This configuration file is used to instantiate your robot across all scripts. We'll cover how this works later on.
 
 **Connect and Configure your Motors**
@@ -298,32 +308,37 @@ Alternatively, you can unplug the power cord, which will automatically disable t
 
 */!\ Warning*: These motors tend to overheat, especially under torque or if left plugged in for too long. Unplug after use.
 
-### b. Teleoperate your Koch v1.1 with KochRobot
+### b. Teleoperate your Koch v1.1 with ManipulatorRobot
 
-**Instantiate the KochRobot**
+**Instantiate the ManipulatorRobot**
 
-Before you can teleoperate your robot, you need to instantiate the  [`KochRobot`](../lerobot/common/robot_devices/robots/koch.py) using the previously defined `leader_arm` and `follower_arm`.
+Before you can teleoperate your robot, you need to instantiate the  [`ManipulatorRobot`](../lerobot/common/robot_devices/robots/manipulator.py) using the previously defined `leader_arm` and `follower_arm`.
 
-For the Koch robot, we only have one leader, so we refer to it as `"main"` and define it as `leader_arms={"main": leader_arm}`. We do the same for the follower arm. For other robots (like the Aloha), which may have two pairs of leader and follower arms, you would define them like this: `leader_arms={"left": left_leader_arm, "right": right_leader_arm},`. Same thing for the follower arms.
+For the Koch v1.1 robot, we only have one leader, so we refer to it as `"main"` and define it as `leader_arms={"main": leader_arm}`. We do the same for the follower arm. For other robots (like the Aloha), which may have two pairs of leader and follower arms, you would define them like this: `leader_arms={"left": left_leader_arm, "right": right_leader_arm},`. Same thing for the follower arms.
 
-You also need to provide a path to a calibration file, such as  `calibration_path=".cache/calibration/koch.pkl"`. More on this in the next section.
+You also need to provide a path to a calibration directory, such as  `calibration_dir=".cache/calibration/koch"`. More on this in the next section.
 
-Run the following code to instantiate your Koch robot:
+Run the following code to instantiate your manipulator robot:
 ```python
-from lerobot.common.robot_devices.robots.koch import KochRobot
+from lerobot.common.robot_devices.robots.manipulator import ManipulatorRobot
 
-robot = KochRobot(
+robot = ManipulatorRobot(
+    robot_type="koch",
     leader_arms={"main": leader_arm},
     follower_arms={"main": follower_arm},
-    calibration_path=".cache/calibration/koch.pkl",
+    calibration_dir=".cache/calibration/koch",
 )
 ```
 
-**Calibrate and Connect the KochRobot**
+The `robot_type="koch"` is used to set the associated settings and calibration process. For instance, we activate the torque of the gripper of the leader Koch v1.1 arm and position it at a 40 degree angle to use it as a trigger.
 
-Next, you'll need to calibrate your robot to ensure that the leader and follower arms have the same position values when they are in the same physical position. This calibration is essential because it allows a neural network trained on one Koch robot to work on another.
+For the [Aloha bimanual robot](https://aloha-2.github.io), we would use `robot_type="aloha"` to set different settings such as a secondary ID for shadow joints (shoulder, elbow). Specific to Aloha, LeRobot comes with default calibration files stored in in `.cache/calibration/aloha_default`. Assuming the motors have been properly assembled, no manual calibration step is expected. If you need to run manual calibration, simply update `calibration_dir` to `.cache/calibration/aloha`.
 
-When you connect your robot for the first time, the [`KochRobot`](../lerobot/common/robot_devices/robots/koch.py) will detect if the calibration file is missing and trigger the calibration procedure. During this process, you will be guided to move each arm to three different positions.
+**Calibrate and Connect the ManipulatorRobot**
+
+Next, you'll need to calibrate your Koch robot to ensure that the leader and follower arms have the same position values when they are in the same physical position. This calibration is essential because it allows a neural network trained on one Koch robot to work on another.
+
+When you connect your robot for the first time, the [`ManipulatorRobot`](../lerobot/common/robot_devices/robots/manipulator.py) will detect if the calibration file is missing and trigger the calibration procedure. During this process, you will be guided to move each arm to three different positions.
 
 Here are the positions you'll move the follower arm to:
 
@@ -354,27 +369,26 @@ The output will look like this:
 ```
 Connecting main follower arm
 Connecting main leader arm
-Missing calibration file '.cache/calibration/koch.pkl'. Starting calibration procedure.
-
-Running calibration of main follower...
 
+Missing calibration file '.cache/calibration/koch/main_follower.json'
+Running calibration of koch main follower...
 Move arm to zero position
 [...]
 Move arm to rotated position
 [...]
 Move arm to rest position
 [...]
+Calibration is done! Saving calibration file '.cache/calibration/koch/main_follower.json'
 
-Running calibration of main leader...
-
+Missing calibration file '.cache/calibration/koch/main_leader.json'
+Running calibration of koch main leader...
 Move arm to zero position
 [...]
 Move arm to rotated position
 [...]
 Move arm to rest position
 [...]
-
-Calibration is done! Saving calibration file '.cache/calibration/koch.pkl'
+Calibration is done! Saving calibration file '.cache/calibration/koch/main_leader.json'
 ```
 
 *Verifying Calibration*
@@ -414,7 +428,7 @@ for _ in tqdm.tqdm(range(seconds*frequency)):
 
 *Using `teleop_step` for Teleoperation*
 
-Alternatively, you can teleoperate the robot using the `teleop_step` method from [`KochRobot`](../lerobot/common/robot_devices/robots/koch.py).
+Alternatively, you can teleoperate the robot using the `teleop_step` method from [`ManipulatorRobot`](../lerobot/common/robot_devices/robots/manipulator.py).
 
 Run this code to teleoperate:
 ```python
@@ -607,10 +621,10 @@ Additionaly, you can set up your robot to work with your cameras.
 
 Modify the following Python code with the appropriate camera names and configurations:
 ```python
-robot = KochRobot(
+robot = ManipulatorRobot(
     leader_arms={"main": leader_arm},
     follower_arms={"main": follower_arm},
-    calibration_path=".cache/calibration/koch.pkl",
+    calibration_dir=".cache/calibration/koch",
     cameras={
         "laptop": OpenCVCamera(0, fps=30, width=640, height=480),
         "phone": OpenCVCamera(1, fps=30, width=640, height=480),
@@ -925,7 +939,7 @@ huggingface-cli upload ${HF_USER}/act_koch_test_${CKPT} \
 
 ## 5. Evaluate your policy
 
-Now that you have a policy checkpoint, you can easily control your robot with it using methods from [`KochRobot`](../lerobot/common/robot_devices/robots/koch.py) and the policy.
+Now that you have a policy checkpoint, you can easily control your robot with it using methods from [`ManipulatorRobot`](../lerobot/common/robot_devices/robots/manipulator.py) and the policy.
 
 Try this code for running inference for 60 seconds at 30 fps:
 ```python

examples/8_use_stretch.md

@@ -0,0 +1,158 @@
+This tutorial explains how to use [Stretch 3](https://hello-robot.com/stretch-3-product) with LeRobot.
+
+## Setup
+
+Familiarize yourself with Stretch by following its [tutorials](https://docs.hello-robot.com/0.3/getting_started/hello_robot/) (recommended).
+
+To use LeRobot on Stretch, 3 options are available:
+- [tethered setup](https://docs.hello-robot.com/0.3/getting_started/connecting_to_stretch/#tethered-setup)
+- [untethered setup](https://docs.hello-robot.com/0.3/getting_started/connecting_to_stretch/#untethered-setup)
+- ssh directly into Stretch (you will first need to install and configure openssh-server on stretch using one of the two above setups)
+
+
+## Install LeRobot
+
+On Stretch's CLI, follow these steps:
+
+1. [Install Miniconda](https://docs.anaconda.com/miniconda/#quick-command-line-install):
+```bash
+mkdir -p ~/miniconda3
+wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda3/miniconda.sh
+bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3
+rm ~/miniconda3/miniconda.sh
+~/miniconda3/bin/conda init bash
+```
+
+2. Comment out these lines in `~/.profile` (this can mess up paths used by conda and ~/.local/bin should already be in your PATH)
+```
+# set PATH so it includes user's private bin if it exists
+if [ -d "$HOME/.local/bin" ] ; then
+    PATH="$HOME/.local/bin:$PATH"
+fi
+```
+
+3. Restart shell or `source ~/.bashrc`
+
+4. Create and activate a fresh conda environment for lerobot
+```bash
+conda create -y -n lerobot python=3.10 && conda activate lerobot
+```
+
+5. Clone LeRobot:
+```bash
+git clone https://github.com/huggingface/lerobot.git ~/lerobot
+```
+
+6. Install LeRobot with stretch dependencies:
+```bash
+cd ~/lerobot && pip install -e ".[stretch]"
+```
+
+> **Note:** If you get this message, you can ignore it: `ERROR: pip's dependency resolver does not currently take into account all the packages that are installed.`
+
+And install extra dependencies for recording datasets on Linux:
+```bash
+conda install -y -c conda-forge ffmpeg
+pip uninstall -y opencv-python
+conda install -y -c conda-forge "opencv>=4.10.0"
+```
+
+7. Run a [system check](https://docs.hello-robot.com/0.3/getting_started/stretch_hardware_overview/#system-check) to make sure your robot is ready:
+```bash
+stretch_system_check.py
+```
+
+> **Note:** You may need to free the "robot process" after booting Stretch by running `stretch_free_robot_process.py`. For more info this Stretch's [doc](https://docs.hello-robot.com/0.3/getting_started/stretch_hardware_overview/#turning-off-gamepad-teleoperation).
+
+You should get something like this:
+```bash
+For use with S T R E T C H (R) from Hello Robot Inc.
+---------------------------------------------------------------------
+
+Model = Stretch 3
+Tool = DexWrist 3 w/ Gripper
+Serial Number = stretch-se3-3054
+
+---- Checking Hardware ----
+[Pass] Comms are ready
+[Pass] Actuators are ready
+[Warn] Sensors not ready (IMU AZ = -10.19 out of range -10.1 to -9.5)
+[Pass] Battery voltage is 13.6 V
+
+---- Checking Software ----
+[Pass] Ubuntu 22.04 is ready
+[Pass] All APT pkgs are setup correctly
+[Pass] Firmware is up-to-date
+[Pass] Python pkgs are up-to-date
+[Pass] ROS2 Humble is ready
+```
+
+## Teleoperate, record a dataset and run a policy
+
+**Calibrate (Optional)**
+Before operating Stretch, you need to [home](https://docs.hello-robot.com/0.3/getting_started/stretch_hardware_overview/#homing) it first. Be mindful about giving Stretch some space as this procedure will move the robot's arm and gripper. Now run this command:
+```bash
+python lerobot/scripts/control_robot.py calibrate \
+    --robot-path lerobot/configs/robot/stretch.yaml
+```
+This is equivalent to running `stretch_robot_home.py`
+
+> **Note:** If you run any of the LeRobot scripts below and Stretch is not poperly homed, it will automatically home/calibrate first.
+
+**Teleoperate**
+Before trying teleoperation, you need activate the gamepad controller by pressing the middle button. For more info, see Stretch's [doc](https://docs.hello-robot.com/0.3/getting_started/hello_robot/#gamepad-teleoperation).
+
+Now try out teleoperation (see above documentation to learn about the gamepad controls):
+```bash
+python lerobot/scripts/control_robot.py teleoperate \
+    --robot-path lerobot/configs/robot/stretch.yaml
+```
+This is essentially the same as running `stretch_gamepad_teleop.py`
+
+**Record a dataset**
+Once you're familiar with the gamepad controls and after a bit of practice, you can try to record your first dataset with Stretch.
+
+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
+```
+
+Store your Hugging Face repository name in a variable to run these commands:
+```bash
+HF_USER=$(huggingface-cli whoami | head -n 1)
+echo $HF_USER
+```
+
+Record one episode:
+```bash
+python lerobot/scripts/control_robot.py record \
+    --robot-path lerobot/configs/robot/stretch.yaml \
+    --fps 20 \
+    --root data \
+    --repo-id ${HF_USER}/stretch_test \
+    --tags stretch tutorial \
+    --warmup-time-s 3 \
+    --episode-time-s 40 \
+    --reset-time-s 10 \
+    --num-episodes 1 \
+    --push-to-hub 0
+```
+
+> **Note:** If you're using ssh to connect to Stretch and run this script, you won't be able to visualize its cameras feed (though they will still be recording). To see the cameras stream, use [tethered](https://docs.hello-robot.com/0.3/getting_started/connecting_to_stretch/#tethered-setup) or [untethered setup](https://docs.hello-robot.com/0.3/getting_started/connecting_to_stretch/#untethered-setup).
+
+**Replay an episode**
+Now try to replay this episode (make sure the robot's initial position is the same):
+```bash
+python lerobot/scripts/control_robot.py replay \
+    --robot-path lerobot/configs/robot/stretch.yaml \
+    --fps 20 \
+    --root data \
+    --repo-id ${HF_USER}/stretch_test \
+    --episode 0
+```
+
+Follow [previous tutorial](https://github.com/huggingface/lerobot/blob/main/examples/7_get_started_with_real_robot.md#4-train-a-policy-on-your-data) to train a policy on your data and run inference on your robot. You will need to adapt the code for Stretch.
+
+> TODO(rcadene, aliberts): Add already setup environment and policy yaml configuration files
+
+If you need help, please reach out on Discord in the channel `#stretch3-mobile-arm`.

lerobot/__init__.py

@@ -27,6 +27,9 @@ Example:
         print(lerobot.available_real_world_datasets)
         print(lerobot.available_policies)
         print(lerobot.available_policies_per_env)
+        print(lerobot.available_robots)
+        print(lerobot.available_cameras)
+        print(lerobot.available_motors)
     ```
 
 When implementing a new dataset loadable with LeRobotDataset follow these steps:
@@ -182,7 +185,7 @@ available_datasets = list(
     itertools.chain(*available_datasets_per_env.values(), available_real_world_datasets)
 )
 
-# lists all available policies from `lerobot/common/policies` by their class attribute: `name`.
+# lists all available policies from `lerobot/common/policies`
 available_policies = [
     "act",
     "diffusion",
@@ -190,12 +193,32 @@ available_policies = [
     "vqbet",
 ]
 
+# lists all available robots from `lerobot/common/robot_devices/robots`
+available_robots = [
+    # "koch",
+    # "koch_bimanual",
+    "aloha",
+]
+
+# lists all available cameras from `lerobot/common/robot_devices/cameras`
+available_cameras = [
+    "opencv",
+    "intelrealsense",
+]
+
+# lists all available motors from `lerobot/common/robot_devices/motors`
+available_motors = [
+    "dynamixel",
+]
+
 # keys and values refer to yaml files
 available_policies_per_env = {
     "aloha": ["act"],
     "pusht": ["diffusion", "vqbet"],
     "xarm": ["tdmpc"],
-    "dora_aloha_real": ["act_real"],
+    "koch_real": ["act_koch_real"],
+    "aloha_real": ["act_aloha_real"],
+    "dora_aloha_real": ["act_aloha_real"],
 }
 
 env_task_pairs = [(env, task) for env, tasks in available_tasks_per_env.items() for task in tasks]

lerobot/common/datasets/compute_stats.py

@@ -68,7 +68,7 @@ def get_stats_einops_patterns(dataset, num_workers=0):
     return stats_patterns
 
 
-def compute_stats(dataset, batch_size=32, num_workers=16, max_num_samples=None):
+def compute_stats(dataset, batch_size=8, num_workers=8, max_num_samples=None):
     """Compute mean/std and min/max statistics of all data keys in a LeRobotDataset."""
     if max_num_samples is None:
         max_num_samples = len(dataset)

lerobot/common/datasets/push_dataset_to_hub/openx/configs.yaml

@@ -8,7 +8,7 @@ OPENX_DATASET_CONFIGS:
       - base_pose_tool_reached
       - gripper_closed
     fps: 3
-  
+
   kuka:
     image_obs_keys:
       - image
@@ -18,7 +18,7 @@ OPENX_DATASET_CONFIGS:
       - clip_function_input/base_pose_tool_reached
       - gripper_closed
     fps: 10
-  
+
   bridge_openx:
     image_obs_keys:
       - image
@@ -28,7 +28,7 @@ OPENX_DATASET_CONFIGS:
       - EEF_state
       - gripper_state
     fps: 5
-  
+
   taco_play:
     image_obs_keys:
       - rgb_static
@@ -40,7 +40,7 @@ OPENX_DATASET_CONFIGS:
       - state_eef
       - state_gripper
     fps: 15
-  
+
   jaco_play:
     image_obs_keys:
       - image
@@ -51,7 +51,7 @@ OPENX_DATASET_CONFIGS:
       - state_eef
       - state_gripper
     fps: 10
-  
+
   berkeley_cable_routing:
     image_obs_keys:
       - image
@@ -72,7 +72,7 @@ OPENX_DATASET_CONFIGS:
     state_obs_keys:
       - null
     fps: 10
-  
+
   nyu_door_opening_surprising_effectiveness:
     image_obs_keys:
       - image
@@ -221,7 +221,7 @@ OPENX_DATASET_CONFIGS:
     image_obs_keys:
       - highres_image
     depth_obs_keys:
-      - null  
+      - null
     state_obs_keys:
       - null
     fps: 10
@@ -234,7 +234,7 @@ OPENX_DATASET_CONFIGS:
     state_obs_keys:
       - joint_state
     fps: 2
-  
+
   ucsd_pick_and_place_dataset_converted_externally_to_rlds:
     image_obs_keys:
       - image
@@ -244,7 +244,7 @@ OPENX_DATASET_CONFIGS:
       - eef_state
       - gripper_state
     fps: 3
-  
+
   spoc:
     image_obs_keys:
       - image
@@ -254,7 +254,7 @@ OPENX_DATASET_CONFIGS:
     state_obs_keys:
       - null
     fps: 3
-  
+
   austin_sailor_dataset_converted_externally_to_rlds:
     image_obs_keys:
       - image
@@ -264,7 +264,7 @@ OPENX_DATASET_CONFIGS:
     state_obs_keys:
       - state
     fps: 20
-  
+
   austin_sirius_dataset_converted_externally_to_rlds:
     image_obs_keys:
       - image
@@ -274,7 +274,7 @@ OPENX_DATASET_CONFIGS:
     state_obs_keys:
       - state
     fps: 20
-  
+
   bc_z:
     image_obs_keys:
       - image
@@ -285,7 +285,7 @@ OPENX_DATASET_CONFIGS:
       - present/axis_angle
       - present/sensed_close
     fps: 10
-  
+
   utokyo_pr2_opening_fridge_converted_externally_to_rlds:
     image_obs_keys:
       - image
@@ -295,7 +295,7 @@ OPENX_DATASET_CONFIGS:
       - eef_state
       - gripper_state
     fps: 10
-  
+
   utokyo_pr2_tabletop_manipulation_converted_externally_to_rlds:
     image_obs_keys:
       - image
@@ -305,7 +305,7 @@ OPENX_DATASET_CONFIGS:
       - eef_state
       - gripper_state
     fps: 10
-  
+
   utokyo_xarm_pick_and_place_converted_externally_to_rlds:
     image_obs_keys:
       - image
@@ -316,7 +316,7 @@ OPENX_DATASET_CONFIGS:
     state_obs_keys:
       - end_effector_pose
     fps: 10
-  
+
   utokyo_xarm_bimanual_converted_externally_to_rlds:
     image_obs_keys:
       - image
@@ -325,7 +325,7 @@ OPENX_DATASET_CONFIGS:
     state_obs_keys:
       - pose_r
     fps: 10
-  
+
   robo_net:
     image_obs_keys:
       - image
@@ -336,7 +336,7 @@ OPENX_DATASET_CONFIGS:
       - eef_state
       - gripper_state
     fps: 1
-  
+
   robo_set:
     image_obs_keys:
       - image_left
@@ -348,7 +348,7 @@ OPENX_DATASET_CONFIGS:
       - state
       - state_velocity
     fps: 5
-  
+
   berkeley_mvp_converted_externally_to_rlds:
     image_obs_keys:
       - hand_image
@@ -359,7 +359,7 @@ OPENX_DATASET_CONFIGS:
       - pose
       - joint_pos
     fps: 5
-  
+
   berkeley_rpt_converted_externally_to_rlds:
     image_obs_keys:
       - hand_image
@@ -369,7 +369,7 @@ OPENX_DATASET_CONFIGS:
       - joint_pos
       - gripper
     fps: 30
-  
+
   kaist_nonprehensile_converted_externally_to_rlds:
     image_obs_keys:
       - image
@@ -378,7 +378,7 @@ OPENX_DATASET_CONFIGS:
     state_obs_keys:
       - state
     fps: 10
-  
+
   stanford_mask_vit_converted_externally_to_rlds:
     image_obs_keys:
       - image
@@ -387,7 +387,7 @@ OPENX_DATASET_CONFIGS:
     state_obs_keys:
       - eef_state
       - gripper_state
-  
+
   tokyo_u_lsmo_converted_externally_to_rlds:
     image_obs_keys:
       - image
@@ -397,7 +397,7 @@ OPENX_DATASET_CONFIGS:
       - eef_state
       - gripper_state
     fps: 10
-  
+
   dlr_sara_pour_converted_externally_to_rlds:
     image_obs_keys:
       - image
@@ -406,16 +406,16 @@ OPENX_DATASET_CONFIGS:
     state_obs_keys:
       - state
     fps: 10
-  
+
   dlr_sara_grid_clamp_converted_externally_to_rlds:
     image_obs_keys:
       - image
     depth_obs_keys:
       - null
     state_obs_keys:
-      - state  
+      - state
     fps: 10
-  
+
   dlr_edan_shared_control_converted_externally_to_rlds:
     image_obs_keys:
       - image
@@ -424,7 +424,7 @@ OPENX_DATASET_CONFIGS:
     state_obs_keys:
       - state
     fps: 5
-  
+
   asu_table_top_converted_externally_to_rlds:
     image_obs_keys:
       - image
@@ -478,7 +478,7 @@ OPENX_DATASET_CONFIGS:
     state_obs_keys:
       - null
     fps: 1
-  
+
   utaustin_mutex:
     image_obs_keys:
       - image
@@ -488,7 +488,7 @@ OPENX_DATASET_CONFIGS:
     state_obs_keys:
       - state
     fps: 20
-  
+
   berkeley_fanuc_manipulation:
     image_obs_keys:
       - image
@@ -499,7 +499,7 @@ OPENX_DATASET_CONFIGS:
       - joint_state
       - gripper_state
     fps: 10
-  
+
   cmu_playing_with_food:
     image_obs_keys:
       - image
@@ -509,7 +509,7 @@ OPENX_DATASET_CONFIGS:
     state_obs_keys:
       - state
     fps: 10
-  
+
   cmu_play_fusion:
     image_obs_keys:
       - image
@@ -518,7 +518,7 @@ OPENX_DATASET_CONFIGS:
     state_obs_keys:
       - state
     fps: 5
-  
+
   cmu_stretch:
     image_obs_keys:
       - image
@@ -528,7 +528,7 @@ OPENX_DATASET_CONFIGS:
       - eef_state
       - gripper_state
     fps: 10
-  
+
   berkeley_gnm_recon:
     image_obs_keys:
       - image
@@ -539,7 +539,7 @@ OPENX_DATASET_CONFIGS:
       - position
       - yaw
     fps: 3
- 
+
   berkeley_gnm_cory_hall:
     image_obs_keys:
       - image
@@ -550,7 +550,7 @@ OPENX_DATASET_CONFIGS:
       - position
       - yaw
     fps: 5
- 
+
   berkeley_gnm_sac_son:
     image_obs_keys:
       - image
@@ -561,7 +561,7 @@ OPENX_DATASET_CONFIGS:
       - position
       - yaw
     fps: 10
-  
+
   droid:
     image_obs_keys:
       - exterior_image_1_left
@@ -572,7 +572,7 @@ OPENX_DATASET_CONFIGS:
     state_obs_keys:
       - proprio
     fps: 15
-  
+
   droid_100:
     image_obs_keys:
       - exterior_image_1_left
@@ -583,7 +583,7 @@ OPENX_DATASET_CONFIGS:
     state_obs_keys:
       - proprio
     fps: 15
-  
+
   fmb:
     image_obs_keys:
       - image_side_1
@@ -598,7 +598,7 @@ OPENX_DATASET_CONFIGS:
     state_obs_keys:
       - proprio
     fps: 10
-  
+
   dobbe:
     image_obs_keys:
       - wrist_image
@@ -607,7 +607,7 @@ OPENX_DATASET_CONFIGS:
     state_obs_keys:
       - proprio
     fps: 3.75
-  
+
   usc_cloth_sim_converted_externally_to_rlds:
     image_obs_keys:
       - image
@@ -616,7 +616,7 @@ OPENX_DATASET_CONFIGS:
     state_obs_keys:
       - null
     fps: 10
-  
+
   plex_robosuite:
     image_obs_keys:
       - image
@@ -626,7 +626,7 @@ OPENX_DATASET_CONFIGS:
     state_obs_keys:
       - state
     fps: 20
-  
+
   conq_hose_manipulation:
     image_obs_keys:
       - frontleft_fisheye_image
@@ -637,4 +637,3 @@ OPENX_DATASET_CONFIGS:
     state_obs_keys:
       - state
     fps: 30
-  

lerobot/common/datasets/push_dataset_to_hub/openx_rlds_format.py

@@ -52,7 +52,7 @@ from lerobot.common.datasets.utils import (
 )
 from lerobot.common.datasets.video_utils import VideoFrame, encode_video_frames
 
-with open("lerobot/common/datasets/push_dataset_to_hub/openx/configs.yaml", "r") as f:
+with open("lerobot/common/datasets/push_dataset_to_hub/openx/configs.yaml") as f:
     _openx_list = yaml.safe_load(f)
 
 OPENX_DATASET_CONFIGS = _openx_list["OPENX_DATASET_CONFIGS"]

lerobot/common/datasets/utils.py

@@ -32,7 +32,7 @@ DATASET_CARD_TEMPLATE = """
 ---
 # Metadata will go there
 ---
-This dataset was created using [🤗 LeRobot](https://github.com/huggingface/lerobot).
+This dataset was created using [LeRobot](https://github.com/huggingface/lerobot).
 
 """
 

lerobot/common/envs/utils.py

@@ -39,7 +39,7 @@ def preprocess_observation(observations: dict[str, np.ndarray]) -> dict[str, Ten
 
             # sanity check that images are channel last
             _, h, w, c = img.shape
-            assert c < h and c < w, f"expect channel first images, but instead {img.shape}"
+            assert c < h and c < w, f"expect channel last images, but instead got {img.shape=}"
 
             # sanity check that images are uint8
             assert img.dtype == torch.uint8, f"expect torch.uint8, but instead {img.dtype=}"

lerobot/common/policies/act/modeling_act.py

@@ -296,7 +296,7 @@ class ACT(nn.Module):
         self.use_images = any(k.startswith("observation.image") for k in config.input_shapes)
         self.use_env_state = "observation.environment_state" in config.input_shapes
         if self.config.use_vae:
-            self.vae_encoder = ACTEncoder(config)
+            self.vae_encoder = ACTEncoder(config, is_vae_encoder=True)
             self.vae_encoder_cls_embed = nn.Embedding(1, config.dim_model)
             # Projection layer for joint-space configuration to hidden dimension.
             if self.use_robot_state:
@@ -521,9 +521,11 @@ class ACT(nn.Module):
 class ACTEncoder(nn.Module):
     """Convenience module for running multiple encoder layers, maybe followed by normalization."""
 
-    def __init__(self, config: ACTConfig):
+    def __init__(self, config: ACTConfig, is_vae_encoder: bool = False):
         super().__init__()
-        self.layers = nn.ModuleList([ACTEncoderLayer(config) for _ in range(config.n_encoder_layers)])
+        self.is_vae_encoder = is_vae_encoder
+        num_layers = config.n_vae_encoder_layers if self.is_vae_encoder else config.n_encoder_layers
+        self.layers = nn.ModuleList([ACTEncoderLayer(config) for _ in range(num_layers)])
         self.norm = nn.LayerNorm(config.dim_model) if config.pre_norm else nn.Identity()
 
     def forward(

lerobot/common/policies/diffusion/configuration_diffusion.py

@@ -196,3 +196,12 @@ class DiffusionConfig:
                 f"`noise_scheduler_type` must be one of {supported_noise_schedulers}. "
                 f"Got {self.noise_scheduler_type}."
             )
+
+        # Check that the horizon size and U-Net downsampling is compatible.
+        # U-Net downsamples by 2 with each stage.
+        downsampling_factor = 2 ** len(self.down_dims)
+        if self.horizon % downsampling_factor != 0:
+            raise ValueError(
+                "The horizon should be an integer multiple of the downsampling factor (which is determined "
+                f"by `len(down_dims)`). Got {self.horizon=} and {self.down_dims=}"
+            )

lerobot/common/policies/factory.py

@@ -51,11 +51,6 @@ def get_policy_and_config_classes(name: str) -> tuple[Policy, object]:
         from lerobot.common.policies.tdmpc.modeling_tdmpc import TDMPCPolicy
 
         return TDMPCPolicy, TDMPCConfig
-    elif name == "tdmpc2":
-        from lerobot.common.policies.tdmpc2.configuration_tdmpc2 import TDMPC2Config
-        from lerobot.common.policies.tdmpc2.modeling_tdmpc2 import TDMPC2Policy
-
-        return TDMPC2Policy, TDMPC2Config
     elif name == "diffusion":
         from lerobot.common.policies.diffusion.configuration_diffusion import DiffusionConfig
         from lerobot.common.policies.diffusion.modeling_diffusion import DiffusionPolicy

lerobot/common/policies/tdmpc2/configuration_tdmpc2.py

@@ -1,217 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2024 Nicklas Hansen, Xiaolong Wang, Hao Su,
-# and 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
-
-
-@dataclass
-class TDMPC2Config:
-    """Configuration class for TDMPC2Policy.
-
-    Defaults are configured for training with xarm_lift_medium_replay providing proprioceptive and single
-    camera observations.
-
-    The parameters you will most likely need to change are the ones which depend on the environment / sensors.
-    Those are: `input_shapes`, `output_shapes`, and perhaps `max_random_shift_ratio`.
-
-    Args:
-        n_action_repeats: The number of times to repeat the action returned by the planning. (hint: Google
-            action repeats in Q-learning or ask your favorite chatbot)
-        horizon: Horizon for model predictive control.
-        n_action_steps: Number of action steps to take from the plan given by model predictive control. This
-            is an alternative to using action repeats. If this is set to more than 1, then we require
-            `n_action_repeats == 1`, `use_mpc == True` and `n_action_steps <= horizon`. Note that this
-            approach of using multiple steps from the plan is not in the original implementation.
-        input_shapes: A dictionary defining the shapes of the input data for the policy. The key represents
-            the input data name, and the value is a list indicating the dimensions of the corresponding data.
-            For example, "observation.image" refers to an input from a camera with dimensions [3, 96, 96],
-            indicating it has three color channels and 96x96 resolution. Importantly, `input_shapes` doesn't
-            include batch dimension or temporal dimension.
-        output_shapes: A dictionary defining the shapes of the output data for the policy. The key represents
-            the output data name, and the value is a list indicating the dimensions of the corresponding data.
-            For example, "action" refers to an output shape of [14], indicating 14-dimensional actions.
-            Importantly, `output_shapes` doesn't include batch dimension or temporal dimension.
-        input_normalization_modes: A dictionary with key representing the modality (e.g. "observation.state"),
-            and the value specifies the normalization mode to apply. The two available modes are "mean_std"
-            which subtracts the mean and divides by the standard deviation and "min_max" which rescale in a
-            [-1, 1] range. Note that here this defaults to None meaning inputs are not normalized. This is to
-            match the original implementation.
-        output_normalization_modes: Similar dictionary as `normalize_input_modes`, but to unnormalize to the
-            original scale. Note that this is also used for normalizing the training targets. NOTE: Clipping
-            to [-1, +1] is used during MPPI/CEM. Therefore, it is recommended that you stick with "min_max"
-            normalization mode here.
-        image_encoder_hidden_dim: Number of channels for the convolutional layers used for image encoding.
-        state_encoder_hidden_dim: Hidden dimension for MLP used for state vector encoding.
-        latent_dim: Observation's latent embedding dimension.
-        q_ensemble_size: Number of Q function estimators to use in an ensemble for uncertainty estimation.
-        mlp_dim: Hidden dimension of MLPs used for modelling the dynamics encoder, reward function, policy
-            (π), Q ensemble, and V.
-        discount: Discount factor (γ) to use for the reinforcement learning formalism.
-        use_mpc: Whether to use model predictive control. The alternative is to just sample the policy model
-            (π) for each step.
-        cem_iterations: Number of iterations for the MPPI/CEM loop in MPC.
-        max_std: Maximum standard deviation for actions sampled from the gaussian PDF in CEM.
-        min_std: Minimum standard deviation for noise applied to actions sampled from the policy model (π).
-            Doubles up as the minimum standard deviation for actions sampled from the gaussian PDF in CEM.
-        n_gaussian_samples: Number of samples to draw from the gaussian distribution every CEM iteration. Must
-            be non-zero.
-        n_pi_samples: Number of samples to draw from the policy / world model rollout every CEM iteration. Can
-            be zero.
-        uncertainty_regularizer_coeff: Coefficient for the uncertainty regularization used when estimating
-            trajectory values (this is the λ coeffiecient in eqn 4 of FOWM).
-        n_elites: The number of elite samples to use for updating the gaussian parameters every CEM iteration.
-        elite_weighting_temperature: The temperature to use for softmax weighting (by trajectory value) of the
-            elites, when updating the gaussian parameters for CEM.
-        gaussian_mean_momentum: Momentum (α) used for EMA updates of the mean parameter μ of the gaussian
-            parameters optimized in CEM. Updates are calculated as μ⁻ ← αμ⁻ + (1-α)μ.
-        max_random_shift_ratio: Maximum random shift (as a proportion of the image size) to apply to the
-            image(s) (in units of pixels) for training-time augmentation. If set to 0, no such augmentation
-            is applied. Note that the input images are assumed to be square for this augmentation.
-        reward_coeff: Loss weighting coefficient for the reward regression loss.
-        expectile_weight: Weighting (τ) used in expectile regression for the state value function (V).
-            v_pred < v_target is weighted by τ and v_pred >= v_target is weighted by (1-τ). τ is expected to
-            be in [0, 1]. Setting τ closer to 1 results in a more "optimistic" V. This is sensible to do
-            because v_target is obtained by evaluating the learned state-action value functions (Q) with
-            in-sample actions that may not be always optimal.
-        value_coeff: Loss weighting coefficient for both the state-action value (Q) TD loss, and the state
-            value (V) expectile regression loss.
-        consistency_coeff: Loss weighting coefficient for the consistency loss.
-        advantage_scaling: A factor by which the advantages are scaled prior to exponentiation for advantage
-            weighted regression of the policy (π) estimator parameters. Note that the exponentiated advantages
-            are clamped at 100.0.
-        pi_coeff: Loss weighting coefficient for the action regression loss.
-        temporal_decay_coeff: Exponential decay coefficient for decaying the loss coefficient for future time-
-            steps. Hint: each loss computation involves `horizon` steps worth of actions starting from the
-            current time step.
-        target_model_momentum: Momentum (α) used for EMA updates of the target models. Updates are calculated
-            as ϕ ← αϕ + (1-α)θ where ϕ are the parameters of the target model and θ are the parameters of the
-            model being trained.
-    """
-
-    num_bins = 101
-    vmin = -10
-    vmax = +10
-    rho: float = 0.5
-    tau: float = 0.01
-    simnorm_dim: int = 8
-
-    # Input / output structure.
-    n_action_repeats: int = 2
-    horizon: int = 5
-    n_action_steps: int = 1
-
-    input_shapes: dict[str, list[int]] = field(
-        default_factory=lambda: {
-            "observation.image": [3, 64, 64],
-            "observation.state": [4],
-        }
-    )
-    output_shapes: dict[str, list[int]] = field(
-        default_factory=lambda: {
-            "action": [4],
-        }
-    )
-
-    # Normalization / Unnormalization
-    input_normalization_modes: dict[str, str] | None = None
-    output_normalization_modes: dict[str, str] = field(
-        default_factory=lambda: {"action": "min_max"},
-    )
-
-    # Architecture / modeling.
-    # Neural networks.
-    image_encoder_hidden_dim: int = 32
-    state_encoder_hidden_dim: int = 256
-    latent_dim: int = 8 #50
-    q_ensemble_size: int = 5
-    mlp_dim: int = 512
-    # Reinforcement learning.
-    discount: float = 0.9
-    lr: float = 3e-4
-    enc_lr_scale: float = 0.3
-
-    num_q: int = 5
-    dropout: float = 0.01
-
-    num_channels = 32
-    num_enc_layers = 2
-    enc_dim = 256
-
-    # Inference.
-    use_mpc: bool = True
-    cem_iterations: int = 6
-    max_std: float = 2.0
-    min_std: float = 0.05
-    n_gaussian_samples: int = 512
-    n_pi_samples: int = 51
-    uncertainty_regularizer_coeff: float = 1.0
-    n_elites: int = 50
-    elite_weighting_temperature: float = 0.5
-    gaussian_mean_momentum: float = 0.1
-
-    # Training and loss computation.
-    grad_clip_norm: float = 20
-
-    max_random_shift_ratio: float = 0.0476
-    # Loss coefficients.
-    consistency_coef: float = 20
-    entropy_coef: float = 1e-4
-
-    reward_coef: float = 0.1
-    expectile_weight: float = 0.9
-    value_coef: float = 0.1
-    consistency_coeff: float = 20.0
-    advantage_scaling: float = 3.0
-    pi_coeff: float = 0.5
-    temporal_decay_coeff: float = 0.5
-    # Target model.
-    target_model_momentum: float = 0.995
-
-    def __post_init__(self):
-        """Input validation (not exhaustive)."""
-        # There should only be one image key.
-        image_keys = {k for k in self.input_shapes if k.startswith("observation.image")}
-        if len(image_keys) > 1:
-            raise ValueError(
-                f"{self.__class__.__name__} handles at most one image for now. Got image keys {image_keys}."
-            )
-        if len(image_keys) > 0:
-            image_key = next(iter(image_keys))
-            if self.input_shapes[image_key][-2] != self.input_shapes[image_key][-1]:
-                # TODO(alexander-soare): This limitation is solely because of code in the random shift
-                # augmentation. It should be able to be removed.
-                raise ValueError(
-                    f"Only square images are handled now. Got image shape {self.input_shapes[image_key]}."
-                )
-        if self.n_gaussian_samples <= 0:
-            raise ValueError(
-                f"The number of guassian samples for CEM should be non-zero. Got `{self.n_gaussian_samples=}`"
-            )
-        if self.output_normalization_modes != {"action": "min_max"}:
-            raise ValueError(
-                "TD-MPC assumes the action space dimensions to all be in [-1, 1]. Therefore it is strongly "
-                f"advised that you stick with the default. See {self.__class__.__name__} docstring for more "
-                "information."
-            )
-        if self.n_action_steps > 1:
-            if self.n_action_repeats != 1:
-                raise ValueError(
-                    "If `n_action_steps > 1`, `n_action_repeats` must be left to its default value of 1."
-                )
-            if not self.use_mpc:
-                raise ValueError("If `n_action_steps > 1`, `use_mpc` must be set to `True`.")
-            if self.n_action_steps > self.horizon:
-                raise ValueError("`n_action_steps` must be less than or equal to `horizon`.")

lerobot/common/policies/tdmpc2/modeling_tdmpc2.py

@@ -1,727 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2024 Nicklas Hansen, Xiaolong Wang, Hao Su,
-# and 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.
-"""Implementation of Finetuning Offline World Models in the Real World.
-
-The comments in this code may sometimes refer to these references:
-    TD-MPC paper: Temporal Difference Learning for Model Predictive Control (https://arxiv.org/abs/2203.04955)
-    FOWM paper: Finetuning Offline World Models in the Real World (https://arxiv.org/abs/2310.16029)
-"""
-
-# ruff: noqa: N806
-
-import logging
-from collections import deque
-from copy import deepcopy
-from functools import partial
-from typing import Callable
-
-import einops
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F  # noqa: N812
-from huggingface_hub import PyTorchModelHubMixin
-from torch import Tensor
-
-import lerobot.common.policies.tdmpc2.tdmpc2_utils as utils
-from lerobot.common.policies.normalize import Normalize, Unnormalize
-from lerobot.common.policies.tdmpc2.configuration_tdmpc2 import TDMPC2Config
-from lerobot.common.policies.utils import get_device_from_parameters, populate_queues
-
-
-class TDMPC2Policy(nn.Module, PyTorchModelHubMixin):
-    """Implementation of TD-MPC2 learning + inference.
-
-    Please note several warnings for this policy.
-        - We have NOT checked that training on LeRobot reproduces SOTA results. This is a TODO.
-    """
-
-    name = "tdmpc2"
-
-    def __init__(
-        self, config: TDMPC2Config | None = None, dataset_stats: dict[str, dict[str, Tensor]] | None = None
-    ):
-        """
-        Args:
-            config: Policy configuration class instance or None, in which case the default instantiation of
-                the configuration class is used.
-            dataset_stats: Dataset statistics to be used for normalization. If not passed here, it is expected
-                that they will be passed with a call to `load_state_dict` before the policy is used.
-        """
-        super().__init__()
-        logging.warning(
-            """
-            Please note several warnings for this policy.
-            - We have NOT checked that training on LeRobot reproduces SOTA results. This is a TODO.
-            """
-        )
-
-        if config is None:
-            config = TDMPC2Config()
-        self.config = config
-        self.model = TDMPC2TOLD(config)
-
-        if config.input_normalization_modes is not None:
-            self.normalize_inputs = Normalize(
-                config.input_shapes, config.input_normalization_modes, dataset_stats
-            )
-        else:
-            self.normalize_inputs = nn.Identity()
-        self.normalize_targets = Normalize(
-            config.output_shapes, config.output_normalization_modes, dataset_stats
-        )
-        self.unnormalize_outputs = Unnormalize(
-            config.output_shapes, config.output_normalization_modes, dataset_stats
-        )
-
-        image_keys = [k for k in config.input_shapes if k.startswith("observation.image")]
-        # Note: This check is covered in the post-init of the config but have a sanity check just in case.
-        self._use_image = False
-        self._use_env_state = False
-        if len(image_keys) > 0:
-            assert len(image_keys) == 1
-            self._use_image = True
-            self.input_image_key = image_keys[0]
-        if "observation.environment_state" in config.input_shapes:
-            self._use_env_state = True
-
-        self.scale = utils.RunningScale(self.config)
-
-        self.queue_keys = None
-
-        self.reset()
-
-    def reset(self):
-        """
-        Clear observation and action queues. Clear previous means for warm starting of MPPI/CEM. Should be
-        called on `env.reset()`
-        """
-        self._queues = {
-            "observation.state": deque(maxlen=1),
-            "action": deque(maxlen=max(self.config.n_action_steps, self.config.n_action_repeats)),
-        }
-        if self._use_image:
-            self._queues["observation.image"] = deque(maxlen=1)
-        if self._use_env_state:
-            self._queues["observation.environment_state"] = deque(maxlen=1)
-        # Previous mean obtained from the cross-entropy method (CEM) used during MPC. It is used to warm start
-        # CEM for the next step.
-        self._prev_mean: torch.Tensor | None = None
-
-    @torch.no_grad()
-    def select_action(self, batch: dict[str, Tensor]) -> Tensor:
-        """Select a single action given environment observations."""
-        batch = self.normalize_inputs(batch)
-        if self._use_image:
-            batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
-            batch["observation.image"] = batch[self.input_image_key]
-
-        self._queues = populate_queues(self._queues, batch)
-
-        if self.queue_keys is None:
-            self.queue_keys = [k for k in batch if k in self._queues]
-
-        # When the action queue is depleted, populate it again by querying the policy.
-        if len(self._queues["action"]) == 0:
-            batch = {key: torch.stack(list(self._queues[key]), dim=1) for key in self.queue_keys}
-
-            # Remove the time dimensions as it is not handled yet.
-            for key in batch:
-                assert batch[key].shape[1] == 1
-                batch[key] = batch[key][:, 0]
-
-            # NOTE: Order of observations matters here.
-            encode_keys = []
-            if self._use_image:
-                encode_keys.append("observation.image")
-            if self._use_env_state:
-                encode_keys.append("observation.environment_state")
-            encode_keys.append("observation.state")
-
-            z = self.model.encode({k: batch[k] for k in encode_keys})
-
-            if self.config.use_mpc:  # noqa: SIM108
-                actions = self.plan(z)  # (horizon, batch, action_dim)
-            else:
-                # Plan with the policy (π) alone. This always returns one action so unsqueeze to get a
-                # sequence dimension like in the MPC branch.
-                actions = self.model.pi(z).unsqueeze(0)
-
-            actions = torch.clamp(actions, -1, +1)
-
-            actions = self.unnormalize_outputs({"action": actions})["action"]
-
-            if self.config.n_action_repeats > 1:
-                for _ in range(self.config.n_action_repeats):
-                    self._queues["action"].append(actions[0])
-            else:
-                # Action queue is (n_action_steps, batch_size, action_dim), so we transpose the action.
-                self._queues["action"].extend(
-                    actions[: self.config.n_action_steps]
-                )  # TDMPC2 does it use n_action_steps?
-
-        action = self._queues["action"].popleft()
-        return action
-
-    @torch.no_grad()
-    def plan(self, z: Tensor) -> Tensor:
-        """Plan sequence of actions using TD-MPC inference.
-
-        Args:
-            z: (batch, latent_dim,) tensor for the initial state.
-        Returns:
-            (horizon, batch, action_dim,) tensor for the planned trajectory of actions.
-        """
-        device = get_device_from_parameters(self)
-
-        batch_size = z.shape[0]
-
-        # Sample Nπ trajectories from the policy.
-        pi_actions = torch.empty(
-            self.config.horizon,
-            self.config.n_pi_samples,
-            batch_size,
-            self.config.output_shapes["action"][0],
-            device=device,
-        )
-        if self.config.n_pi_samples > 0:
-            _z = einops.repeat(z, "b d -> n b d", n=self.config.n_pi_samples)
-            for t in range(self.config.horizon):
-                # Note: Adding a small amount of noise here doesn't hurt during inference and may even be
-                # helpful for CEM.
-                pi_actions[t] = self.model.pi_action(_z)
-                _z = self.model.latent_dynamics(_z, pi_actions[t])
-
-        # In the CEM loop we will need this for a call to estimate_value with the gaussian sampled
-        # trajectories.
-        z = einops.repeat(z, "b d -> n b d", n=self.config.n_gaussian_samples + self.config.n_pi_samples)
-
-        # Model Predictive Path Integral (MPPI) with the cross-entropy method (CEM) as the optimization
-        # algorithm.
-        # The initial mean and standard deviation for the cross-entropy method (CEM).
-        mean = torch.zeros(
-            self.config.horizon, batch_size, self.config.output_shapes["action"][0], device=device
-        )
-        # Maybe warm start CEM with the mean from the previous step.
-        if self._prev_mean is not None:
-            mean[:-1] = self._prev_mean[1:]
-        std = self.config.max_std * torch.ones_like(mean)
-
-        for _ in range(self.config.cem_iterations):
-            # Randomly sample action trajectories for the gaussian distribution.
-            std_normal_noise = torch.randn(
-                self.config.horizon,
-                self.config.n_gaussian_samples,
-                batch_size,
-                self.config.output_shapes["action"][0],
-                device=std.device,
-            )
-            gaussian_actions = torch.clamp(mean.unsqueeze(1) + std.unsqueeze(1) * std_normal_noise, -1, 1)
-
-            # Compute elite actions.
-            actions = torch.cat([gaussian_actions, pi_actions], dim=1)
-            value = self.estimate_value(z, actions).nan_to_num_(0).squeeze(-1)
-            elite_idxs = torch.topk(value, self.config.n_elites, dim=0).indices  # (n_elites, batch)
-            # from IPython import embed; embed()
-            elite_value = value.take_along_dim(elite_idxs, dim=0)  # (n_elites, batch)
-            # (horizon, n_elites, batch, action_dim)
-            elite_actions = actions.take_along_dim(einops.rearrange(elite_idxs, "n b -> 1 n b 1"), dim=1)
-
-            # Update gaussian PDF parameters to be the (weighted) mean and standard deviation of the elites.
-            max_value = elite_value.max(0, keepdim=True)[0]  # (1, batch)
-            # The weighting is a softmax over trajectory values. Note that this is not the same as the usage
-            # of Ω in eqn 4 of the TD-MPC paper. Instead it is the normalized version of it: s = Ω/ΣΩ. This
-            # makes the equations: μ = Σ(s⋅Γ), σ = Σ(s⋅(Γ-μ)²).
-            score = torch.exp(self.config.elite_weighting_temperature * (elite_value - max_value))
-            score /= score.sum(axis=0, keepdim=True)
-            # (horizon, batch, action_dim)
-            _mean = torch.sum(einops.rearrange(score, "n b -> n b 1") * elite_actions, dim=1)
-            _std = torch.sqrt(
-                torch.sum(
-                    einops.rearrange(score, "n b -> n b 1")
-                    * (elite_actions - einops.rearrange(_mean, "h b d -> h 1 b d")) ** 2,
-                    dim=1,
-                )
-            )
-            # Update mean with an exponential moving average, and std with a direct replacement.
-            mean = (
-                self.config.gaussian_mean_momentum * mean + (1 - self.config.gaussian_mean_momentum) * _mean
-            )
-            std = _std.clamp_(self.config.min_std, self.config.max_std)
-
-        # Keep track of the mean for warm-starting subsequent steps.
-        self._prev_mean = mean
-
-        # Randomly select one of the elite actions from the last iteration of MPPI/CEM using the softmax
-        # scores from the last iteration.
-        actions = elite_actions[:, torch.multinomial(score.T, 1).squeeze(), torch.arange(batch_size)]
-
-        return actions
-
-    @torch.no_grad()
-    def estimate_value(self, z, actions):
-        """Estimate value of a trajectory starting at latent state z and executing given actions."""
-        G, discount = 0, 1
-        for t in range(self.config.horizon):
-            reward = utils.two_hot_inv(self.model._reward(torch.cat([z, actions[t]], dim=-1)), self.config)
-            z = self.model.next(z, actions[t])
-            G += discount * reward
-            discount *= self.config.discount
-        return G + discount * self.model.Qs(z, self.model.pi(z)[1], return_type="avg")
-
-    def forward(self, batch: dict[str, Tensor]) -> dict[str, Tensor | float]:
-        """Run the batch through the model and compute the loss.
-
-        Returns a dictionary with loss as a tensor, and other information as native floats.
-        """
-        device = get_device_from_parameters(self)
-
-        batch = self.normalize_inputs(batch)
-        if self._use_image:
-            batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
-            batch["observation.image"] = batch[self.input_image_key]
-        batch = self.normalize_targets(batch)
-
-        # (b, t) -> (t, b)
-        for key in batch:
-            if batch[key].ndim > 1:
-                batch[key] = batch[key].transpose(1, 0)
-
-        action = batch["action"]  # (t, b, action_dim)
-        reward = batch["next.reward"]  # (t, b)
-        reward = reward.unsqueeze(-1)  # (t, b, 1)
-        observations = {k: v for k, v in batch.items() if k.startswith("observation.")}
-
-        # Apply random image augmentations.
-        if self._use_image and self.config.max_random_shift_ratio > 0:
-            observations["observation.image"] = flatten_forward_unflatten(
-                partial(random_shifts_aug, max_random_shift_ratio=self.config.max_random_shift_ratio),
-                observations["observation.image"],
-            )
-
-        # Get the current observation for predicting trajectories, and all future observations for use in
-        # the latent consistency loss and TD loss.
-        current_observation, next_observations = {}, {}
-        for k in observations:
-            current_observation[k] = observations[k][0]
-            next_observations[k] = observations[k][1:]
-        horizon, batch_size = next_observations[
-            "observation.image" if self._use_image else "observation.environment_state"
-        ].shape[:2]
-
-        # Compute targets
-        with torch.no_grad():
-            next_z = self.model.encode(next_observations)
-            curr_z = self.model.encode(current_observation).unsqueeze(
-                0
-            )  # TODO: not necessary to do the whole thing
-            # get the next targets # _td_target in the original code
-            pi = self.model.pi(next_z)[1]
-            discount = self.config.discount
-
-            td_targets = reward + discount * self.model.Qs(next_z, pi, return_type="min", target=True)
-
-        #self.model.train()
-
-        # Latent rollout
-        zs = torch.empty(self.config.horizon + 1, batch_size, self.config.latent_dim, device=device)
-        zs[0] = z = curr_z[0]
-        consistency_loss = 0
-        for t in range(self.config.horizon):
-            x = torch.cat([z, action[t]], dim=-1)
-            z = self.model._dynamics(x)
-            consistency_loss += F.mse_loss(z, next_z[t]) * self.config.rho**t
-            zs[t + 1] = z
-
-        # Predictions
-        _zs = zs[:-1]
-        qs = self.model.Qs(_zs, action, return_type="all")
-        reward_preds = self.model._reward(torch.cat([_zs, action], dim=-1))
-
-        # Compute losses
-        reward_loss, value_loss = 0, 0
-        for t in range(self.config.horizon):
-            reward_loss += utils.soft_ce(reward_preds[t], reward[t], self.config).mean() * self.config.rho**t
-            for q in range(self.config.num_q):
-                value_loss += utils.soft_ce(qs[q][t], td_targets[t], self.config).mean() * self.config.rho**t
-        consistency_loss *= 1 / self.config.horizon
-        reward_loss *= 1 / self.config.horizon
-        value_loss *= 1 / (self.config.horizon * self.config.num_q)
-
-        ############################ deviation from NHansen
-        # total_loss = (
-        #    self.config.consistency_coef * consistency_loss
-        #    + self.config.reward_coef * reward_loss
-        #    + self.config.value_coef * value_loss
-        # )
-
-        # Update model########################
-        # total_loss.backward()
-        # grad_norm = torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.config.grad_clip_norm)
-        # self.optim.step()
-        ##########################################
-
-        # Deviation from Hansen, since the optimizer step is called in train.py
-        # Update the policy using a sequence of latent states.
-        zs_for_pi = zs.detach()
-        # self.pi_optim.zero_grad(set_to_none=True)
-        self.model.track_q_grad(False)
-        _, pis, log_pis, _ = self.model.pi(zs_for_pi)
-        qs = self.model.Qs(zs_for_pi, pis, return_type="avg")
-        self.scale.update(qs[0])
-        qs = self.scale(qs)
-
-        # Loss is a weighted sum of Q-values
-        rho = torch.pow(self.config.rho, torch.arange(len(qs), device=device))
-        pi_loss = ((self.config.entropy_coef * log_pis - qs).mean(dim=(1, 2)) * rho).mean()
-        # pi_loss.backward()
-        # torch.nn.utils.clip_grad_norm_(self.model._pi.parameters(), self.config.grad_clip_norm)
-        # self.pi_optim.step()
-
-        # self.model.track_q_grad(True)
-
-        # pi_loss = pi_loss.item()
-
-        loss = (
-            self.config.consistency_coef * consistency_loss
-            + self.config.reward_coef * reward_loss
-            + self.config.value_coef * value_loss
-            + self.config.pi_coeff * pi_loss
-        )
-
-        # Update target Q-functions
-        # """
-        # Soft-update target Q-networks using Polyak averaging.
-        # """
-        # with torch.no_grad():
-        #    for p, p_target in zip(self.model._Qs.parameters(), self.model._target_Qs.parameters()):
-        #        p_target.data.lerp_(p.data, self.config.tau)
-
-        # Return training statistics
-        self.model.eval()
-        info = {
-            "loss": loss,
-            "consistency_loss": consistency_loss.mean().item(),
-            "reward_loss": reward_loss.mean().item(),
-            "value_loss": value_loss.mean().item(),
-            "pi_loss": pi_loss.item(),
-            "pi_scale": self.scale.value,
-        }
-
-        # Undo (b, t) -> (t, b).
-        for key in batch:
-            if batch[key].ndim > 1:
-                batch[key] = batch[key].transpose(1, 0)
-
-        return info
-
-    def update(self):
-        """Soft-update target Q-networks using Polyak averaging."""
-        with torch.no_grad():
-            for p, p_target in zip(
-                self.model._Qs.parameters(), self.model._target_Qs.parameters(), strict=False
-            ):
-                p_target.data.lerp_(p.data, self.config.tau)
-
-
-class TDMPC2TOLD(nn.Module):
-    """Task-Oriented Latent Dynamics (TOLD) model used in TD-MPC2."""
-
-    def __init__(self, config: TDMPC2Config):
-        super().__init__()
-        self.config = config
-
-        self.config.bin_size = (config.vmax - config.vmin) / (
-            config.num_bins - 1
-        )  # Bin size for discrete regression
-
-        action_dim = config.output_shapes["action"][0]
-
-        self._encoder = TDMPC2ObservationEncoder(config)
-        self._dynamics = utils.mlp(
-            config.latent_dim + action_dim,
-            2 * [config.mlp_dim],
-            config.latent_dim,
-            act=utils.SimNorm(config),
-        )
-        self._reward = utils.mlp(
-            config.latent_dim + action_dim, 2 * [config.mlp_dim], max(config.num_bins, 1)
-        )
-        self._pi = utils.mlp(config.latent_dim, 2 * [config.mlp_dim], 2 * action_dim)
-        self._Qs = utils.Ensemble(
-            [
-                utils.mlp(
-                    config.latent_dim + action_dim,
-                    2 * [config.mlp_dim],
-                    max(config.num_bins, 1),
-                    dropout=config.dropout,
-                )
-                for _ in range(config.num_q)
-            ]
-        )
-
-        self.apply(self.weight_init)
-        for p in [self._reward[-1].weight, self._Qs.params[-2]]:
-            p.data.fill_(0)
-
-        self._target_Qs = deepcopy(self._Qs).requires_grad_(False)
-        log_std_min, log_std_max = -10, 2  # TODO: add to config
-        self.log_std_min = torch.tensor(log_std_min)
-        self.log_std_dif = torch.tensor(log_std_max) - self.log_std_min
-
-    def track_q_grad(self, mode=True):
-        """
-        Enables/disables gradient tracking of Q-networks.
-        Avoids unnecessary computation during policy optimization.
-        This method also enables/disables gradients for task embeddings.
-        """
-        for p in self._Qs.parameters():
-            p.requires_grad_(mode)
-
-    def weight_init(self, m):  # lifted from Nicklas' code
-        """Custom weight initialization for TD-MPC2."""
-        if isinstance(m, nn.Linear):
-            nn.init.trunc_normal_(m.weight, std=0.02)
-            if m.bias is not None:
-                nn.init.constant_(m.bias, 0)
-        elif isinstance(m, nn.Embedding):
-            nn.init.uniform_(m.weight, -0.02, 0.02)
-        elif isinstance(m, nn.ParameterList):
-            for i, p in enumerate(m):
-                if p.dim() == 3:  # Linear
-                    nn.init.trunc_normal_(p, std=0.02)  # Weight
-                    nn.init.constant_(m[i + 1], 0)  # Bias
-
-    def encode(self, obs: dict[str, Tensor]) -> Tensor:
-        """Encodes an observation into its latent representation."""
-        # from IPython import embed; embed()
-        # print(obs["observation.state"].shape, obs["observation.image"].shape)
-        return self._encoder(obs)
-
-    def latent_dynamics_and_reward(self, z: Tensor, a: Tensor) -> tuple[Tensor, Tensor]:
-        """Predict the next state's latent representation and the reward given a current latent and action.
-
-        Args:
-            z: (*, latent_dim) tensor for the current state's latent representation.
-            a: (*, action_dim) tensor for the action to be applied.
-        Returns:
-            A tuple containing:
-                - (*, latent_dim) tensor for the next state's latent representation.
-                - (*,) tensor for the estimated reward.
-        """
-        x = torch.cat([z, a], dim=-1)
-        r = self._reward(x)
-        r = utils.two_hot_inv(r, self.config).squeeze(-1)
-        # from IPython import embed; embed()
-        return self._dynamics(x), r
-
-    def latent_dynamics(self, z: Tensor, a: Tensor) -> Tensor:
-        """Predict the next state's latent representation given a current latent and action.
-
-        Args:
-            z: (*, latent_dim) tensor for the current state's latent representation.
-            a: (*, action_dim) tensor for the action to be applied.
-        Returns:
-            (*, latent_dim) tensor for the next state's latent representation.
-        """
-        x = torch.cat([z, a], dim=-1)
-        return self._dynamics(x)
-
-    def next(self, z: Tensor, a: Tensor) -> Tensor:
-        return self.latent_dynamics(z, a)  # just a wrapper
-
-    def pi(self, z):  # lifted from Nicklas' code
-        """
-        Samples an action from the policy prior.
-        The policy prior is a Gaussian distribution with
-        mean and (log) std predicted by a neural network.
-        """
-        # Gaussian policy prior
-        mu, log_std = self._pi(z).chunk(2, dim=-1)
-        log_std = utils.log_std_fn(log_std, self.log_std_min, self.log_std_dif)
-        eps = torch.randn_like(mu)
-
-        # No masking
-        action_dims = None
-
-        log_pi = utils.gaussian_logprob(eps, log_std, size=action_dims)
-        pi = mu + eps * log_std.exp()
-        mu, pi, log_pi = utils.squash(mu, pi, log_pi)
-
-        return mu, pi, log_pi, log_std
-
-    def pi_action(self, z):
-        return self.pi(z)[1]  # just return the action
-
-    def Qs(self, z: Tensor, a: Tensor, return_type: str = "min", target: bool = False) -> Tensor:  # noqa: N802
-        """Predict state-action value for all of the learned Q functions.
-
-        Args:
-            z: (*, latent_dim) tensor for the current state's latent representation.
-            a: (*, action_dim) tensor for the action to be applied.
-            return_type can be one of [`min`, `avg`, `all`]:
-                - `min`: return the minimum of two randomly subsampled Q-values.
-                - `avg`: return the average of two randomly subsampled Q-values.
-                - `all`: return all Q-values.
-            target: Set to true to use the target Q functions.
-        Returns:
-            (q_ensemble, *) tensor for the value predictions of each learned Q function in the ensemble OR
-            (*,) tensor if return_min=True.
-        """
-        assert return_type in {"min", "avg", "all"}
-
-        z = torch.cat([z, a], dim=-1)
-        out = (self._target_Qs if target else self._Qs)(z)
-
-        if return_type == "all":
-            return out
-
-        Q1, Q2 = out[np.random.choice(self.config.num_q, 2, replace=False)]
-        Q1, Q2 = utils.two_hot_inv(Q1, self.config), utils.two_hot_inv(Q2, self.config)
-        return torch.min(Q1, Q2) if return_type == "min" else (Q1 + Q2) / 2
-
-
-class TDMPC2ObservationEncoder(nn.Module):
-    """Encode image and/or state vector observations."""
-
-    def __init__(self, config: TDMPC2Config):
-        """
-        Creates encoders for pixel and/or state modalities.
-        TODO(alexander-soare): The original work allows for multiple images by concatenating them along the
-            channel dimension. Re-implement this capability.
-        """
-        super().__init__()
-        self.config = config
-
-        for k in config.input_shapes:
-            if "observation.environment_state" in k:
-                obs_dim = config.input_shapes["observation.environment_state"][0]
-                self.env_state_enc_layers = utils.mlp(
-                    obs_dim,
-                    max(config.num_enc_layers - 1, 1) * [config.enc_dim],
-                    config.latent_dim,
-                    act=utils.SimNorm(config),
-                )
-            elif "observation.state" in k:
-                obs_dim = config.input_shapes["observation.state"][0]
-                self.state_enc_layers = utils.mlp(
-                    obs_dim,
-                    max(config.num_enc_layers - 1, 1) * [config.enc_dim],
-                    config.latent_dim,
-                    act=utils.SimNorm(config),
-                )
-            elif "observation.image" in k:
-                obs_shape = config.input_shapes["observation.image"]
-                self.image_enc_layers = utils.conv(obs_shape, config.num_channels, act=utils.SimNorm(config))
-                dummy_batch = torch.zeros(1, *config.input_shapes["observation.image"])
-                with torch.no_grad():
-                    out_shape = self.image_enc_layers(dummy_batch).shape[1]
-                self.image_enc_layers.extend(
-                    utils.mlp(
-                            out_shape,
-                            max(config.num_enc_layers - 1, 1) * [config.enc_dim],
-                            config.latent_dim,
-                            act=utils.SimNorm(config),
-                            ))
-            
-
-
-    def forward(self, obs_dict: dict[str, Tensor]) -> Tensor:
-        """Encode the image and/or state vector.
-
-        Each modality is encoded into a feature vector of size (latent_dim,) and then a uniform mean is taken
-        over all features.
-        """
-        feat = []
-        # NOTE: Order of observations matters here.
-        if "observation.image" in self.config.input_shapes:
-            feat.append(flatten_forward_unflatten(self.image_enc_layers, obs_dict["observation.image"]))
-        if "observation.environment_state" in self.config.input_shapes:
-            feat.append(self.env_state_enc_layers(obs_dict["observation.environment_state"]))
-        if "observation.state" in self.config.input_shapes:
-            feat.append(self.state_enc_layers(obs_dict["observation.state"]))
-
-        return torch.stack(feat, dim=0).mean(0)
-
-
-def random_shifts_aug(x: Tensor, max_random_shift_ratio: float) -> Tensor:
-    """Randomly shifts images horizontally and vertically.
-
-    Adapted from https://github.com/facebookresearch/drqv2
-    """
-    b, _, h, w = x.size()
-    assert h == w, "non-square images not handled yet"
-    pad = int(round(max_random_shift_ratio * h))
-    x = F.pad(x, tuple([pad] * 4), "replicate")
-    eps = 1.0 / (h + 2 * pad)
-    arange = torch.linspace(
-        -1.0 + eps,
-        1.0 - eps,
-        h + 2 * pad,
-        device=x.device,
-        dtype=torch.float32,
-    )[:h]
-    arange = einops.repeat(arange, "w -> h w 1", h=h)
-    base_grid = torch.cat([arange, arange.transpose(1, 0)], dim=2)
-    base_grid = einops.repeat(base_grid, "h w c -> b h w c", b=b)
-    # A random shift in units of pixels and within the boundaries of the padding.
-    shift = torch.randint(
-        0,
-        2 * pad + 1,
-        size=(b, 1, 1, 2),
-        device=x.device,
-        dtype=torch.float32,
-    )
-    shift *= 2.0 / (h + 2 * pad)
-    grid = base_grid + shift
-    return F.grid_sample(x, grid, padding_mode="zeros", align_corners=False)
-
-
-def update_ema_parameters(ema_net: nn.Module, net: nn.Module, alpha: float):
-    """Update EMA parameters in place with ema_param <- alpha * ema_param + (1 - alpha) * param."""
-    for ema_module, module in zip(ema_net.modules(), net.modules(), strict=True):
-        for (n_p_ema, p_ema), (n_p, p) in zip(
-            ema_module.named_parameters(recurse=False), module.named_parameters(recurse=False), strict=True
-        ):
-            assert n_p_ema == n_p, "Parameter names don't match for EMA model update"
-            if isinstance(p, dict):
-                raise RuntimeError("Dict parameter not supported")
-            if isinstance(module, nn.modules.batchnorm._BatchNorm) or not p.requires_grad:
-                # Copy BatchNorm parameters, and non-trainable parameters directly.
-                p_ema.copy_(p.to(dtype=p_ema.dtype).data)
-            with torch.no_grad():
-                p_ema.mul_(alpha)
-                p_ema.add_(p.to(dtype=p_ema.dtype).data, alpha=1 - alpha)
-
-
-def flatten_forward_unflatten(fn: Callable[[Tensor], Tensor], image_tensor: Tensor) -> Tensor:
-    """Helper to temporarily flatten extra dims at the start of the image tensor.
-
-    Args:
-        fn: Callable that the image tensor will be passed to. It should accept (B, C, H, W) and return
-            (B, *), where * is any number of dimensions.
-        image_tensor: An image tensor of shape (**, C, H, W), where ** is any number of dimensions, generally
-            different from *.
-    Returns:
-        A return value from the callable reshaped to (**, *).
-    """
-    if image_tensor.ndim == 4:
-        return fn(image_tensor)
-    start_dims = image_tensor.shape[:-3]
-    inp = torch.flatten(image_tensor, end_dim=-4)
-    flat_out = fn(inp)
-    return torch.reshape(flat_out, (*start_dims, *flat_out.shape[1:]))

lerobot/common/policies/tdmpc2/tdmpc2_utils.py

@@ -1,305 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F  # noqa: N812
-from functorch import combine_state_for_ensemble
-
-# Lifted directly from https://github.com/nicklashansen/tdmpc2
-DREG_BINS = None
-
-
-def soft_ce(pred, target, cfg):
-    """Computes the cross entropy loss between predictions and soft targets."""
-    pred = F.log_softmax(pred, dim=-1)
-    target = two_hot(target, cfg)
-    return -(target * pred).sum(-1, keepdim=True)
-
-
-@torch.jit.script
-def log_std(x, low, dif):
-    return low + 0.5 * dif * (torch.tanh(x) + 1)
-
-
-@torch.jit.script
-def _gaussian_residual(eps, log_std):
-    return -0.5 * eps.pow(2) - log_std
-
-
-@torch.jit.script
-def _gaussian_logprob(residual):
-    return residual - 0.5 * torch.log(2 * torch.pi)
-
-
-def gaussian_logprob(eps, log_std, size=None):
-    """Compute Gaussian log probability."""
-    residual = _gaussian_residual(eps, log_std).sum(-1, keepdim=True)
-    if size is None:
-        size = eps.size(-1)
-    return _gaussian_logprob(residual) * size
-
-
-@torch.jit.script
-def _squash(pi):
-    return torch.log(F.relu(1 - pi.pow(2)) + 1e-6)
-
-
-def squash(mu, pi, log_pi):
-    """Apply squashing function."""
-    mu = torch.tanh(mu)
-    pi = torch.tanh(pi)
-    log_pi -= _squash(pi).sum(-1, keepdim=True)
-    return mu, pi, log_pi
-
-
-@torch.jit.script
-def symexp(x):
-    """
-    Symmetric exponential function.
-    Adapted from https://github.com/danijar/dreamerv3.
-    """
-    return torch.sign(x) * (torch.exp(torch.abs(x)) - 1)
-
-
-@torch.jit.script
-def symlog(x):
-    """
-    Symmetric logarithmic function.
-    Adapted from https://github.com/danijar/dreamerv3.
-    """
-    return torch.sign(x) * torch.log(1 + torch.abs(x))
-
-
-@torch.jit.script
-def log_std_fn(x, low, dif):
-    return low + 0.5 * dif * (torch.tanh(x) + 1)
-
-
-def two_hot(x, cfg):
-    """Converts a batch of scalars to soft two-hot encoded targets for discrete regression."""
-    if cfg.num_bins == 0:
-        return x
-    elif cfg.num_bins == 1:
-        return symlog(x)
-    x = torch.clamp(symlog(x), cfg.vmin, cfg.vmax).squeeze(1)
-    bin_idx = torch.floor((x - cfg.vmin) / cfg.bin_size).long()
-    bin_offset = ((x - cfg.vmin) / cfg.bin_size - bin_idx.float()).unsqueeze(-1)
-    soft_two_hot = torch.zeros(x.size(0), cfg.num_bins, device=x.device)
-
-    # print("x shape:", x.shape)
-    # print("bin_idx shape:", bin_idx.shape)
-    # print("bin_offset shape:", bin_offset.shape)
-    # print("soft_two_hot shape:", soft_two_hot.shape)
-
-    # from IPython import embed; embed()
-
-    soft_two_hot.scatter_(1, bin_idx.unsqueeze(1), 1 - bin_offset)
-    soft_two_hot.scatter_(1, (bin_idx.unsqueeze(1) + 1) % cfg.num_bins, bin_offset)
-    return soft_two_hot
-
-
-def two_hot_inv(x, cfg):
-    """Converts a batch of soft two-hot encoded vectors to scalars."""
-    global DREG_BINS
-    if cfg.num_bins == 0:
-        return x
-    elif cfg.num_bins == 1:
-        return symexp(x)
-    if DREG_BINS is None:
-        DREG_BINS = torch.linspace(cfg.vmin, cfg.vmax, cfg.num_bins, device=x.device)
-    x = F.softmax(x, dim=-1)
-
-    # cloning bins to avoid the inference tensor errodr
-    x = torch.sum(x * DREG_BINS.clone(), dim=-1, keepdim=True)
-
-    return symexp(x)
-
-
-class Ensemble(nn.Module):
-    """
-    Vectorized ensemble of modules.
-    """
-
-    def __init__(self, modules, **kwargs):
-        super().__init__()
-        modules = nn.ModuleList(modules)
-        fn, params, _ = combine_state_for_ensemble(modules)
-        self.vmap = torch.vmap(fn, in_dims=(0, 0, None), randomness="different", **kwargs)
-        self.params = nn.ParameterList([nn.Parameter(p) for p in params])
-        self._repr = str(modules)
-
-    def forward(self, *args, **kwargs):
-        return self.vmap(list(self.params), (), *args, **kwargs)
-
-    def __repr__(self):
-        return "Vectorized " + self._repr
-
-
-class ShiftAug(nn.Module):
-    """
-    Random shift image augmentation.
-    Adapted from https://github.com/facebookresearch/drqv2
-    """
-
-    def __init__(self, pad=3):
-        super().__init__()
-        self.pad = pad
-
-    def forward(self, x):
-        x = x.float()
-        n, _, h, w = x.size()
-        assert h == w
-        padding = tuple([self.pad] * 4)
-        x = F.pad(x, padding, "replicate")
-        eps = 1.0 / (h + 2 * self.pad)
-        arange = torch.linspace(-1.0 + eps, 1.0 - eps, h + 2 * self.pad, device=x.device, dtype=x.dtype)[:h]
-        arange = arange.unsqueeze(0).repeat(h, 1).unsqueeze(2)
-        base_grid = torch.cat([arange, arange.transpose(1, 0)], dim=2)
-        base_grid = base_grid.unsqueeze(0).repeat(n, 1, 1, 1)
-        shift = torch.randint(0, 2 * self.pad + 1, size=(n, 1, 1, 2), device=x.device, dtype=x.dtype)
-        shift *= 2.0 / (h + 2 * self.pad)
-        grid = base_grid + shift
-        return F.grid_sample(x, grid, padding_mode="zeros", align_corners=False)
-
-
-class PixelPreprocess(nn.Module):
-    """
-    Normalizes pixel observations to [-0.5, 0.5].
-    """
-
-    def __init__(self):
-        super().__init__()
-
-    def forward(self, x):
-        return x.div_(255.0).sub_(0.5)
-
-
-class SimNorm(nn.Module):
-    """
-    Simplicial normalization.
-    Adapted from https://arxiv.org/abs/2204.00616.
-    """
-
-    def __init__(self, cfg):
-        super().__init__()
-        self.dim = cfg.simnorm_dim
-
-    def forward(self, x):
-        shp = x.shape
-        x = x.view(*shp[:-1], -1, self.dim)
-
-        x = F.softmax(x, dim=-1)
-        return x.view(*shp)
-
-    def __repr__(self):
-        return f"SimNorm(dim={self.dim})"
-
-
-class NormedLinear(nn.Linear):
-    """
-    Linear layer with LayerNorm, activation, and optionally dropout.
-    """
-
-    def __init__(self, *args, dropout=0.0, act=None, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.ln = nn.LayerNorm(self.out_features)
-        self.act = nn.Mish(inplace=True) if act is None else act
-        self.dropout = nn.Dropout(dropout, inplace=True) if dropout else None
-
-    def forward(self, x):
-        x = super().forward(x)
-        if self.dropout:
-            x = self.dropout(x)
-        return self.act(self.ln(x))
-
-    def __repr__(self):
-        repr_dropout = f", dropout={self.dropout.p}" if self.dropout else ""
-        return (
-            f"NormedLinear(in_features={self.in_features}, "
-            f"out_features={self.out_features}, "
-            f"bias={self.bias is not None}{repr_dropout}, "
-            f"act={self.act.__class__.__name__})"
-        )
-
-
-def mlp(in_dim, mlp_dims, out_dim, act=None, dropout=0.0):
-    """
-    Basic building block of TD-MPC2.
-    MLP with LayerNorm, Mish activations, and optionally dropout.
-    """
-    if isinstance(mlp_dims, int):
-        mlp_dims = [mlp_dims]
-    dims = [in_dim] + mlp_dims + [out_dim]
-    mlp = nn.ModuleList()
-    for i in range(len(dims) - 2):
-        mlp.append(NormedLinear(dims[i], dims[i + 1], dropout=dropout * (i == 0)))
-    mlp.append(NormedLinear(dims[-2], dims[-1], act=act) if act else nn.Linear(dims[-2], dims[-1]))
-    return nn.Sequential(*mlp)
-
-
-def conv(in_shape, num_channels, act=None):
-    """
-    Basic convolutional encoder for TD-MPC2 with raw image observations.
-    4 layers of convolution with ReLU activations, followed by a linear layer.
-    """
-    #assert in_shape[-1] == 64  # assumes rgb observations to be 64x64
-    layers = [
-        ShiftAug(),
-        PixelPreprocess(),
-        nn.Conv2d(in_shape[0], num_channels, 7, stride=2),
-        nn.ReLU(inplace=True),
-        nn.Conv2d(num_channels, num_channels, 5, stride=2),
-        nn.ReLU(inplace=True),
-        nn.Conv2d(num_channels, num_channels, 3, stride=2),
-        nn.ReLU(inplace=True),
-        nn.Conv2d(num_channels, num_channels, 3, stride=1),
-        nn.Flatten(),
-    ]
-    if act:
-        layers.append(act)
-    return nn.Sequential(*layers)
-
-
-class RunningScale:
-    """Running trimmed scale estimator."""
-
-    def __init__(self, cfg):
-        self.cfg = cfg
-        self._value = torch.ones(1, dtype=torch.float32, device=torch.device("cuda"))
-        self._percentiles = torch.tensor([5, 95], dtype=torch.float32, device=torch.device("cuda"))
-
-    def state_dict(self):
-        return {"value": self._value, "percentiles": self._percentiles}
-
-    def load_state_dict(self, state_dict):
-        self._value.data.copy_(state_dict["value"])
-        self._percentiles.data.copy_(state_dict["percentiles"])
-
-    @property
-    def value(self):
-        return self._value.cpu().item()
-
-    def _percentile(self, x):
-        x_dtype, x_shape = x.dtype, x.shape
-        x = x.view(x.shape[0], -1)
-        in_sorted, _ = torch.sort(x, dim=0)
-        positions = self._percentiles * (x.shape[0] - 1) / 100
-        floored = torch.floor(positions)
-        ceiled = floored + 1
-        ceiled[ceiled > x.shape[0] - 1] = x.shape[0] - 1
-        weight_ceiled = positions - floored
-        weight_floored = 1.0 - weight_ceiled
-        d0 = in_sorted[floored.long(), :] * weight_floored[:, None]
-        d1 = in_sorted[ceiled.long(), :] * weight_ceiled[:, None]
-        return (d0 + d1).view(-1, *x_shape[1:]).type(x_dtype)
-
-    def update(self, x):
-        percentiles = self._percentile(x.detach())
-        value = torch.clamp(percentiles[1] - percentiles[0], min=1.0)
-        self._value.data.lerp_(value, self.cfg.tau)
-
-    def __call__(self, x, update=False):
-        if update:
-            self.update(x)
-        return x * (1 / self.value)
-
-    def __repr__(self):
-        return f"RunningScale(S: {self.value})"

lerobot/common/policies/vqbet/modeling_vqbet.py

@@ -350,17 +350,22 @@ class VQBeTModel(nn.Module):
 
         # get action features (pass through GPT)
         features = self.policy(input_tokens)
-        # len(self.config.input_shapes) is the number of different observation modes. this line gets the index of action prompt tokens.
+        # len(self.config.input_shapes) is the number of different observation modes.
+        # this line gets the index of action prompt tokens.
         historical_act_pred_index = np.arange(0, n_obs_steps) * (len(self.config.input_shapes) + 1) + len(
             self.config.input_shapes
         )
 
         # only extract the output tokens at the position of action query:
-        # Behavior Transformer (BeT), and VQ-BeT are both sequence-to-sequence prediction models, mapping sequential observation to sequential action (please refer to section 2.2 in BeT paper https://arxiv.org/pdf/2206.11251).
-        # Thus, it predict historical action sequence, in addition to current and future actions (predicting future actions : optional).
-        features = torch.cat(
-            [features[:, historical_act_pred_index], features[:, -len_additional_action_token:]], dim=1
-        )
+        # Behavior Transformer (BeT), and VQ-BeT are both sequence-to-sequence prediction models,
+        # mapping sequential observation to sequential action (please refer to section 2.2 in BeT paper https://arxiv.org/pdf/2206.11251).
+        # Thus, it predicts a historical action sequence, in addition to current and future actions (predicting future actions : optional).
+        if len_additional_action_token > 0:
+            features = torch.cat(
+                [features[:, historical_act_pred_index], features[:, -len_additional_action_token:]], dim=1
+            )
+        else:
+            features = features[:, historical_act_pred_index]
         # pass through action head
         action_head_output = self.action_head(features)
         # if rollout, VQ-BeT don't calculate loss

lerobot/common/robot_devices/cameras/intelrealsense.py

@@ -0,0 +1,557 @@
+"""
+This file contains utilities for recording frames from Intel Realsense cameras.
+"""
+
+import argparse
+import concurrent.futures
+import logging
+import math
+import shutil
+import threading
+import time
+import traceback
+from collections import Counter
+from dataclasses import dataclass, replace
+from pathlib import Path
+from threading import Thread
+
+import numpy as np
+from PIL import Image
+
+from lerobot.common.robot_devices.utils import (
+    RobotDeviceAlreadyConnectedError,
+    RobotDeviceNotConnectedError,
+)
+from lerobot.common.utils.utils import capture_timestamp_utc
+from lerobot.scripts.control_robot import busy_wait
+
+SERIAL_NUMBER_INDEX = 1
+
+
+def find_cameras(raise_when_empty=True, mock=False) -> list[dict]:
+    """
+    Find the names and the serial numbers of the Intel RealSense cameras
+    connected to the computer.
+    """
+    if mock:
+        import tests.mock_pyrealsense2 as rs
+    else:
+        import pyrealsense2 as rs
+
+    cameras = []
+    for device in rs.context().query_devices():
+        serial_number = int(device.get_info(rs.camera_info(SERIAL_NUMBER_INDEX)))
+        name = device.get_info(rs.camera_info.name)
+        cameras.append(
+            {
+                "serial_number": serial_number,
+                "name": name,
+            }
+        )
+
+    if raise_when_empty and len(cameras) == 0:
+        raise OSError(
+            "Not a single camera was detected. Try re-plugging, or re-installing `librealsense` and its python wrapper `pyrealsense2`, or updating the firmware."
+        )
+
+    return cameras
+
+
+def save_image(img_array, serial_number, frame_index, images_dir):
+    try:
+        img = Image.fromarray(img_array)
+        path = images_dir / f"camera_{serial_number}_frame_{frame_index:06d}.png"
+        path.parent.mkdir(parents=True, exist_ok=True)
+        img.save(str(path), quality=100)
+        logging.info(f"Saved image: {path}")
+    except Exception as e:
+        logging.error(f"Failed to save image for camera {serial_number} frame {frame_index}: {e}")
+
+
+def save_images_from_cameras(
+    images_dir: Path,
+    serial_numbers: list[int] | None = None,
+    fps=None,
+    width=None,
+    height=None,
+    record_time_s=2,
+    mock=False,
+):
+    """
+    Initializes all the cameras and saves images to the directory. Useful to visually identify the camera
+    associated to a given serial number.
+    """
+    if serial_numbers is None or len(serial_numbers) == 0:
+        camera_infos = find_cameras(mock=mock)
+        serial_numbers = [cam["serial_number"] for cam in camera_infos]
+
+    if mock:
+        import tests.mock_cv2 as cv2
+    else:
+        import cv2
+
+    print("Connecting cameras")
+    cameras = []
+    for cam_sn in serial_numbers:
+        print(f"{cam_sn=}")
+        camera = IntelRealSenseCamera(cam_sn, fps=fps, width=width, height=height, mock=mock)
+        camera.connect()
+        print(
+            f"IntelRealSenseCamera({camera.serial_number}, fps={camera.fps}, width={camera.width}, height={camera.height}, color_mode={camera.color_mode})"
+        )
+        cameras.append(camera)
+
+    images_dir = Path(images_dir)
+    if images_dir.exists():
+        shutil.rmtree(
+            images_dir,
+        )
+    images_dir.mkdir(parents=True, exist_ok=True)
+
+    print(f"Saving images to {images_dir}")
+    frame_index = 0
+    start_time = time.perf_counter()
+    try:
+        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
+            while True:
+                now = time.perf_counter()
+
+                for camera in cameras:
+                    # If we use async_read when fps is None, the loop will go full speed, and we will end up
+                    # saving the same images from the cameras multiple times until the RAM/disk is full.
+                    image = camera.read() if fps is None else camera.async_read()
+                    if image is None:
+                        print("No Frame")
+
+                    bgr_converted_image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+
+                    executor.submit(
+                        save_image,
+                        bgr_converted_image,
+                        camera.serial_number,
+                        frame_index,
+                        images_dir,
+                    )
+
+                if fps is not None:
+                    dt_s = time.perf_counter() - now
+                    busy_wait(1 / fps - dt_s)
+
+                if time.perf_counter() - start_time > record_time_s:
+                    break
+
+                print(f"Frame: {frame_index:04d}\tLatency (ms): {(time.perf_counter() - now) * 1000:.2f}")
+
+                frame_index += 1
+    finally:
+        print(f"Images have been saved to {images_dir}")
+        for camera in cameras:
+            camera.disconnect()
+
+
+@dataclass
+class IntelRealSenseCameraConfig:
+    """
+    Example of tested options for Intel Real Sense D405:
+
+    ```python
+    IntelRealSenseCameraConfig(30, 640, 480)
+    IntelRealSenseCameraConfig(60, 640, 480)
+    IntelRealSenseCameraConfig(90, 640, 480)
+    IntelRealSenseCameraConfig(30, 1280, 720)
+    IntelRealSenseCameraConfig(30, 640, 480, use_depth=True)
+    IntelRealSenseCameraConfig(30, 640, 480, rotation=90)
+    ```
+    """
+
+    fps: int | None = None
+    width: int | None = None
+    height: int | None = None
+    color_mode: str = "rgb"
+    use_depth: bool = False
+    force_hardware_reset: bool = True
+    rotation: int | None = None
+    mock: bool = False
+
+    def __post_init__(self):
+        if self.color_mode not in ["rgb", "bgr"]:
+            raise ValueError(
+                f"`color_mode` is expected to be 'rgb' or 'bgr', but {self.color_mode} is provided."
+            )
+
+        at_least_one_is_not_none = self.fps is not None or self.width is not None or self.height is not None
+        at_least_one_is_none = self.fps is None or self.width is None or self.height is None
+        if at_least_one_is_not_none and at_least_one_is_none:
+            raise ValueError(
+                "For `fps`, `width` and `height`, either all of them need to be set, or none of them, "
+                f"but {self.fps=}, {self.width=}, {self.height=} were provided."
+            )
+
+        if self.rotation not in [-90, None, 90, 180]:
+            raise ValueError(f"`rotation` must be in [-90, None, 90, 180] (got {self.rotation})")
+
+
+class IntelRealSenseCamera:
+    """
+    The IntelRealSenseCamera class is similar to OpenCVCamera class but adds additional features for Intel Real Sense cameras:
+    - is instantiated with the serial number of the camera - won't randomly change as it can be the case of OpenCVCamera for Linux,
+    - can also be instantiated with the camera's name — if it's unique — using IntelRealSenseCamera.init_from_name(),
+    - depth map can be returned.
+
+    To find the camera indices of your cameras, you can run our utility script that will save a few frames for each camera:
+    ```bash
+    python lerobot/common/robot_devices/cameras/intelrealsense.py --images-dir outputs/images_from_intelrealsense_cameras
+    ```
+
+    When an IntelRealSenseCamera is instantiated, if no specific config is provided, the default fps, width, height and color_mode
+    of the given camera will be used.
+
+    Example of usage:
+    ```python
+    # Instantiate with its serial number
+    camera = IntelRealSenseCamera(128422271347)
+    # Or by its name if it's unique
+    camera = IntelRealSenseCamera.init_from_name("Intel RealSense D405")
+    camera.connect()
+    color_image = camera.read()
+    # when done using the camera, consider disconnecting
+    camera.disconnect()
+    ```
+
+    Example of changing default fps, width, height and color_mode:
+    ```python
+    camera = IntelRealSenseCamera(serial_number, fps=30, width=1280, height=720)
+    camera = connect()  # applies the settings, might error out if these settings are not compatible with the camera
+
+    camera = IntelRealSenseCamera(serial_number, fps=90, width=640, height=480)
+    camera = connect()
+
+    camera = IntelRealSenseCamera(serial_number, fps=90, width=640, height=480, color_mode="bgr")
+    camera = connect()
+    ```
+
+    Example of returning depth:
+    ```python
+    camera = IntelRealSenseCamera(serial_number, use_depth=True)
+    camera.connect()
+    color_image, depth_map = camera.read()
+    ```
+    """
+
+    def __init__(
+        self,
+        serial_number: int,
+        config: IntelRealSenseCameraConfig | None = None,
+        **kwargs,
+    ):
+        if config is None:
+            config = IntelRealSenseCameraConfig()
+
+        # Overwrite the config arguments using kwargs
+        config = replace(config, **kwargs)
+
+        self.serial_number = serial_number
+        self.fps = config.fps
+        self.width = config.width
+        self.height = config.height
+        self.color_mode = config.color_mode
+        self.use_depth = config.use_depth
+        self.force_hardware_reset = config.force_hardware_reset
+        self.mock = config.mock
+
+        self.camera = None
+        self.is_connected = False
+        self.thread = None
+        self.stop_event = None
+        self.color_image = None
+        self.depth_map = None
+        self.logs = {}
+
+        if self.mock:
+            import tests.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
+        elif config.rotation == 90:
+            self.rotation = cv2.ROTATE_90_CLOCKWISE
+        elif config.rotation == 180:
+            self.rotation = cv2.ROTATE_180
+
+    @classmethod
+    def init_from_name(cls, name: str, config: IntelRealSenseCameraConfig | None = None, **kwargs):
+        camera_infos = find_cameras()
+        camera_names = [cam["name"] for cam in camera_infos]
+        this_name_count = Counter(camera_names)[name]
+        if this_name_count > 1:
+            # TODO(aliberts): Test this with multiple identical cameras (Aloha)
+            raise ValueError(
+                f"Multiple {name} cameras have been detected. Please use their serial number to instantiate them."
+            )
+
+        name_to_serial_dict = {cam["name"]: cam["serial_number"] for cam in camera_infos}
+        cam_sn = name_to_serial_dict[name]
+
+        if config is None:
+            config = IntelRealSenseCameraConfig()
+
+        # Overwrite the config arguments using kwargs
+        config = replace(config, **kwargs)
+
+        return cls(serial_number=cam_sn, config=config, **kwargs)
+
+    def connect(self):
+        if self.is_connected:
+            raise RobotDeviceAlreadyConnectedError(
+                f"IntelRealSenseCamera({self.serial_number}) is already connected."
+            )
+
+        if self.mock:
+            import tests.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:
+            # TODO(rcadene): can we set rgb8 directly?
+            config.enable_stream(rs.stream.color, self.width, self.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)
+            else:
+                config.enable_stream(rs.stream.depth)
+
+        self.camera = rs.pipeline()
+        try:
+            profile = self.camera.start(config)
+            is_camera_open = True
+        except RuntimeError:
+            is_camera_open = False
+            traceback.print_exc()
+
+        # If the camera doesn't work, display the camera indices corresponding to
+        # valid cameras.
+        if not is_camera_open:
+            # Verify that the provided `serial_number` is valid before printing the traceback
+            camera_infos = find_cameras()
+            serial_numbers = [cam["serial_number"] for cam in camera_infos]
+            if self.serial_number not in serial_numbers:
+                raise ValueError(
+                    f"`serial_number` is expected to be one of these available cameras {serial_numbers}, but {self.serial_number} is provided instead. "
+                    "To find the serial number you should use, run `python lerobot/common/robot_devices/cameras/intelrealsense.py`."
+                )
+
+            raise OSError(f"Can't access IntelRealSenseCamera({self.serial_number}).")
+
+        color_stream = profile.get_stream(rs.stream.color)
+        color_profile = color_stream.as_video_stream_profile()
+        actual_fps = color_profile.fps()
+        actual_width = color_profile.width()
+        actual_height = color_profile.height()
+
+        # Using `math.isclose` since actual fps can be a float (e.g. 29.9 instead of 30)
+        if self.fps is not None and not math.isclose(self.fps, actual_fps, rel_tol=1e-3):
+            # Using `OSError` since it's a broad that encompasses issues related to device communication
+            raise OSError(
+                f"Can't set {self.fps=} for IntelRealSenseCamera({self.serial_number}). Actual value is {actual_fps}."
+            )
+        if self.width is not None and self.width != actual_width:
+            raise OSError(
+                f"Can't set {self.width=} for IntelRealSenseCamera({self.serial_number}). Actual value is {actual_width}."
+            )
+        if self.height is not None and self.height != actual_height:
+            raise OSError(
+                f"Can't set {self.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.is_connected = True
+
+    def read(self, temporary_color: str | None = None) -> np.ndarray | tuple[np.ndarray, np.ndarray]:
+        """Read a frame from the camera returned in the format height x width x channels (e.g. 480 x 640 x 3)
+        of type `np.uint8`, contrarily to the pytorch format which is float channel first.
+
+        When `use_depth=True`, returns a tuple `(color_image, depth_map)` with a depth map in the format
+        height x width (e.g. 480 x 640) of type np.uint16.
+
+        Note: Reading a frame is done every `camera.fps` times per second, and it is blocking.
+        If you are reading data from other sensors, we advise to use `camera.async_read()` which is non blocking version of `camera.read()`.
+        """
+        if not self.is_connected:
+            raise RobotDeviceNotConnectedError(
+                f"IntelRealSenseCamera({self.serial_number}) is not connected. Try running `camera.connect()` first."
+            )
+
+        if self.mock:
+            import tests.mock_cv2 as cv2
+        else:
+            import cv2
+
+        start_time = time.perf_counter()
+
+        frame = self.camera.wait_for_frames(timeout_ms=5000)
+
+        color_frame = frame.get_color_frame()
+
+        if not color_frame:
+            raise OSError(f"Can't capture color image from IntelRealSenseCamera({self.serial_number}).")
+
+        color_image = np.asanyarray(color_frame.get_data())
+
+        requested_color_mode = self.color_mode if temporary_color is None else temporary_color
+        if requested_color_mode not in ["rgb", "bgr"]:
+            raise ValueError(
+                f"Expected color values are 'rgb' or 'bgr', but {requested_color_mode} is provided."
+            )
+
+        # IntelRealSense uses RGB format as default (red, green, blue).
+        if requested_color_mode == "bgr":
+            color_image = cv2.cvtColor(color_image, cv2.COLOR_RGB2BGR)
+
+        h, w, _ = color_image.shape
+        if h != self.height or w != self.width:
+            raise OSError(
+                f"Can't capture color image with expected height and width ({self.height} x {self.width}). ({h} x {w}) returned instead."
+            )
+
+        if self.rotation is not None:
+            color_image = cv2.rotate(color_image, self.rotation)
+
+        # log the number of seconds it took to read the image
+        self.logs["delta_timestamp_s"] = time.perf_counter() - start_time
+
+        # log the utc time at which the image was received
+        self.logs["timestamp_utc"] = capture_timestamp_utc()
+
+        if self.use_depth:
+            depth_frame = frame.get_depth_frame()
+            if not depth_frame:
+                raise OSError(f"Can't capture depth image from IntelRealSenseCamera({self.serial_number}).")
+
+            depth_map = np.asanyarray(depth_frame.get_data())
+
+            h, w = depth_map.shape
+            if h != self.height or w != self.width:
+                raise OSError(
+                    f"Can't capture depth map with expected height and width ({self.height} x {self.width}). ({h} x {w}) returned instead."
+                )
+
+            if self.rotation is not None:
+                depth_map = cv2.rotate(depth_map, self.rotation)
+
+            return color_image, depth_map
+        else:
+            return color_image
+
+    def read_loop(self):
+        while not self.stop_event.is_set():
+            if self.use_depth:
+                self.color_image, self.depth_map = self.read()
+            else:
+                self.color_image = self.read()
+
+    def async_read(self):
+        """Access the latest color image"""
+        if not self.is_connected:
+            raise RobotDeviceNotConnectedError(
+                f"IntelRealSenseCamera({self.serial_number}) is not connected. Try running `camera.connect()` first."
+            )
+
+        if self.thread is None:
+            self.stop_event = threading.Event()
+            self.thread = Thread(target=self.read_loop, args=())
+            self.thread.daemon = True
+            self.thread.start()
+
+        num_tries = 0
+        while self.color_image is None:
+            # TODO(rcadene, aliberts): intelrealsense has diverged compared to opencv over here
+            num_tries += 1
+            time.sleep(1 / self.fps)
+            if num_tries > self.fps and (self.thread.ident is None or not self.thread.is_alive()):
+                raise Exception(
+                    "The thread responsible for `self.async_read()` took too much time to start. There might be an issue. Verify that `self.thread.start()` has been called."
+                )
+
+        if self.use_depth:
+            return self.color_image, self.depth_map
+        else:
+            return self.color_image
+
+    def disconnect(self):
+        if not self.is_connected:
+            raise RobotDeviceNotConnectedError(
+                f"IntelRealSenseCamera({self.serial_number}) is not connected. Try running `camera.connect()` first."
+            )
+
+        if self.thread is not None and self.thread.is_alive():
+            # wait for the thread to finish
+            self.stop_event.set()
+            self.thread.join()
+            self.thread = None
+            self.stop_event = None
+
+        self.camera.stop()
+        self.camera = None
+
+        self.is_connected = False
+
+    def __del__(self):
+        if getattr(self, "is_connected", False):
+            self.disconnect()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Save a few frames using `IntelRealSenseCamera` for all cameras connected to the computer, or a selected subset."
+    )
+    parser.add_argument(
+        "--serial-numbers",
+        type=int,
+        nargs="*",
+        default=None,
+        help="List of serial numbers used to instantiate the `IntelRealSenseCamera`. If not provided, find and use all available camera indices.",
+    )
+    parser.add_argument(
+        "--fps",
+        type=int,
+        default=30,
+        help="Set the number of frames recorded per seconds for all cameras. If not provided, use the default fps of each camera.",
+    )
+    parser.add_argument(
+        "--width",
+        type=str,
+        default=640,
+        help="Set the width for all cameras. If not provided, use the default width of each camera.",
+    )
+    parser.add_argument(
+        "--height",
+        type=str,
+        default=480,
+        help="Set the height for all cameras. If not provided, use the default height of each camera.",
+    )
+    parser.add_argument(
+        "--images-dir",
+        type=Path,
+        default="outputs/images_from_intelrealsense_cameras",
+        help="Set directory to save a few frames for each camera.",
+    )
+    parser.add_argument(
+        "--record-time-s",
+        type=float,
+        default=2.0,
+        help="Set the number of seconds used to record the frames. By default, 2 seconds.",
+    )
+    args = parser.parse_args()
+    save_images_from_cameras(**vars(args))

lerobot/common/robot_devices/cameras/opencv.py

@@ -13,17 +13,15 @@ from dataclasses import dataclass, replace
 from pathlib import Path
 from threading import Thread
 
-import cv2
 import numpy as np
 from PIL import Image
 
-from lerobot.common.robot_devices.utils import RobotDeviceAlreadyConnectedError, RobotDeviceNotConnectedError
+from lerobot.common.robot_devices.utils import (
+    RobotDeviceAlreadyConnectedError,
+    RobotDeviceNotConnectedError,
+    busy_wait,
+)
 from lerobot.common.utils.utils import capture_timestamp_utc
-from lerobot.scripts.control_robot import busy_wait
-
-# Use 1 thread to avoid blocking the main thread. Especially useful during data collection
-# when other threads are used to save the images.
-cv2.setNumThreads(1)
 
 # The maximum opencv device index depends on your operating system. For instance,
 # if you have 3 cameras, they should be associated to index 0, 1, and 2. This is the case
@@ -33,20 +31,44 @@ cv2.setNumThreads(1)
 MAX_OPENCV_INDEX = 60
 
 
-def find_camera_indices(raise_when_empty=False, max_index_search_range=MAX_OPENCV_INDEX):
+def find_cameras(raise_when_empty=False, max_index_search_range=MAX_OPENCV_INDEX, mock=False) -> list[dict]:
+    cameras = []
     if platform.system() == "Linux":
-        # Linux uses camera ports
         print("Linux detected. Finding available camera indices through scanning '/dev/video*' ports")
-        possible_camera_ids = []
-        for port in Path("/dev").glob("video*"):
-            camera_idx = int(str(port).replace("/dev/video", ""))
-            possible_camera_ids.append(camera_idx)
+        possible_ports = [str(port) for port in Path("/dev").glob("video*")]
+        ports = _find_cameras(possible_ports, mock=mock)
+        for port in ports:
+            cameras.append(
+                {
+                    "port": port,
+                    "index": int(port.removeprefix("/dev/video")),
+                }
+            )
     else:
         print(
             "Mac or Windows detected. Finding available camera indices through "
             f"scanning all indices from 0 to {MAX_OPENCV_INDEX}"
         )
-        possible_camera_ids = range(max_index_search_range)
+        possible_indices = range(max_index_search_range)
+        indices = _find_cameras(possible_indices, mock=mock)
+        for index in indices:
+            cameras.append(
+                {
+                    "port": None,
+                    "index": index,
+                }
+            )
+
+    return cameras
+
+
+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
+    else:
+        import cv2
 
     camera_ids = []
     for camera_idx in possible_camera_ids:
@@ -67,6 +89,16 @@ def find_camera_indices(raise_when_empty=False, max_index_search_range=MAX_OPENC
     return camera_ids
 
 
+def is_valid_unix_path(path: str) -> bool:
+    """Note: if 'path' points to a symlink, this will return True only if the target exists"""
+    p = Path(path)
+    return p.is_absolute() and p.exists()
+
+
+def get_camera_index_from_unix_port(port: Path) -> int:
+    return int(str(port.resolve()).removeprefix("/dev/video"))
+
+
 def save_image(img_array, camera_index, frame_index, images_dir):
     img = Image.fromarray(img_array)
     path = images_dir / f"camera_{camera_index:02d}_frame_{frame_index:06d}.png"
@@ -75,15 +107,26 @@ def save_image(img_array, camera_index, frame_index, images_dir):
 
 
 def save_images_from_cameras(
-    images_dir: Path, camera_ids: list[int] | None = None, fps=None, width=None, height=None, record_time_s=2
+    images_dir: Path,
+    camera_ids: list | None = None,
+    fps=None,
+    width=None,
+    height=None,
+    record_time_s=2,
+    mock=False,
 ):
-    if camera_ids is None:
-        camera_ids = find_camera_indices()
+    """
+    Initializes all the cameras and saves images to the directory. Useful to visually identify the camera
+    associated to a given camera index.
+    """
+    if camera_ids is None or len(camera_ids) == 0:
+        camera_infos = find_cameras(mock=mock)
+        camera_ids = [cam["index"] for cam in camera_infos]
 
     print("Connecting cameras")
     cameras = []
     for cam_idx in camera_ids:
-        camera = OpenCVCamera(cam_idx, fps=fps, width=width, height=height)
+        camera = OpenCVCamera(cam_idx, fps=fps, width=width, height=height, mock=mock)
         camera.connect()
         print(
             f"OpenCVCamera({camera.camera_index}, fps={camera.fps}, width={camera.width}, "
@@ -101,7 +144,7 @@ def save_images_from_cameras(
     print(f"Saving images to {images_dir}")
     frame_index = 0
     start_time = time.perf_counter()
-    with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
+    with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
         while True:
             now = time.perf_counter()
 
@@ -113,7 +156,7 @@ def save_images_from_cameras(
                 executor.submit(
                     save_image,
                     image,
-                    camera.camera_index,
+                    camera.index,
                     frame_index,
                     images_dir,
                 )
@@ -122,11 +165,11 @@ def save_images_from_cameras(
                 dt_s = time.perf_counter() - now
                 busy_wait(1 / fps - dt_s)
 
+            print(f"Frame: {frame_index:04d}\tLatency (ms): {(time.perf_counter() - now) * 1000:.2f}")
+
             if time.perf_counter() - start_time > record_time_s:
                 break
 
-            print(f"Frame: {frame_index:04d}\tLatency (ms): {(time.perf_counter() - now) * 1000:.2f}")
-
             frame_index += 1
 
     print(f"Images have been saved to {images_dir}")
@@ -149,13 +192,18 @@ class OpenCVCameraConfig:
     width: int | None = None
     height: int | None = None
     color_mode: str = "rgb"
+    rotation: int | None = None
+    mock: bool = False
 
     def __post_init__(self):
         if self.color_mode not in ["rgb", "bgr"]:
             raise ValueError(
-                f"Expected color_mode values are 'rgb' or 'bgr', but {self.color_mode} is provided."
+                f"`color_mode` is expected to be 'rgb' or 'bgr', but {self.color_mode} is provided."
             )
 
+        if self.rotation not in [-90, None, 90, 180]:
+            raise ValueError(f"`rotation` must be in [-90, None, 90, 180] (got {self.rotation})")
+
 
 class OpenCVCamera:
     """
@@ -196,17 +244,32 @@ class OpenCVCamera:
     ```
     """
 
-    def __init__(self, camera_index: int, config: OpenCVCameraConfig | None = None, **kwargs):
+    def __init__(self, camera_index: int | str, config: OpenCVCameraConfig | None = None, **kwargs):
         if config is None:
             config = OpenCVCameraConfig()
+
         # Overwrite config arguments using kwargs
         config = replace(config, **kwargs)
 
         self.camera_index = camera_index
+        self.port = None
+
+        # Linux uses ports for connecting to cameras
+        if platform.system() == "Linux":
+            if isinstance(self.camera_index, int):
+                self.port = Path(f"/dev/video{self.camera_index}")
+            elif isinstance(self.camera_index, str) and is_valid_unix_path(self.camera_index):
+                self.port = Path(self.camera_index)
+                # Retrieve the camera index from a potentially symlinked path
+                self.camera_index = get_camera_index_from_unix_port(self.port)
+            else:
+                raise ValueError(f"Please check the provided camera_index: {camera_index}")
+
         self.fps = config.fps
         self.width = config.width
         self.height = config.height
         self.color_mode = config.color_mode
+        self.mock = config.mock
 
         self.camera = None
         self.is_connected = False
@@ -215,43 +278,60 @@ class OpenCVCamera:
         self.color_image = None
         self.logs = {}
 
+        if self.mock:
+            import tests.mock_cv2 as cv2
+        else:
+            import cv2
+
+        # TODO(aliberts): Do we keep original width/height or do we define them after rotation?
+        self.rotation = None
+        if config.rotation == -90:
+            self.rotation = cv2.ROTATE_90_COUNTERCLOCKWISE
+        elif config.rotation == 90:
+            self.rotation = cv2.ROTATE_90_CLOCKWISE
+        elif config.rotation == 180:
+            self.rotation = cv2.ROTATE_180
+
     def connect(self):
         if self.is_connected:
-            raise RobotDeviceAlreadyConnectedError(f"Camera {self.camera_index} is already connected.")
+            raise RobotDeviceAlreadyConnectedError(f"OpenCVCamera({self.camera_index}) is already connected.")
 
+        if self.mock:
+            import tests.mock_cv2 as cv2
+        else:
+            import cv2
+
+            # Use 1 thread to avoid blocking the main thread. Especially useful during data collection
+            # when other threads are used to save the images.
+            cv2.setNumThreads(1)
+
+        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`.
-
-        if platform.system() == "Linux":
-            # Linux uses ports for connecting to cameras
-            tmp_camera = cv2.VideoCapture(f"/dev/video{self.camera_index}")
-        else:
-            tmp_camera = cv2.VideoCapture(self.camera_index)
-
+        tmp_camera = cv2.VideoCapture(camera_idx)
         is_camera_open = tmp_camera.isOpened()
         # Release camera to make it accessible for `find_camera_indices`
+        tmp_camera.release()
         del tmp_camera
 
         # If the camera doesn't work, display the camera indices corresponding to
         # valid cameras.
         if not is_camera_open:
             # Verify that the provided `camera_index` is valid before printing the traceback
-            available_cam_ids = find_camera_indices()
+            cameras_info = find_cameras()
+            available_cam_ids = [cam["index"] for cam in cameras_info]
             if self.camera_index not in available_cam_ids:
                 raise ValueError(
                     f"`camera_index` is expected to be one of these available cameras {available_cam_ids}, but {self.camera_index} is provided instead. "
                     "To find the camera index you should use, run `python lerobot/common/robot_devices/cameras/opencv.py`."
                 )
 
-            raise OSError(f"Can't access camera {self.camera_index}.")
+            raise OSError(f"Can't access OpenCVCamera({camera_idx}).")
 
         # 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.
-        if platform.system() == "Linux":
-            self.camera = cv2.VideoCapture(f"/dev/video{self.camera_index}")
-        else:
-            self.camera = cv2.VideoCapture(self.camera_index)
+        self.camera = cv2.VideoCapture(camera_idx)
 
         if self.fps is not None:
             self.camera.set(cv2.CAP_PROP_FPS, self.fps)
@@ -264,28 +344,30 @@ class OpenCVCamera:
         actual_width = self.camera.get(cv2.CAP_PROP_FRAME_WIDTH)
         actual_height = self.camera.get(cv2.CAP_PROP_FRAME_HEIGHT)
 
+        # Using `math.isclose` since actual fps can be a float (e.g. 29.9 instead of 30)
         if self.fps is not None and not math.isclose(self.fps, actual_fps, rel_tol=1e-3):
+            # Using `OSError` since it's a broad that encompasses issues related to device communication
             raise OSError(
-                f"Can't set {self.fps=} for camera {self.camera_index}. Actual value is {actual_fps}."
+                f"Can't set {self.fps=} for OpenCVCamera({self.camera_index}). Actual value is {actual_fps}."
             )
-        if self.width is not None and self.width != actual_width:
+        if self.width is not None and not math.isclose(self.width, actual_width, rel_tol=1e-3):
             raise OSError(
-                f"Can't set {self.width=} for camera {self.camera_index}. Actual value is {actual_width}."
+                f"Can't set {self.width=} for OpenCVCamera({self.camera_index}). Actual value is {actual_width}."
             )
-        if self.height is not None and self.height != actual_height:
+        if self.height is not None and not math.isclose(self.height, actual_height, rel_tol=1e-3):
             raise OSError(
-                f"Can't set {self.height=} for camera {self.camera_index}. Actual value is {actual_height}."
+                f"Can't set {self.height=} for OpenCVCamera({self.camera_index}). Actual value is {actual_height}."
             )
 
-        self.fps = actual_fps
-        self.width = actual_width
-        self.height = actual_height
+        self.fps = round(actual_fps)
+        self.width = round(actual_width)
+        self.height = round(actual_height)
 
         self.is_connected = True
 
     def read(self, temporary_color_mode: str | None = None) -> np.ndarray:
         """Read a frame from the camera returned in the format (height, width, channels)
-        (e.g. (640, 480, 3)), contrarily to the pytorch format which is channel first.
+        (e.g. 480 x 640 x 3), contrarily to the pytorch format which is channel first.
 
         Note: Reading a frame is done every `camera.fps` times per second, and it is blocking.
         If you are reading data from other sensors, we advise to use `camera.async_read()` which is non blocking version of `camera.read()`.
@@ -298,6 +380,7 @@ class OpenCVCamera:
         start_time = time.perf_counter()
 
         ret, color_image = self.camera.read()
+
         if not ret:
             raise OSError(f"Can't capture color image from camera {self.camera_index}.")
 
@@ -308,10 +391,15 @@ class OpenCVCamera:
                 f"Expected color values are 'rgb' or 'bgr', but {requested_color_mode} is provided."
             )
 
-        # OpenCV uses BGR format as default (blue, green red) for all operations, including displaying images.
+        # OpenCV uses BGR format as default (blue, green, red) for all operations, including displaying images.
         # However, Deep Learning framework such as LeRobot uses RGB format as default to train neural networks,
         # so we convert the image color from BGR to RGB.
         if requested_color_mode == "rgb":
+            if self.mock:
+                import tests.mock_cv2 as cv2
+            else:
+                import cv2
+
             color_image = cv2.cvtColor(color_image, cv2.COLOR_BGR2RGB)
 
         h, w, _ = color_image.shape
@@ -320,17 +408,25 @@ class OpenCVCamera:
                 f"Can't capture color image with expected height and width ({self.height} x {self.width}). ({h} x {w}) returned instead."
             )
 
+        if self.rotation is not None:
+            color_image = cv2.rotate(color_image, self.rotation)
+
         # log the number of seconds it took to read the image
         self.logs["delta_timestamp_s"] = time.perf_counter() - start_time
 
         # log the utc time at which the image was received
         self.logs["timestamp_utc"] = capture_timestamp_utc()
 
+        self.color_image = color_image
+
         return color_image
 
     def read_loop(self):
-        while self.stop_event is None or not self.stop_event.is_set():
-            self.color_image = self.read()
+        while not self.stop_event.is_set():
+            try:
+                self.color_image = self.read()
+            except Exception as e:
+                print(f"Error reading in thread: {e}")
 
     def async_read(self):
         if not self.is_connected:
@@ -345,15 +441,14 @@ class OpenCVCamera:
             self.thread.start()
 
         num_tries = 0
-        while self.color_image is None:
-            num_tries += 1
-            time.sleep(1 / self.fps)
-            if num_tries > self.fps and (self.thread.ident is None or not self.thread.is_alive()):
-                raise Exception(
-                    "The thread responsible for `self.async_read()` took too much time to start. There might be an issue. Verify that `self.thread.start()` has been called."
-                )
+        while True:
+            if self.color_image is not None:
+                return self.color_image
 
-        return self.color_image
+            time.sleep(1 / self.fps)
+            num_tries += 1
+            if num_tries > self.fps * 2:
+                raise TimeoutError("Timed out waiting for async_read() to start.")
 
     def disconnect(self):
         if not self.is_connected:
@@ -361,16 +456,14 @@ class OpenCVCamera:
                 f"OpenCVCamera({self.camera_index}) is not connected. Try running `camera.connect()` first."
             )
 
-        if self.thread is not None and self.thread.is_alive():
-            # wait for the thread to finish
+        if self.thread is not None:
             self.stop_event.set()
-            self.thread.join()
+            self.thread.join()  # wait for the thread to finish
             self.thread = None
             self.stop_event = None
 
         self.camera.release()
         self.camera = None
-
         self.is_connected = False
 
     def __del__(self):
@@ -416,7 +509,7 @@ if __name__ == "__main__":
     parser.add_argument(
         "--record-time-s",
         type=float,
-        default=2.0,
+        default=4.0,
         help="Set the number of seconds used to record the frames. By default, 2 seconds.",
     )
     args = parser.parse_args()

lerobot/common/robot_devices/cameras/utils.py

@@ -1,55 +1,8 @@
-from pathlib import Path
 from typing import Protocol
 
-import cv2
-import einops
 import numpy as np
 
 
-def write_shape_on_image_inplace(image):
-    height, width = image.shape[:2]
-    text = f"Width: {width} Height: {height}"
-
-    # Define the font, scale, color, and thickness
-    font = cv2.FONT_HERSHEY_SIMPLEX
-    font_scale = 1
-    color = (255, 0, 0)  # Blue in BGR
-    thickness = 2
-
-    position = (10, height - 10)  # 10 pixels from the bottom-left corner
-    cv2.putText(image, text, position, font, font_scale, color, thickness)
-
-
-def save_color_image(image, path, write_shape=False):
-    path = Path(path)
-    path.parent.mkdir(parents=True, exist_ok=True)
-    if write_shape:
-        write_shape_on_image_inplace(image)
-    cv2.imwrite(str(path), image)
-
-
-def save_depth_image(depth, path, write_shape=False):
-    path = Path(path)
-    path.parent.mkdir(parents=True, exist_ok=True)
-
-    # Apply colormap on depth image (image must be converted to 8-bit per pixel first)
-    depth_image = cv2.applyColorMap(cv2.convertScaleAbs(depth, alpha=0.03), cv2.COLORMAP_JET)
-
-    if write_shape:
-        write_shape_on_image_inplace(depth_image)
-    cv2.imwrite(str(path), depth_image)
-
-
-def convert_torch_image_to_cv2(tensor, rgb_to_bgr=True):
-    assert tensor.ndim == 3
-    c, h, w = tensor.shape
-    assert c < h and c < w
-    color_image = einops.rearrange(tensor, "c h w -> h w c").numpy()
-    if rgb_to_bgr:
-        color_image = cv2.cvtColor(color_image, cv2.COLOR_RGB2BGR)
-    return color_image
-
-
 # Defines a camera type
 class Camera(Protocol):
     def connect(self): ...

lerobot/common/robot_devices/motors/dynamixel.py

@@ -1,4 +1,6 @@
 import enum
+import logging
+import math
 import time
 import traceback
 from copy import deepcopy
@@ -6,17 +8,6 @@ from pathlib import Path
 
 import numpy as np
 import tqdm
-from dynamixel_sdk import (
-    COMM_SUCCESS,
-    DXL_HIBYTE,
-    DXL_HIWORD,
-    DXL_LOBYTE,
-    DXL_LOWORD,
-    GroupSyncRead,
-    GroupSyncWrite,
-    PacketHandler,
-    PortHandler,
-)
 
 from lerobot.common.robot_devices.utils import RobotDeviceAlreadyConnectedError, RobotDeviceNotConnectedError
 from lerobot.common.utils.utils import capture_timestamp_utc
@@ -27,11 +18,28 @@ TIMEOUT_MS = 1000
 
 MAX_ID_RANGE = 252
 
+# The following bounds define the lower and upper joints range (after calibration).
+# For joints in degree (i.e. revolute joints), their nominal range is [-180, 180] degrees
+# which corresponds to a half rotation on the left and half rotation on the right.
+# Some joints might require higher range, so we allow up to [-270, 270] degrees until
+# an error is raised.
+LOWER_BOUND_DEGREE = -270
+UPPER_BOUND_DEGREE = 270
+# For joints in percentage (i.e. joints that move linearly like the prismatic joint of a gripper),
+# their nominal range is [0, 100] %. For instance, for Aloha gripper, 0% is fully
+# closed, and 100% is fully open. To account for slight calibration issue, we allow up to
+# [-10, 110] until an error is raised.
+LOWER_BOUND_LINEAR = -10
+UPPER_BOUND_LINEAR = 110
+
+HALF_TURN_DEGREE = 180
+
 # https://emanual.robotis.com/docs/en/dxl/x/xl330-m077
 # https://emanual.robotis.com/docs/en/dxl/x/xl330-m288
 # https://emanual.robotis.com/docs/en/dxl/x/xl430-w250
 # https://emanual.robotis.com/docs/en/dxl/x/xm430-w350
 # https://emanual.robotis.com/docs/en/dxl/x/xm540-w270
+# https://emanual.robotis.com/docs/en/dxl/x/xc430-w150
 
 # data_name: (address, size_byte)
 X_SERIES_CONTROL_TABLE = {
@@ -109,6 +117,7 @@ MODEL_CONTROL_TABLE = {
     "xl430-w250": X_SERIES_CONTROL_TABLE,
     "xm430-w350": X_SERIES_CONTROL_TABLE,
     "xm540-w270": X_SERIES_CONTROL_TABLE,
+    "xc430-w150": X_SERIES_CONTROL_TABLE,
 }
 
 MODEL_RESOLUTION = {
@@ -118,6 +127,7 @@ MODEL_RESOLUTION = {
     "xl430-w250": 4096,
     "xm430-w350": 4096,
     "xm540-w270": 4096,
+    "xc430-w150": 4096,
 }
 
 MODEL_BAUDRATE_TABLE = {
@@ -127,44 +137,47 @@ MODEL_BAUDRATE_TABLE = {
     "xl430-w250": X_SERIES_BAUDRATE_TABLE,
     "xm430-w350": X_SERIES_BAUDRATE_TABLE,
     "xm540-w270": X_SERIES_BAUDRATE_TABLE,
+    "xc430-w150": X_SERIES_BAUDRATE_TABLE,
 }
 
 NUM_READ_RETRY = 10
 NUM_WRITE_RETRY = 10
 
 
-def convert_degrees_to_steps(degrees: float | np.ndarray, models: str | list[str]):
-    """This function convert the degree range to the step range for indicating motors rotation.
-    It assums a motor achieves a full rotation by going from -180 degree position to +180.
+def convert_degrees_to_steps(degrees: float | np.ndarray, models: str | list[str]) -> np.ndarray:
+    """This function converts the degree range to the step range for indicating motors rotation.
+    It assumes a motor achieves a full rotation by going from -180 degree position to +180.
     The motor resolution (e.g. 4096) corresponds to the number of steps needed to achieve a full rotation.
     """
-    if isinstance(degrees, float):
-        degrees = np.array(degrees)
-
     resolutions = [MODEL_RESOLUTION[model] for model in models]
     steps = degrees / 180 * np.array(resolutions) / 2
     steps = steps.astype(int)
     return steps
 
 
-def convert_to_bytes(value, bytes):
+def convert_to_bytes(value, bytes, mock=False):
+    if mock:
+        return value
+
+    import dynamixel_sdk as dxl
+
     # Note: No need to convert back into unsigned int, since this byte preprocessing
     # already handles it for us.
     if bytes == 1:
         data = [
-            DXL_LOBYTE(DXL_LOWORD(value)),
+            dxl.DXL_LOBYTE(dxl.DXL_LOWORD(value)),
         ]
     elif bytes == 2:
         data = [
-            DXL_LOBYTE(DXL_LOWORD(value)),
-            DXL_HIBYTE(DXL_LOWORD(value)),
+            dxl.DXL_LOBYTE(dxl.DXL_LOWORD(value)),
+            dxl.DXL_HIBYTE(dxl.DXL_LOWORD(value)),
         ]
     elif bytes == 4:
         data = [
-            DXL_LOBYTE(DXL_LOWORD(value)),
-            DXL_HIBYTE(DXL_LOWORD(value)),
-            DXL_LOBYTE(DXL_HIWORD(value)),
-            DXL_HIBYTE(DXL_HIWORD(value)),
+            dxl.DXL_LOBYTE(dxl.DXL_LOWORD(value)),
+            dxl.DXL_HIBYTE(dxl.DXL_LOWORD(value)),
+            dxl.DXL_LOBYTE(dxl.DXL_HIWORD(value)),
+            dxl.DXL_HIBYTE(dxl.DXL_HIWORD(value)),
         ]
     else:
         raise NotImplementedError(
@@ -250,20 +263,24 @@ class TorqueMode(enum.Enum):
     DISABLED = 0
 
 
-class OperatingMode(enum.Enum):
-    VELOCITY = 1
-    POSITION = 3
-    EXTENDED_POSITION = 4
-    CURRENT_CONTROLLED_POSITION = 5
-    PWM = 16
-    UNKNOWN = -1
-
-
 class DriveMode(enum.Enum):
     NON_INVERTED = 0
     INVERTED = 1
 
 
+class CalibrationMode(enum.Enum):
+    # Joints with rotational motions are expressed in degrees in nominal range of [-180, 180]
+    DEGREE = 0
+    # Joints with linear motions (like gripper of Aloha) are experessed in nominal range of [0, 100]
+    LINEAR = 1
+
+
+class JointOutOfRangeError(Exception):
+    def __init__(self, message="Joint is out of range"):
+        self.message = message
+        super().__init__(self.message)
+
+
 class DynamixelMotorsBus:
     # TODO(rcadene): Add a script to find the motor indices without DynamixelWizzard2
     """
@@ -310,9 +327,11 @@ class DynamixelMotorsBus:
         motors: dict[str, tuple[int, str]],
         extra_model_control_table: dict[str, list[tuple]] | None = None,
         extra_model_resolution: dict[str, int] | None = None,
+        mock=False,
     ):
         self.port = port
         self.motors = motors
+        self.mock = mock
 
         self.model_ctrl_table = deepcopy(MODEL_CONTROL_TABLE)
         if extra_model_control_table:
@@ -336,8 +355,13 @@ class DynamixelMotorsBus:
                 f"DynamixelMotorsBus({self.port}) is already connected. Do not call `motors_bus.connect()` twice."
             )
 
-        self.port_handler = PortHandler(self.port)
-        self.packet_handler = PacketHandler(PROTOCOL_VERSION)
+        if self.mock:
+            import tests.mock_dynamixel_sdk as dxl
+        else:
+            import dynamixel_sdk as dxl
+
+        self.port_handler = dxl.PortHandler(self.port)
+        self.packet_handler = dxl.PacketHandler(PROTOCOL_VERSION)
 
         try:
             if not self.port_handler.openPort():
@@ -369,10 +393,17 @@ class DynamixelMotorsBus:
             self.configure_motors()
 
     def reconnect(self):
-        self.port_handler = PortHandler(self.port)
-        self.packet_handler = PacketHandler(PROTOCOL_VERSION)
+        if self.mock:
+            import tests.mock_dynamixel_sdk as dxl
+        else:
+            import dynamixel_sdk as dxl
+
+        self.port_handler = dxl.PortHandler(self.port)
+        self.packet_handler = dxl.PacketHandler(PROTOCOL_VERSION)
+
         if not self.port_handler.openPort():
             raise OSError(f"Failed to open port '{self.port}'.")
+
         self.is_connected = True
 
     def are_motors_configured(self):
@@ -531,9 +562,22 @@ class DynamixelMotorsBus:
     def motor_indices(self) -> list[int]:
         return [idx for idx, _ in self.motors.values()]
 
-    def set_calibration(self, calibration: dict[str, tuple[int, bool]]):
+    def set_calibration(self, calibration: dict[str, list]):
         self.calibration = calibration
 
+    def apply_calibration_autocorrect(self, values: np.ndarray | list, motor_names: list[str] | None):
+        """This function applies the calibration, automatically detects out of range errors for motors values and attempts to correct.
+
+        For more info, see docstring of `apply_calibration` and `autocorrect_calibration`.
+        """
+        try:
+            values = self.apply_calibration(values, motor_names)
+        except JointOutOfRangeError as e:
+            print(e)
+            self.autocorrect_calibration(values, motor_names)
+            values = self.apply_calibration(values, motor_names)
+        return values
+
     def apply_calibration(self, values: np.ndarray | list, motor_names: list[str] | None):
         """Convert from unsigned int32 joint position range [0, 2**32[ to the universal float32 nominal degree range ]-180.0, 180.0[ with
         a "zero position" at 0 degree.
@@ -551,56 +595,205 @@ class DynamixelMotorsBus:
         if motor_names is None:
             motor_names = self.motor_names
 
-        # Convert from unsigned int32 original range [0, 2**32[ to centered signed int32 range [-2**31, 2**31[
-        values = values.astype(np.int32)
-
-        for i, name in enumerate(motor_names):
-            homing_offset, drive_mode = self.calibration[name]
-
-            # Update direction of rotation of the motor to match between leader and follower. In fact, the motor of the leader for a given joint
-            # can be assembled in an opposite direction in term of rotation than the motor of the follower on the same joint.
-            if drive_mode:
-                values[i] *= -1
-
-            # Convert from range [-2**31, 2**31[ to nominal range ]-resolution, resolution[ (e.g. ]-2048, 2048[)
-            values[i] += homing_offset
-
-        # Convert from range ]-resolution, resolution[ to the universal float32 centered degree range ]-180, 180[
+        # Convert from unsigned int32 original range [0, 2**32] to signed float32 range
         values = values.astype(np.float32)
+
         for i, name in enumerate(motor_names):
-            _, model = self.motors[name]
-            resolution = self.model_resolution[model]
-            values[i] = values[i] / (resolution // 2) * 180
+            calib_idx = self.calibration["motor_names"].index(name)
+            calib_mode = self.calibration["calib_mode"][calib_idx]
+
+            if CalibrationMode[calib_mode] == CalibrationMode.DEGREE:
+                drive_mode = self.calibration["drive_mode"][calib_idx]
+                homing_offset = self.calibration["homing_offset"][calib_idx]
+                _, model = self.motors[name]
+                resolution = self.model_resolution[model]
+
+                # Update direction of rotation of the motor to match between leader and follower.
+                # In fact, the motor of the leader for a given joint can be assembled in an
+                # opposite direction in term of rotation than the motor of the follower on the same joint.
+                if drive_mode:
+                    values[i] *= -1
+
+                # Convert from range [-2**31, 2**31] to
+                # nominal range [-resolution//2, resolution//2] (e.g. [-2048, 2048])
+                values[i] += homing_offset
+
+                # Convert from range [-resolution//2, resolution//2] to
+                # universal float32 centered degree range [-180, 180]
+                # (e.g. 2048 / (4096 // 2) * 180 = 180)
+                values[i] = values[i] / (resolution // 2) * HALF_TURN_DEGREE
+
+                if (values[i] < LOWER_BOUND_DEGREE) or (values[i] > UPPER_BOUND_DEGREE):
+                    raise JointOutOfRangeError(
+                        f"Wrong motor position range detected for {name}. "
+                        f"Expected to be in nominal range of [-{HALF_TURN_DEGREE}, {HALF_TURN_DEGREE}] degrees (a full rotation), "
+                        f"with a maximum range of [{LOWER_BOUND_DEGREE}, {UPPER_BOUND_DEGREE}] degrees to account for joints that can rotate a bit more, "
+                        f"but present value is {values[i]} degree. "
+                        "This might be due to a cable connection issue creating an artificial 360 degrees jump in motor values. "
+                        "You need to recalibrate by running: `python lerobot/scripts/control_robot.py calibrate`"
+                    )
+
+            elif CalibrationMode[calib_mode] == CalibrationMode.LINEAR:
+                start_pos = self.calibration["start_pos"][calib_idx]
+                end_pos = self.calibration["end_pos"][calib_idx]
+
+                # Rescale the present position to a nominal range [0, 100] %,
+                # useful for joints with linear motions like Aloha gripper
+                values[i] = (values[i] - start_pos) / (end_pos - start_pos) * 100
+
+                if (values[i] < LOWER_BOUND_LINEAR) or (values[i] > UPPER_BOUND_LINEAR):
+                    raise JointOutOfRangeError(
+                        f"Wrong motor position range detected for {name}. "
+                        f"Expected to be in nominal range of [0, 100] % (a full linear translation), "
+                        f"with a maximum range of [{LOWER_BOUND_LINEAR}, {UPPER_BOUND_LINEAR}] % to account for some imprecision during calibration, "
+                        f"but present value is {values[i]} %. "
+                        "This might be due to a cable connection issue creating an artificial jump in motor values. "
+                        "You need to recalibrate by running: `python lerobot/scripts/control_robot.py calibrate`"
+                    )
 
         return values
 
+    def autocorrect_calibration(self, values: np.ndarray | list, motor_names: list[str] | None):
+        """This function automatically detects issues with values of motors after calibration, and correct for these issues.
+
+        Some motors might have values outside of expected maximum bounds after calibration.
+        For instance, for a joint in degree, its value can be outside [-270, 270] degrees, which is totally unexpected given
+        a nominal range of [-180, 180] degrees, which represents half a turn to the left or right starting from zero position.
+
+        Known issues:
+        #1: Motor value randomly shifts of a full turn, caused by hardware/connection errors.
+        #2: Motor internal homing offset is shifted by a full turn, caused by using default calibration (e.g Aloha).
+        #3: motor internal homing offset is shifted by less or more than a full turn, caused by using default calibration
+            or by human error during manual calibration.
+
+        Issues #1 and #2 can be solved by shifting the calibration homing offset by a full turn.
+        Issue #3 will be visually detected by user and potentially captured by the safety feature `max_relative_target`,
+        that will slow down the motor, raise an error asking to recalibrate. Manual recalibrating will solve the issue.
+
+        Note: A full turn corresponds to 360 degrees but also to 4096 steps for a motor resolution of 4096.
+        """
+        if motor_names is None:
+            motor_names = self.motor_names
+
+        # Convert from unsigned int32 original range [0, 2**32] to signed float32 range
+        values = values.astype(np.float32)
+
+        for i, name in enumerate(motor_names):
+            calib_idx = self.calibration["motor_names"].index(name)
+            calib_mode = self.calibration["calib_mode"][calib_idx]
+
+            if CalibrationMode[calib_mode] == CalibrationMode.DEGREE:
+                drive_mode = self.calibration["drive_mode"][calib_idx]
+                homing_offset = self.calibration["homing_offset"][calib_idx]
+                _, model = self.motors[name]
+                resolution = self.model_resolution[model]
+
+                # Update direction of rotation of the motor to match between leader and follower.
+                # In fact, the motor of the leader for a given joint can be assembled in an
+                # opposite direction in term of rotation than the motor of the follower on the same joint.
+                if drive_mode:
+                    values[i] *= -1
+
+                # Convert from initial range to range [-180, 180] degrees
+                calib_val = (values[i] + homing_offset) / (resolution // 2) * HALF_TURN_DEGREE
+                in_range = (calib_val > LOWER_BOUND_DEGREE) and (calib_val < UPPER_BOUND_DEGREE)
+
+                # Solve this inequality to find the factor to shift the range into [-180, 180] degrees
+                # values[i] = (values[i] + homing_offset + resolution * factor) / (resolution // 2) * HALF_TURN_DEGREE
+                # - HALF_TURN_DEGREE <= (values[i] + homing_offset + resolution * factor) / (resolution // 2) * HALF_TURN_DEGREE <= HALF_TURN_DEGREE
+                # (- (resolution // 2) - values[i] - homing_offset) / resolution <= factor <= ((resolution // 2) - values[i] - homing_offset) / resolution
+                low_factor = (-(resolution // 2) - values[i] - homing_offset) / resolution
+                upp_factor = ((resolution // 2) - values[i] - homing_offset) / resolution
+
+            elif CalibrationMode[calib_mode] == CalibrationMode.LINEAR:
+                start_pos = self.calibration["start_pos"][calib_idx]
+                end_pos = self.calibration["end_pos"][calib_idx]
+
+                # Convert from initial range to range [0, 100] in %
+                calib_val = (values[i] - start_pos) / (end_pos - start_pos) * 100
+                in_range = (calib_val > LOWER_BOUND_LINEAR) and (calib_val < UPPER_BOUND_LINEAR)
+
+                # Solve this inequality to find the factor to shift the range into [0, 100] %
+                # values[i] = (values[i] - start_pos + resolution * factor) / (end_pos + resolution * factor - start_pos - resolution * factor) * 100
+                # values[i] = (values[i] - start_pos + resolution * factor) / (end_pos - start_pos) * 100
+                # 0 <= (values[i] - start_pos + resolution * factor) / (end_pos - start_pos) * 100 <= 100
+                # (start_pos - values[i]) / resolution <= factor <= (end_pos - values[i]) / resolution
+                low_factor = (start_pos - values[i]) / resolution
+                upp_factor = (end_pos - values[i]) / resolution
+
+            if not in_range:
+                # Get first integer between the two bounds
+                if low_factor < upp_factor:
+                    factor = math.ceil(low_factor)
+
+                    if factor > upp_factor:
+                        raise ValueError(f"No integer found between bounds [{low_factor=}, {upp_factor=}]")
+                else:
+                    factor = math.ceil(upp_factor)
+
+                    if factor > low_factor:
+                        raise ValueError(f"No integer found between bounds [{low_factor=}, {upp_factor=}]")
+
+                if CalibrationMode[calib_mode] == CalibrationMode.DEGREE:
+                    out_of_range_str = f"{LOWER_BOUND_DEGREE} < {calib_val} < {UPPER_BOUND_DEGREE} degrees"
+                    in_range_str = f"{LOWER_BOUND_DEGREE} < {calib_val} < {UPPER_BOUND_DEGREE} degrees"
+                elif CalibrationMode[calib_mode] == CalibrationMode.LINEAR:
+                    out_of_range_str = f"{LOWER_BOUND_LINEAR} < {calib_val} < {UPPER_BOUND_LINEAR} %"
+                    in_range_str = f"{LOWER_BOUND_LINEAR} < {calib_val} < {UPPER_BOUND_LINEAR} %"
+
+                logging.warning(
+                    f"Auto-correct calibration of motor '{name}' by shifting value by {abs(factor)} full turns, "
+                    f"from '{out_of_range_str}' to '{in_range_str}'."
+                )
+
+                # A full turn corresponds to 360 degrees but also to 4096 steps for a motor resolution of 4096.
+                self.calibration["homing_offset"][calib_idx] += resolution * factor
+
     def revert_calibration(self, values: np.ndarray | list, motor_names: list[str] | None):
         """Inverse of `apply_calibration`."""
         if motor_names is None:
             motor_names = self.motor_names
 
-        # Convert from the universal float32 centered degree range ]-180, 180[ to resolution range ]-resolution, resolution[
         for i, name in enumerate(motor_names):
-            _, model = self.motors[name]
-            resolution = self.model_resolution[model]
-            values[i] = values[i] / 180 * (resolution // 2)
+            calib_idx = self.calibration["motor_names"].index(name)
+            calib_mode = self.calibration["calib_mode"][calib_idx]
+
+            if CalibrationMode[calib_mode] == CalibrationMode.DEGREE:
+                drive_mode = self.calibration["drive_mode"][calib_idx]
+                homing_offset = self.calibration["homing_offset"][calib_idx]
+                _, model = self.motors[name]
+                resolution = self.model_resolution[model]
+
+                # Convert from nominal 0-centered degree range [-180, 180] to
+                # 0-centered resolution range (e.g. [-2048, 2048] for resolution=4096)
+                values[i] = values[i] / HALF_TURN_DEGREE * (resolution // 2)
+
+                # Substract the homing offsets to come back to actual motor range of values
+                # which can be arbitrary.
+                values[i] -= homing_offset
+
+                # Remove drive mode, which is the rotation direction of the motor, to come back to
+                # actual motor rotation direction which can be arbitrary.
+                if drive_mode:
+                    values[i] *= -1
+
+            elif CalibrationMode[calib_mode] == CalibrationMode.LINEAR:
+                start_pos = self.calibration["start_pos"][calib_idx]
+                end_pos = self.calibration["end_pos"][calib_idx]
+
+                # Convert from nominal lnear range of [0, 100] % to
+                # actual motor range of values which can be arbitrary.
+                values[i] = values[i] / 100 * (end_pos - start_pos) + start_pos
 
         values = np.round(values).astype(np.int32)
-
-        # Convert from nominal range ]-resolution, resolution[ to centered signed int32 range [-2**31, 2**31[
-        for i, name in enumerate(motor_names):
-            homing_offset, drive_mode = self.calibration[name]
-            values[i] -= homing_offset
-
-            # Update direction of rotation of the motor that was matching between leader and follower to their original direction.
-            # In fact, the motor of the leader for a given joint can be assembled in an opposite direction in term of rotation
-            # than the motor of the follower on the same joint.
-            if drive_mode:
-                values[i] *= -1
-
         return values
 
     def _read_with_motor_ids(self, motor_models, motor_ids, data_name):
+        if self.mock:
+            import tests.mock_dynamixel_sdk as dxl
+        else:
+            import dynamixel_sdk as dxl
+
         return_list = True
         if not isinstance(motor_ids, list):
             return_list = False
@@ -608,12 +801,12 @@ class DynamixelMotorsBus:
 
         assert_same_address(self.model_ctrl_table, self.motor_models, data_name)
         addr, bytes = self.model_ctrl_table[motor_models[0]][data_name]
-        group = GroupSyncRead(self.port_handler, self.packet_handler, addr, bytes)
+        group = dxl.GroupSyncRead(self.port_handler, self.packet_handler, addr, bytes)
         for idx in motor_ids:
             group.addParam(idx)
 
         comm = group.txRxPacket()
-        if comm != COMM_SUCCESS:
+        if comm != dxl.COMM_SUCCESS:
             raise ConnectionError(
                 f"Read failed due to communication error on port {self.port_handler.port_name} for indices {motor_ids}: "
                 f"{self.packet_handler.getTxRxResult(comm)}"
@@ -637,6 +830,11 @@ class DynamixelMotorsBus:
 
         start_time = time.perf_counter()
 
+        if self.mock:
+            import tests.mock_dynamixel_sdk as dxl
+        else:
+            import dynamixel_sdk as dxl
+
         if motor_names is None:
             motor_names = self.motor_names
 
@@ -656,16 +854,18 @@ class DynamixelMotorsBus:
 
         if data_name not in self.group_readers:
             # create new group reader
-            self.group_readers[group_key] = GroupSyncRead(self.port_handler, self.packet_handler, addr, bytes)
+            self.group_readers[group_key] = dxl.GroupSyncRead(
+                self.port_handler, self.packet_handler, addr, bytes
+            )
             for idx in motor_ids:
                 self.group_readers[group_key].addParam(idx)
 
         for _ in range(NUM_READ_RETRY):
             comm = self.group_readers[group_key].txRxPacket()
-            if comm == COMM_SUCCESS:
+            if comm == dxl.COMM_SUCCESS:
                 break
 
-        if comm != COMM_SUCCESS:
+        if comm != dxl.COMM_SUCCESS:
             raise ConnectionError(
                 f"Read failed due to communication error on port {self.port} for group_key {group_key}: "
                 f"{self.packet_handler.getTxRxResult(comm)}"
@@ -683,19 +883,7 @@ class DynamixelMotorsBus:
             values = values.astype(np.int32)
 
         if data_name in CALIBRATION_REQUIRED and self.calibration is not None:
-            values = self.apply_calibration(values, motor_names)
-
-            # We expect our motors to stay in a nominal range of [-180, 180] degrees
-            # which corresponds to a half turn rotation.
-            # However, some motors can turn a bit more, hence we extend the nominal range to [-270, 270]
-            # which is less than a full 360 degree rotation.
-            if not np.all((values > -270) & (values < 270)):
-                raise ValueError(
-                    f"Wrong motor position range detected. "
-                    f"Expected to be in [-270, +270] but in [{values.min()}, {values.max()}]. "
-                    "This might be due to a cable connection issue creating an artificial 360 degrees jump in motor values. "
-                    "You need to recalibrate by running: `python lerobot/scripts/control_robot.py calibrate`"
-                )
+            values = self.apply_calibration_autocorrect(values, motor_names)
 
         # log the number of seconds it took to read the data from the motors
         delta_ts_name = get_log_name("delta_timestamp_s", "read", data_name, motor_names)
@@ -708,6 +896,11 @@ class DynamixelMotorsBus:
         return values
 
     def _write_with_motor_ids(self, motor_models, motor_ids, data_name, values):
+        if self.mock:
+            import tests.mock_dynamixel_sdk as dxl
+        else:
+            import dynamixel_sdk as dxl
+
         if not isinstance(motor_ids, list):
             motor_ids = [motor_ids]
         if not isinstance(values, list):
@@ -715,13 +908,13 @@ class DynamixelMotorsBus:
 
         assert_same_address(self.model_ctrl_table, motor_models, data_name)
         addr, bytes = self.model_ctrl_table[motor_models[0]][data_name]
-        group = GroupSyncWrite(self.port_handler, self.packet_handler, addr, bytes)
+        group = dxl.GroupSyncWrite(self.port_handler, self.packet_handler, addr, bytes)
         for idx, value in zip(motor_ids, values, strict=True):
-            data = convert_to_bytes(value, bytes)
+            data = convert_to_bytes(value, bytes, self.mock)
             group.addParam(idx, data)
 
         comm = group.txPacket()
-        if comm != COMM_SUCCESS:
+        if comm != dxl.COMM_SUCCESS:
             raise ConnectionError(
                 f"Write failed due to communication error on port {self.port_handler.port_name} for indices {motor_ids}: "
                 f"{self.packet_handler.getTxRxResult(comm)}"
@@ -735,6 +928,11 @@ class DynamixelMotorsBus:
 
         start_time = time.perf_counter()
 
+        if self.mock:
+            import tests.mock_dynamixel_sdk as dxl
+        else:
+            import dynamixel_sdk as dxl
+
         if motor_names is None:
             motor_names = self.motor_names
 
@@ -764,19 +962,19 @@ class DynamixelMotorsBus:
 
         init_group = data_name not in self.group_readers
         if init_group:
-            self.group_writers[group_key] = GroupSyncWrite(
+            self.group_writers[group_key] = dxl.GroupSyncWrite(
                 self.port_handler, self.packet_handler, addr, bytes
             )
 
         for idx, value in zip(motor_ids, values, strict=True):
-            data = convert_to_bytes(value, bytes)
+            data = convert_to_bytes(value, bytes, self.mock)
             if init_group:
                 self.group_writers[group_key].addParam(idx, data)
             else:
                 self.group_writers[group_key].changeParam(idx, data)
 
         comm = self.group_writers[group_key].txPacket()
-        if comm != COMM_SUCCESS:
+        if comm != dxl.COMM_SUCCESS:
             raise ConnectionError(
                 f"Write failed due to communication error on port {self.port} for group_key {group_key}: "
                 f"{self.packet_handler.getTxRxResult(comm)}"

lerobot/common/robot_devices/robots/factory.py

@@ -1,7 +1,9 @@
 import hydra
 from omegaconf import DictConfig
 
+from lerobot.common.robot_devices.robots.utils import Robot
 
-def make_robot(cfg: DictConfig):
+
+def make_robot(cfg: DictConfig) -> Robot:
     robot = hydra.utils.instantiate(cfg)
     return robot

lerobot/common/robot_devices/robots/manipulator.py

@@ -1,6 +1,7 @@
+import json
 import logging
-import pickle
 import time
+import warnings
 from dataclasses import dataclass, field, replace
 from pathlib import Path
 from typing import Sequence
@@ -10,11 +11,12 @@ import torch
 
 from lerobot.common.robot_devices.cameras.utils import Camera
 from lerobot.common.robot_devices.motors.dynamixel import (
-    OperatingMode,
+    CalibrationMode,
     TorqueMode,
     convert_degrees_to_steps,
 )
 from lerobot.common.robot_devices.motors.utils import MotorsBus
+from lerobot.common.robot_devices.robots.utils import get_arm_id
 from lerobot.common.robot_devices.utils import RobotDeviceAlreadyConnectedError, RobotDeviceNotConnectedError
 
 ########################################################################
@@ -25,7 +27,8 @@ URL_TEMPLATE = (
     "https://raw.githubusercontent.com/huggingface/lerobot/main/media/{robot}/{arm}_{position}.webp"
 )
 
-# In nominal degree range ]-180, +180[
+# The following positions are provided in nominal degree range ]-180, +180[
+# For more info on these constants, see comments in the code where they get used.
 ZERO_POSITION_DEGREE = 0
 ROTATED_POSITION_DEGREE = 90
 
@@ -45,27 +48,13 @@ def apply_drive_mode(position, drive_mode):
     return position
 
 
-def reset_torque_mode(arm: MotorsBus):
-    # To be configured, all servos must be in "torque disable" mode
-    arm.write("Torque_Enable", TorqueMode.DISABLED.value)
-
-    # Use 'extended position mode' for all motors except gripper, because in joint mode the servos can't
-    # rotate more than 360 degrees (from 0 to 4095) And some mistake can happen while assembling the arm,
-    # you could end up with a servo with a position 0 or 4095 at a crucial point See [
-    # https://emanual.robotis.com/docs/en/dxl/x/x_series/#operating-mode11]
-    all_motors_except_gripper = [name for name in arm.motor_names if name != "gripper"]
-    if len(all_motors_except_gripper) > 0:
-        arm.write("Operating_Mode", OperatingMode.EXTENDED_POSITION.value, all_motors_except_gripper)
-
-    # Use 'position control current based' for gripper to be limited by the limit of the current.
-    # For the follower gripper, it means it can grasp an object without forcing too much even tho,
-    # it's goal position is a complete grasp (both gripper fingers are ordered to join and reach a touch).
-    # For the leader gripper, it means we can use it as a physical trigger, since we can force with our finger
-    # to make it move, and it will move back to its original target position when we release the force.
-    arm.write("Operating_Mode", OperatingMode.CURRENT_CONTROLLED_POSITION.value, "gripper")
+def compute_nearest_rounded_position(position, models):
+    delta_turn = convert_degrees_to_steps(ROTATED_POSITION_DEGREE, models)
+    nearest_pos = np.round(position.astype(float) / delta_turn) * delta_turn
+    return nearest_pos.astype(position.dtype)
 
 
-def run_arm_calibration(arm: MotorsBus, name: str, arm_type: str):
+def run_arm_calibration(arm: MotorsBus, robot_type: str, arm_name: str, arm_type: str):
     """This function ensures that a neural network trained on data collected on a given robot
     can work on another robot. For instance before calibration, setting a same goal position
     for each motor of two different robots will get two very different positions. But after calibration,
@@ -84,38 +73,27 @@ def run_arm_calibration(arm: MotorsBus, name: str, arm_type: str):
 
     Example of usage:
     ```python
-    run_arm_calibration(arm, "left", "follower")
+    run_arm_calibration(arm, "koch", "left", "follower")
     ```
     """
-    reset_torque_mode(arm)
+    if (arm.read("Torque_Enable") != TorqueMode.DISABLED.value).any():
+        raise ValueError("To run calibration, the torque must be disabled on all motors.")
 
-    print(f"\nRunning calibration of {name} {arm_type}...")
+    print(f"\nRunning calibration of {robot_type} {arm_name} {arm_type}...")
 
     print("\nMove arm to zero position")
-    print("See: " + URL_TEMPLATE.format(robot="koch", arm=arm_type, position="zero"))
+    print("See: " + URL_TEMPLATE.format(robot=robot_type, arm=arm_type, position="zero"))
     input("Press Enter to continue...")
 
-    # We arbitrarely choosed our zero target position to be a straight horizontal position with gripper upwards and closed.
+    # We arbitrarily chose our zero target position to be a straight horizontal position with gripper upwards and closed.
     # It is easy to identify and all motors are in a "quarter turn" position. Once calibration is done, this position will
-    # corresponds to every motor angle being 0. If you set all 0 as Goal Position, the arm will move in this position.
-    zero_position = convert_degrees_to_steps(ZERO_POSITION_DEGREE, arm.motor_models)
-
-    def _compute_nearest_rounded_position(position, models):
-        # TODO(rcadene): Rework this function since some motors cant physically rotate a quarter turn
-        # (e.g. the gripper of Aloha arms can only rotate ~50 degree)
-        quarter_turn_degree = 90
-        quarter_turn = convert_degrees_to_steps(quarter_turn_degree, models)
-        nearest_pos = np.round(position.astype(float) / quarter_turn) * quarter_turn
-        return nearest_pos.astype(position.dtype)
+    # correspond to every motor angle being 0. If you set all 0 as Goal Position, the arm will move in this position.
+    zero_target_pos = convert_degrees_to_steps(ZERO_POSITION_DEGREE, arm.motor_models)
 
     # Compute homing offset so that `present_position + homing_offset ~= target_position`.
-    position = arm.read("Present_Position")
-    position = _compute_nearest_rounded_position(position, arm.motor_models)
-    homing_offset = zero_position - position
-
-    print("\nMove arm to rotated target position")
-    print("See: " + URL_TEMPLATE.format(robot="koch", arm=arm_type, position="rotated"))
-    input("Press Enter to continue...")
+    zero_pos = arm.read("Present_Position")
+    zero_nearest_pos = compute_nearest_rounded_position(zero_pos, arm.motor_models)
+    homing_offset = zero_target_pos - zero_nearest_pos
 
     # The rotated target position corresponds to a rotation of a quarter turn from the zero position.
     # This allows to identify the rotation direction of each motor.
@@ -124,44 +102,83 @@ def run_arm_calibration(arm: MotorsBus, name: str, arm_type: str):
     # Sometimes, there is only one possible rotation direction. For instance, if the gripper is closed, there is only one direction which
     # corresponds to opening the gripper. When the rotation direction is ambiguous, we arbitrarely rotate clockwise from the point of view
     # of the previous motor in the kinetic chain.
-    rotated_position = convert_degrees_to_steps(ROTATED_POSITION_DEGREE, arm.motor_models)
+    print("\nMove arm to rotated target position")
+    print("See: " + URL_TEMPLATE.format(robot=robot_type, arm=arm_type, position="rotated"))
+    input("Press Enter to continue...")
+
+    rotated_target_pos = convert_degrees_to_steps(ROTATED_POSITION_DEGREE, arm.motor_models)
 
     # Find drive mode by rotating each motor by a quarter of a turn.
     # Drive mode indicates if the motor rotation direction should be inverted (=1) or not (=0).
-    position = arm.read("Present_Position")
-    position += homing_offset
-    position = _compute_nearest_rounded_position(position, arm.motor_models)
-    drive_mode = (position != rotated_position).astype(np.int32)
+    rotated_pos = arm.read("Present_Position")
+    drive_mode = (rotated_pos < zero_pos).astype(np.int32)
 
     # Re-compute homing offset to take into account drive mode
-    position = arm.read("Present_Position")
-    position = apply_drive_mode(position, drive_mode)
-    position = _compute_nearest_rounded_position(position, arm.motor_models)
-    homing_offset = rotated_position - position
+    rotated_drived_pos = apply_drive_mode(rotated_pos, drive_mode)
+    rotated_nearest_pos = compute_nearest_rounded_position(rotated_drived_pos, arm.motor_models)
+    homing_offset = rotated_target_pos - rotated_nearest_pos
 
     print("\nMove arm to rest position")
-    print("See: " + URL_TEMPLATE.format(robot="koch", arm=arm_type, position="rest"))
+    print("See: " + URL_TEMPLATE.format(robot=robot_type, arm=arm_type, position="rest"))
     input("Press Enter to continue...")
     print()
 
-    return homing_offset, drive_mode
+    # Joints with rotational motions are expressed in degrees in nominal range of [-180, 180]
+    calib_mode = [CalibrationMode.DEGREE.name] * len(arm.motor_names)
+
+    # TODO(rcadene): make type of joints (DEGREE or LINEAR) configurable from yaml?
+    if robot_type == "aloha" and "gripper" in arm.motor_names:
+        # Joints with linear motions (like gripper of Aloha) are experessed in nominal range of [0, 100]
+        calib_idx = arm.motor_names.index("gripper")
+        calib_mode[calib_idx] = CalibrationMode.LINEAR.name
+
+    calib_data = {
+        "homing_offset": homing_offset.tolist(),
+        "drive_mode": drive_mode.tolist(),
+        "start_pos": zero_pos.tolist(),
+        "end_pos": rotated_pos.tolist(),
+        "calib_mode": calib_mode,
+        "motor_names": arm.motor_names,
+    }
+    return calib_data
+
+
+def ensure_safe_goal_position(
+    goal_pos: torch.Tensor, present_pos: torch.Tensor, max_relative_target: float | list[float]
+):
+    # Cap relative action target magnitude for safety.
+    diff = goal_pos - present_pos
+    max_relative_target = torch.tensor(max_relative_target)
+    safe_diff = torch.minimum(diff, max_relative_target)
+    safe_diff = torch.maximum(safe_diff, -max_relative_target)
+    safe_goal_pos = present_pos + safe_diff
+
+    if not torch.allclose(goal_pos, safe_goal_pos):
+        logging.warning(
+            "Relative goal position magnitude had to be clamped to be safe.\n"
+            f"  requested relative goal position target: {diff}\n"
+            f"    clamped relative goal position target: {safe_diff}"
+        )
+
+    return safe_goal_pos
 
 
 ########################################################################
-# Alexander Koch robot arm
+# Manipulator robot
 ########################################################################
 
 
 @dataclass
-class KochRobotConfig:
+class ManipulatorRobotConfig:
     """
     Example of usage:
     ```python
-    KochRobotConfig()
+    ManipulatorRobotConfig()
     ```
     """
 
     # Define all components of the robot
+    robot_type: str | None = None
     leader_arms: dict[str, MotorsBus] = field(default_factory=lambda: {})
     follower_arms: dict[str, MotorsBus] = field(default_factory=lambda: {})
     cameras: dict[str, Camera] = field(default_factory=lambda: {})
@@ -191,14 +208,15 @@ class KochRobotConfig:
         super().__setattr__(prop, val)
 
 
-class KochRobot:
+class ManipulatorRobot:
     # TODO(rcadene): Implement force feedback
-    """This class allows to control any Koch robot of various number of motors.
+    """This class allows to control any manipulator robot of various number of motors.
 
-    A few versions are available:
-    - [Koch v1.0](https://github.com/AlexanderKoch-Koch/low_cost_robot), with and without the wrist-to-elbow expansion, which was developed
-    by Alexander Koch from [Tau Robotics](https://tau-robotics.com): [Github for sourcing and assembly](
-    - [Koch v1.1])https://github.com/jess-moss/koch-v1-1), which was developed by Jess Moss.
+    Non exaustive list of robots:
+    - [Koch v1.0](https://github.com/AlexanderKoch-Koch/low_cost_robot), with and without the wrist-to-elbow expansion, developed
+    by Alexander Koch from [Tau Robotics](https://tau-robotics.com)
+    - [Koch v1.1](https://github.com/jess-moss/koch-v1-1) developed by Jess Moss
+    - [Aloha](https://www.trossenrobotics.com/aloha-kits) developed by Trossen Robotics
 
     Example of highest frequency teleoperation without camera:
     ```python
@@ -231,7 +249,9 @@ class KochRobot:
             },
         ),
     }
-    robot = KochRobot(
+    robot = ManipulatorRobot(
+        robot_type="koch",
+        calibration_dir=".cache/calibration/koch",
         leader_arms=leader_arms,
         follower_arms=follower_arms,
     )
@@ -246,7 +266,9 @@ class KochRobot:
     Example of highest frequency data collection without camera:
     ```python
     # Assumes leader and follower arms have been instantiated already (see first example)
-    robot = KochRobot(
+    robot = ManipulatorRobot(
+        robot_type="koch",
+        calibration_dir=".cache/calibration/koch",
         leader_arms=leader_arms,
         follower_arms=follower_arms,
     )
@@ -267,7 +289,9 @@ class KochRobot:
     }
 
     # Assumes leader and follower arms have been instantiated already (see first example)
-    robot = KochRobot(
+    robot = ManipulatorRobot(
+        robot_type="koch",
+        calibration_dir=".cache/calibration/koch",
         leader_arms=leader_arms,
         follower_arms=follower_arms,
         cameras=cameras,
@@ -280,7 +304,9 @@ class KochRobot:
     Example of controlling the robot with a policy (without running multiple policies in parallel to ensure highest frequency):
     ```python
     # Assumes leader and follower arms + cameras have been instantiated already (see previous example)
-    robot = KochRobot(
+    robot = ManipulatorRobot(
+        robot_type="koch",
+        calibration_dir=".cache/calibration/koch",
         leader_arms=leader_arms,
         follower_arms=follower_arms,
         cameras=cameras,
@@ -306,16 +332,17 @@ class KochRobot:
 
     def __init__(
         self,
-        config: KochRobotConfig | None = None,
-        calibration_path: Path = ".cache/calibration/koch.pkl",
+        config: ManipulatorRobotConfig | None = None,
+        calibration_dir: Path = ".cache/calibration/koch",
         **kwargs,
     ):
         if config is None:
-            config = KochRobotConfig()
+            config = ManipulatorRobotConfig()
         # Overwrite config arguments using kwargs
         self.config = replace(config, **kwargs)
-        self.calibration_path = Path(calibration_path)
+        self.calibration_dir = Path(calibration_dir)
 
+        self.robot_type = self.config.robot_type
         self.leader_arms = self.config.leader_arms
         self.follower_arms = self.config.follower_arms
         self.cameras = self.config.cameras
@@ -325,12 +352,12 @@ class KochRobot:
     def connect(self):
         if self.is_connected:
             raise RobotDeviceAlreadyConnectedError(
-                "KochRobot is already connected. Do not run `robot.connect()` twice."
+                "ManipulatorRobot is already connected. Do not run `robot.connect()` twice."
             )
 
         if not self.leader_arms and not self.follower_arms and not self.cameras:
             raise ValueError(
-                "KochRobot doesn't have any device to connect. See example of usage in docstring of the class."
+                "ManipulatorRobot doesn't have any device to connect. See example of usage in docstring of the class."
             )
 
         # Connect the arms
@@ -340,38 +367,22 @@ class KochRobot:
             print(f"Connecting {name} leader arm.")
             self.leader_arms[name].connect()
 
-        # Reset the arms and load or run calibration
-        if self.calibration_path.exists():
-            # Reset all arms before setting calibration
-            for name in self.follower_arms:
-                reset_torque_mode(self.follower_arms[name])
-            for name in self.leader_arms:
-                reset_torque_mode(self.leader_arms[name])
-
-            with open(self.calibration_path, "rb") as f:
-                calibration = pickle.load(f)
-        else:
-            print(f"Missing calibration file '{self.calibration_path}'. Starting calibration precedure.")
-            # Run calibration process which begins by reseting all arms
-            calibration = self.run_calibration()
-
-            print(f"Calibration is done! Saving calibration file '{self.calibration_path}'")
-            self.calibration_path.parent.mkdir(parents=True, exist_ok=True)
-            with open(self.calibration_path, "wb") as f:
-                pickle.dump(calibration, f)
-
-        # Set calibration
+        # We assume that at connection time, arms are in a rest position, and torque can
+        # be safely disabled to run calibration and/or set robot preset configurations.
         for name in self.follower_arms:
-            self.follower_arms[name].set_calibration(calibration[f"follower_{name}"])
+            self.follower_arms[name].write("Torque_Enable", TorqueMode.DISABLED.value)
         for name in self.leader_arms:
-            self.leader_arms[name].set_calibration(calibration[f"leader_{name}"])
+            self.leader_arms[name].write("Torque_Enable", TorqueMode.DISABLED.value)
 
-        # Set better PID values to close the gap between recored states and actions
-        # TODO(rcadene): Implement an automatic procedure to set optimial PID values for each motor
-        for name in self.follower_arms:
-            self.follower_arms[name].write("Position_P_Gain", 1500, "elbow_flex")
-            self.follower_arms[name].write("Position_I_Gain", 0, "elbow_flex")
-            self.follower_arms[name].write("Position_D_Gain", 600, "elbow_flex")
+        self.activate_calibration()
+
+        # Set robot preset (e.g. torque in leader gripper for Koch v1.1)
+        if self.robot_type == "koch":
+            self.set_koch_robot_preset()
+        elif self.robot_type == "aloha":
+            self.set_aloha_robot_preset()
+        else:
+            warnings.warn(f"No preset found for robot type: {self.robot_type}", stacklevel=1)
 
         # Enable torque on all motors of the follower arms
         for name in self.follower_arms:
@@ -391,31 +402,132 @@ class KochRobot:
 
         self.is_connected = True
 
-    def run_calibration(self):
-        calibration = {}
+    def activate_calibration(self):
+        """After calibration all motors function in human interpretable ranges.
+        Rotations are expressed in degrees in nominal range of [-180, 180],
+        and linear motions (like gripper of Aloha) in nominal range of [0, 100].
+        """
+
+        def load_or_run_calibration_(name, arm, arm_type):
+            arm_id = get_arm_id(name, arm_type)
+            arm_calib_path = self.calibration_dir / f"{arm_id}.json"
+
+            if arm_calib_path.exists():
+                with open(arm_calib_path) as f:
+                    calibration = json.load(f)
+            else:
+                print(f"Missing calibration file '{arm_calib_path}'")
+                calibration = run_arm_calibration(arm, self.robot_type, name, arm_type)
+
+                print(f"Calibration is done! Saving calibration file '{arm_calib_path}'")
+                arm_calib_path.parent.mkdir(parents=True, exist_ok=True)
+                with open(arm_calib_path, "w") as f:
+                    json.dump(calibration, f)
+
+            return calibration
+
+        for name, arm in self.follower_arms.items():
+            calibration = load_or_run_calibration_(name, arm, "follower")
+            arm.set_calibration(calibration)
+        for name, arm in self.leader_arms.items():
+            calibration = load_or_run_calibration_(name, arm, "leader")
+            arm.set_calibration(calibration)
+
+    def set_koch_robot_preset(self):
+        def set_operating_mode_(arm):
+            if (arm.read("Torque_Enable") != TorqueMode.DISABLED.value).any():
+                raise ValueError("To run set robot preset, the torque must be disabled on all motors.")
+
+            # Use 'extended position mode' for all motors except gripper, because in joint mode the servos can't
+            # rotate more than 360 degrees (from 0 to 4095) And some mistake can happen while assembling the arm,
+            # you could end up with a servo with a position 0 or 4095 at a crucial point See [
+            # https://emanual.robotis.com/docs/en/dxl/x/x_series/#operating-mode11]
+            all_motors_except_gripper = [name for name in arm.motor_names if name != "gripper"]
+            if len(all_motors_except_gripper) > 0:
+                # 4 corresponds to Extended Position on Koch motors
+                arm.write("Operating_Mode", 4, all_motors_except_gripper)
+
+            # Use 'position control current based' for gripper to be limited by the limit of the current.
+            # For the follower gripper, it means it can grasp an object without forcing too much even tho,
+            # it's goal position is a complete grasp (both gripper fingers are ordered to join and reach a touch).
+            # For the leader gripper, it means we can use it as a physical trigger, since we can force with our finger
+            # to make it move, and it will move back to its original target position when we release the force.
+            # 5 corresponds to Current Controlled Position on Koch gripper motors "xl330-m077, xl330-m288"
+            arm.write("Operating_Mode", 5, "gripper")
 
         for name in self.follower_arms:
-            homing_offset, drive_mode = run_arm_calibration(self.follower_arms[name], name, "follower")
+            set_operating_mode_(self.follower_arms[name])
 
-            calibration[f"follower_{name}"] = {}
-            for idx, motor_name in enumerate(self.follower_arms[name].motor_names):
-                calibration[f"follower_{name}"][motor_name] = (homing_offset[idx], drive_mode[idx])
+            # Set better PID values to close the gap between recorded states and actions
+            # TODO(rcadene): Implement an automatic procedure to set optimial PID values for each motor
+            self.follower_arms[name].write("Position_P_Gain", 1500, "elbow_flex")
+            self.follower_arms[name].write("Position_I_Gain", 0, "elbow_flex")
+            self.follower_arms[name].write("Position_D_Gain", 600, "elbow_flex")
+
+        if self.config.gripper_open_degree is not None:
+            for name in self.leader_arms:
+                set_operating_mode_(self.leader_arms[name])
+
+                # Enable torque on the gripper of the leader arms, and move it to 45 degrees,
+                # so that we can use it as a trigger to close the gripper of the follower arms.
+                self.leader_arms[name].write("Torque_Enable", 1, "gripper")
+                self.leader_arms[name].write("Goal_Position", self.config.gripper_open_degree, "gripper")
+
+    def set_aloha_robot_preset(self):
+        def set_shadow_(arm):
+            # Set secondary/shadow ID for shoulder and elbow. These joints have two motors.
+            # As a result, if only one of them is required to move to a certain position,
+            # the other will follow. This is to avoid breaking the motors.
+            if "shoulder_shadow" in arm.motor_names:
+                shoulder_idx = arm.read("ID", "shoulder")
+                arm.write("Secondary_ID", shoulder_idx, "shoulder_shadow")
+
+            if "elbow_shadow" in arm.motor_names:
+                elbow_idx = arm.read("ID", "elbow")
+                arm.write("Secondary_ID", elbow_idx, "elbow_shadow")
+
+        for name in self.follower_arms:
+            set_shadow_(self.follower_arms[name])
 
         for name in self.leader_arms:
-            homing_offset, drive_mode = run_arm_calibration(self.leader_arms[name], name, "leader")
+            set_shadow_(self.leader_arms[name])
 
-            calibration[f"leader_{name}"] = {}
-            for idx, motor_name in enumerate(self.leader_arms[name].motor_names):
-                calibration[f"leader_{name}"][motor_name] = (homing_offset[idx], drive_mode[idx])
+        for name in self.follower_arms:
+            # Set a velocity limit of 131 as advised by Trossen Robotics
+            self.follower_arms[name].write("Velocity_Limit", 131)
 
-        return calibration
+            # Use 'extended position mode' for all motors except gripper, because in joint mode the servos can't
+            # rotate more than 360 degrees (from 0 to 4095) And some mistake can happen while assembling the arm,
+            # you could end up with a servo with a position 0 or 4095 at a crucial point See [
+            # https://emanual.robotis.com/docs/en/dxl/x/x_series/#operating-mode11]
+            all_motors_except_gripper = [
+                name for name in self.follower_arms[name].motor_names if name != "gripper"
+            ]
+            if len(all_motors_except_gripper) > 0:
+                # 4 corresponds to Extended Position on Aloha motors
+                self.follower_arms[name].write("Operating_Mode", 4, all_motors_except_gripper)
+
+            # Use 'position control current based' for follower gripper to be limited by the limit of the current.
+            # It can grasp an object without forcing too much even tho,
+            # it's goal position is a complete grasp (both gripper fingers are ordered to join and reach a touch).
+            # 5 corresponds to Current Controlled Position on Aloha gripper follower "xm430-w350"
+            self.follower_arms[name].write("Operating_Mode", 5, "gripper")
+
+            # Note: We can't enable torque on the leader gripper since "xc430-w150" doesn't have
+            # a Current Controlled Position mode.
+
+        if self.config.gripper_open_degree is not None:
+            warnings.warn(
+                f"`gripper_open_degree` is set to {self.config.gripper_open_degree}, but None is expected for Aloha instead",
+                stacklevel=1,
+            )
 
     def teleop_step(
         self, record_data=False
     ) -> None | tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
         if not self.is_connected:
             raise RobotDeviceNotConnectedError(
-                "KochRobot is not connected. You need to run `robot.connect()`."
+                "ManipulatorRobot is not connected. You need to run `robot.connect()`."
             )
 
         # Prepare to assign the position of the leader to the follower
@@ -423,16 +535,27 @@ class KochRobot:
         for name in self.leader_arms:
             before_lread_t = time.perf_counter()
             leader_pos[name] = self.leader_arms[name].read("Present_Position")
+            leader_pos[name] = torch.from_numpy(leader_pos[name])
             self.logs[f"read_leader_{name}_pos_dt_s"] = time.perf_counter() - before_lread_t
 
+        # Send goal position to the follower
         follower_goal_pos = {}
-        for name in self.leader_arms:
-            follower_goal_pos[name] = leader_pos[name]
-
-        # Send action
         for name in self.follower_arms:
             before_fwrite_t = time.perf_counter()
-            self.send_action(torch.tensor(follower_goal_pos[name]), [name])
+            goal_pos = leader_pos[name]
+
+            # Cap goal position when too far away from present position.
+            # Slower fps expected due to reading from the follower.
+            if self.config.max_relative_target is not None:
+                present_pos = self.follower_arms[name].read("Present_Position")
+                present_pos = torch.from_numpy(present_pos)
+                goal_pos = ensure_safe_goal_position(goal_pos, present_pos, self.config.max_relative_target)
+
+            # Used when record_data=True
+            follower_goal_pos[name] = goal_pos
+
+            goal_pos = goal_pos.numpy().astype(np.int32)
+            self.follower_arms[name].write("Goal_Position", goal_pos)
             self.logs[f"write_follower_{name}_goal_pos_dt_s"] = time.perf_counter() - before_fwrite_t
 
         # Early exit when recording data is not requested
@@ -445,6 +568,7 @@ class KochRobot:
         for name in self.follower_arms:
             before_fread_t = time.perf_counter()
             follower_pos[name] = self.follower_arms[name].read("Present_Position")
+            follower_pos[name] = torch.from_numpy(follower_pos[name])
             self.logs[f"read_follower_{name}_pos_dt_s"] = time.perf_counter() - before_fread_t
 
         # Create state by concatenating follower current position
@@ -452,29 +576,30 @@ class KochRobot:
         for name in self.follower_arms:
             if name in follower_pos:
                 state.append(follower_pos[name])
-        state = np.concatenate(state)
+        state = torch.cat(state)
 
         # Create action by concatenating follower goal position
         action = []
         for name in self.follower_arms:
             if name in follower_goal_pos:
                 action.append(follower_goal_pos[name])
-        action = np.concatenate(action)
+        action = torch.cat(action)
 
         # Capture images from cameras
         images = {}
         for name in self.cameras:
             before_camread_t = time.perf_counter()
             images[name] = self.cameras[name].async_read()
+            images[name] = torch.from_numpy(images[name])
             self.logs[f"read_camera_{name}_dt_s"] = self.cameras[name].logs["delta_timestamp_s"]
             self.logs[f"async_read_camera_{name}_dt_s"] = time.perf_counter() - before_camread_t
 
-        # Populate output dictionnaries and format to pytorch
+        # Populate output dictionnaries
         obs_dict, action_dict = {}, {}
-        obs_dict["observation.state"] = torch.from_numpy(state)
-        action_dict["action"] = torch.from_numpy(action)
+        obs_dict["observation.state"] = state
+        action_dict["action"] = action
         for name in self.cameras:
-            obs_dict[f"observation.images.{name}"] = torch.from_numpy(images[name])
+            obs_dict[f"observation.images.{name}"] = images[name]
 
         return obs_dict, action_dict
 
@@ -482,7 +607,7 @@ class KochRobot:
         """The returned observations do not have a batch dimension."""
         if not self.is_connected:
             raise RobotDeviceNotConnectedError(
-                "KochRobot is not connected. You need to run `robot.connect()`."
+                "ManipulatorRobot is not connected. You need to run `robot.connect()`."
             )
 
         # Read follower position
@@ -490,6 +615,7 @@ class KochRobot:
         for name in self.follower_arms:
             before_fread_t = time.perf_counter()
             follower_pos[name] = self.follower_arms[name].read("Present_Position")
+            follower_pos[name] = torch.from_numpy(follower_pos[name])
             self.logs[f"read_follower_{name}_pos_dt_s"] = time.perf_counter() - before_fread_t
 
         # Create state by concatenating follower current position
@@ -497,80 +623,72 @@ class KochRobot:
         for name in self.follower_arms:
             if name in follower_pos:
                 state.append(follower_pos[name])
-        state = np.concatenate(state)
+        state = torch.cat(state)
 
         # Capture images from cameras
         images = {}
         for name in self.cameras:
             before_camread_t = time.perf_counter()
             images[name] = self.cameras[name].async_read()
+            images[name] = torch.from_numpy(images[name])
             self.logs[f"read_camera_{name}_dt_s"] = self.cameras[name].logs["delta_timestamp_s"]
             self.logs[f"async_read_camera_{name}_dt_s"] = time.perf_counter() - before_camread_t
 
         # Populate output dictionnaries and format to pytorch
         obs_dict = {}
-        obs_dict["observation.state"] = torch.from_numpy(state)
+        obs_dict["observation.state"] = state
         for name in self.cameras:
-            obs_dict[f"observation.images.{name}"] = torch.from_numpy(images[name])
+            obs_dict[f"observation.images.{name}"] = images[name]
         return obs_dict
 
-    def send_action(self, action: torch.Tensor, follower_names: list[str] | None = None):
+    def send_action(self, action: torch.Tensor) -> torch.Tensor:
         """Command the follower arms to move to a target joint configuration.
 
         The relative action magnitude may be clipped depending on the configuration parameter
-        `max_relative_target`.
+        `max_relative_target`. In this case, the action sent differs from original action.
+        Thus, this function always returns the action actually sent.
 
         Args:
-            action: tensor containing the concatenated joint positions for the follower arms.
-            follower_names: Pass follower arm names to only control a subset of all the follower arms.
+            action: tensor containing the concatenated goal positions for the follower arms.
         """
         if not self.is_connected:
             raise RobotDeviceNotConnectedError(
-                "KochRobot is not connected. You need to run `robot.connect()`."
-            )
-
-        if follower_names is None:
-            follower_names = list(self.follower_arms)
-        elif not set(follower_names).issubset(self.follower_arms):
-            raise ValueError(
-                f"You provided {follower_names=} but only the following arms are registered: "
-                f"{list(self.follower_arms)}"
+                "ManipulatorRobot is not connected. You need to run `robot.connect()`."
             )
 
         from_idx = 0
         to_idx = 0
-        follower_goal_pos = {}
-        for name in follower_names:
+        action_sent = []
+        for name in self.follower_arms:
+            # Get goal position of each follower arm by splitting the action vector
             to_idx += len(self.follower_arms[name].motor_names)
-            this_action = action[from_idx:to_idx]
-
-            if self.config.max_relative_target is not None:
-                if not isinstance(self.config.max_relative_target, list):
-                    max_relative_target = [self.config.max_relative_target for _ in range(from_idx, to_idx)]
-                max_relative_target = torch.tensor(self.config.max_relative_target)
-                # Cap relative action target magnitude for safety.
-                current_pos = torch.tensor(self.follower_arms[name].read("Present_Position"))
-                diff = this_action - current_pos
-                safe_diff = torch.minimum(diff, max_relative_target)
-                safe_diff = torch.maximum(safe_diff, -max_relative_target)
-                safe_action = current_pos + safe_diff
-                if not torch.allclose(safe_action, action):
-                    logging.warning(
-                        "Relative action magnitude had to be clamped to be safe.\n"
-                        f"  requested relative action target: {diff}\n"
-                        f"    clamped relative action target: {safe_diff}"
-                    )
-
-            follower_goal_pos[name] = safe_action.numpy()
+            goal_pos = action[from_idx:to_idx]
             from_idx = to_idx
 
-        for name in self.follower_arms:
-            self.follower_arms[name].write("Goal_Position", follower_goal_pos[name].astype(np.int32))
+            # Cap goal position when too far away from present position.
+            # Slower fps expected due to reading from the follower.
+            if self.config.max_relative_target is not None:
+                present_pos = self.follower_arms[name].read("Present_Position")
+                present_pos = torch.from_numpy(present_pos)
+                goal_pos = ensure_safe_goal_position(goal_pos, present_pos, self.config.max_relative_target)
+
+            # Save tensor to concat and return
+            action_sent.append(goal_pos)
+
+            # Send goal position to each follower
+            goal_pos = goal_pos.numpy().astype(np.int32)
+            self.follower_arms[name].write("Goal_Position", goal_pos)
+
+        return torch.cat(action_sent)
+
+    def print_logs(self):
+        pass
+        # TODO(aliberts): move robot-specific logs logic here
 
     def disconnect(self):
         if not self.is_connected:
             raise RobotDeviceNotConnectedError(
-                "KochRobot is not connected. You need to run `robot.connect()` before disconnecting."
+                "ManipulatorRobot is not connected. You need to run `robot.connect()` before disconnecting."
             )
 
         for name in self.follower_arms:

lerobot/common/robot_devices/robots/stretch.py

@@ -0,0 +1,216 @@
+#!/usr/bin/env python
+
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import time
+from dataclasses import dataclass, field, replace
+
+import torch
+from stretch_body.gamepad_teleop import GamePadTeleop
+from stretch_body.robot import Robot as StretchAPI
+from stretch_body.robot_params import RobotParams
+
+from lerobot.common.robot_devices.cameras.utils import Camera
+
+
+@dataclass
+class StretchRobotConfig:
+    robot_type: str | None = "stretch"
+    cameras: dict[str, Camera] = field(default_factory=lambda: {})
+    # TODO(aliberts): add feature with max_relative target
+    # TODO(aliberts): add comment on max_relative target
+    max_relative_target: list[float] | float | None = None
+
+
+class StretchRobot(StretchAPI):
+    """Wrapper of stretch_body.robot.Robot"""
+
+    def __init__(self, config: StretchRobotConfig | None = None, **kwargs):
+        super().__init__()
+        if config is None:
+            config = StretchRobotConfig()
+        # Overwrite config arguments using kwargs
+        self.config = replace(config, **kwargs)
+
+        self.robot_type = self.config.robot_type
+        self.cameras = self.config.cameras
+        self.is_connected = False
+        self.teleop = None
+        self.logs = {}
+
+        # TODO(aliberts): test this
+        RobotParams.set_logging_level("WARNING")
+        RobotParams.set_logging_formatter("brief_console_formatter")
+
+        self.state_keys = None
+        self.action_keys = None
+
+    def connect(self) -> None:
+        self.is_connected = self.startup()
+        if not self.is_connected:
+            print("Another process is already using Stretch. Try running 'stretch_free_robot_process.py'")
+            raise ConnectionError()
+
+        for name in self.cameras:
+            self.cameras[name].connect()
+            self.is_connected = self.is_connected and self.cameras[name].is_connected
+
+        if not self.is_connected:
+            print("Could not connect to the cameras, check that all cameras are plugged-in.")
+            raise ConnectionError()
+
+        self.run_calibration()
+
+    def run_calibration(self) -> None:
+        if not self.is_homed():
+            self.home()
+
+    def teleop_step(
+        self, record_data=False
+    ) -> None | tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
+        # TODO(aliberts): return ndarrays instead of torch.Tensors
+        if not self.is_connected:
+            raise ConnectionError()
+
+        if self.teleop is None:
+            self.teleop = GamePadTeleop(robot_instance=False)
+            self.teleop.startup(robot=self)
+
+        before_read_t = time.perf_counter()
+        state = self.get_state()
+        action = self.teleop.gamepad_controller.get_state()
+        self.logs["read_pos_dt_s"] = time.perf_counter() - before_read_t
+
+        before_write_t = time.perf_counter()
+        self.teleop.do_motion(robot=self)
+        self.push_command()
+        self.logs["write_pos_dt_s"] = time.perf_counter() - before_write_t
+
+        if self.state_keys is None:
+            self.state_keys = list(state)
+
+        if not record_data:
+            return
+
+        state = torch.as_tensor(list(state.values()))
+        action = torch.as_tensor(list(action.values()))
+
+        # Capture images from cameras
+        images = {}
+        for name in self.cameras:
+            before_camread_t = time.perf_counter()
+            images[name] = self.cameras[name].async_read()
+            images[name] = torch.from_numpy(images[name])
+            self.logs[f"read_camera_{name}_dt_s"] = self.cameras[name].logs["delta_timestamp_s"]
+            self.logs[f"async_read_camera_{name}_dt_s"] = time.perf_counter() - before_camread_t
+
+        # Populate output dictionnaries
+        obs_dict, action_dict = {}, {}
+        obs_dict["observation.state"] = state
+        action_dict["action"] = action
+        for name in self.cameras:
+            obs_dict[f"observation.images.{name}"] = images[name]
+
+        return obs_dict, action_dict
+
+    def get_state(self) -> dict:
+        status = self.get_status()
+        return {
+            "head_pan.pos": status["head"]["head_pan"]["pos"],
+            "head_tilt.pos": status["head"]["head_tilt"]["pos"],
+            "lift.pos": status["lift"]["pos"],
+            "arm.pos": status["arm"]["pos"],
+            "wrist_pitch.pos": status["end_of_arm"]["wrist_pitch"]["pos"],
+            "wrist_roll.pos": status["end_of_arm"]["wrist_roll"]["pos"],
+            "wrist_yaw.pos": status["end_of_arm"]["wrist_yaw"]["pos"],
+            "gripper.pos": status["end_of_arm"]["stretch_gripper"]["pos"],
+            "base_x.vel": status["base"]["x_vel"],
+            "base_y.vel": status["base"]["y_vel"],
+            "base_theta.vel": status["base"]["theta_vel"],
+        }
+
+    def capture_observation(self) -> dict:
+        # TODO(aliberts): return ndarrays instead of torch.Tensors
+        before_read_t = time.perf_counter()
+        state = self.get_state()
+        self.logs["read_pos_dt_s"] = time.perf_counter() - before_read_t
+
+        if self.state_keys is None:
+            self.state_keys = list(state)
+
+        state = torch.as_tensor(list(state.values()))
+
+        # Capture images from cameras
+        images = {}
+        for name in self.cameras:
+            before_camread_t = time.perf_counter()
+            images[name] = self.cameras[name].async_read()
+            images[name] = torch.from_numpy(images[name])
+            self.logs[f"read_camera_{name}_dt_s"] = self.cameras[name].logs["delta_timestamp_s"]
+            self.logs[f"async_read_camera_{name}_dt_s"] = time.perf_counter() - before_camread_t
+
+        # Populate output dictionnaries
+        obs_dict = {}
+        obs_dict["observation.state"] = state
+        for name in self.cameras:
+            obs_dict[f"observation.images.{name}"] = images[name]
+
+        return obs_dict
+
+    def send_action(self, action: torch.Tensor) -> torch.Tensor:
+        # TODO(aliberts): return ndarrays instead of torch.Tensors
+        if not self.is_connected:
+            raise ConnectionError()
+
+        if self.teleop is None:
+            self.teleop = GamePadTeleop(robot_instance=False)
+            self.teleop.startup(robot=self)
+
+        if self.action_keys is None:
+            dummy_action = self.teleop.gamepad_controller.get_state()
+            self.action_keys = list(dummy_action.keys())
+
+        action_dict = dict(zip(self.action_keys, action.tolist(), strict=True))
+
+        before_write_t = time.perf_counter()
+        self.teleop.do_motion(state=action_dict, robot=self)
+        self.push_command()
+        self.logs["write_pos_dt_s"] = time.perf_counter() - before_write_t
+
+        # TODO(aliberts): return action_sent when motion is limited
+        return action
+
+    def print_logs(self) -> None:
+        pass
+        # TODO(aliberts): move robot-specific logs logic here
+
+    def teleop_safety_stop(self) -> None:
+        if self.teleop is not None:
+            self.teleop._safety_stop(robot=self)
+
+    def disconnect(self) -> None:
+        self.stop()
+        if self.teleop is not None:
+            self.teleop.gamepad_controller.stop()
+            self.teleop.stop()
+
+        if len(self.cameras) > 0:
+            for cam in self.cameras.values():
+                cam.disconnect()
+
+        self.is_connected = False
+
+    def __del__(self):
+        self.disconnect()

lerobot/common/robot_devices/robots/utils.py

@@ -1,9 +1,20 @@
 from typing import Protocol
 
 
+def get_arm_id(name, arm_type):
+    """Returns the string identifier of a robot arm. For instance, for a bimanual manipulator
+    like Aloha, it could be left_follower, right_follower, left_leader, or right_leader.
+    """
+    return f"{name}_{arm_type}"
+
+
 class Robot(Protocol):
-    def init_teleop(self): ...
+    # TODO(rcadene, aliberts): Add unit test checking the protocol is implemented in the corresponding classes
+    robot_type: str
+
+    def connect(self): ...
     def run_calibration(self): ...
     def teleop_step(self, record_data=False): ...
     def capture_observation(self): ...
     def send_action(self, action): ...
+    def disconnect(self): ...

lerobot/common/robot_devices/utils.py

@@ -1,3 +1,35 @@
+import platform
+import time
+
+
+def busy_wait(seconds):
+    if platform.system() == "Darwin":
+        # On Mac, `time.sleep` is not accurate and we need to use this while loop trick,
+        # but it consumes CPU cycles.
+        # TODO(rcadene): find an alternative: from python 11, time.sleep is precise
+        end_time = time.perf_counter() + seconds
+        while time.perf_counter() < end_time:
+            pass
+    else:
+        # On Linux time.sleep is accurate
+        if seconds > 0:
+            time.sleep(seconds)
+
+
+def safe_disconnect(func):
+    # TODO(aliberts): Allow to pass custom exceptions
+    # (e.g. ThreadServiceExit, KeyboardInterrupt, SystemExit, UnpluggedError, DynamixelCommError)
+    def wrapper(robot, *args, **kwargs):
+        try:
+            return func(robot, *args, **kwargs)
+        except Exception as e:
+            if robot.is_connected:
+                robot.disconnect()
+            raise e
+
+    return wrapper
+
+
 class RobotDeviceNotConnectedError(Exception):
     """Exception raised when the robot device is not connected."""
 

lerobot/configs/default.yaml

@@ -120,7 +120,7 @@ eval:
   # `batch_size` specifies the number of environments to use in a gym.vector.VectorEnv.
   batch_size: 1
   # `use_async_envs` specifies whether to use asynchronous environments (multiprocessing).
-  use_async_envs: true
+  use_async_envs: false
 
 wandb:
   enable: false

lerobot/configs/env/aloha.yaml

@@ -2,11 +2,6 @@
 
 fps: 50
 
-eval:
-  # `use_async_envs` specifies whether to use asynchronous environments (multiprocessing).
-  # set it to false to avoid some problems of the aloha env
-  use_async_envs: false
-
 env:
   name: aloha
   task: AlohaInsertion-v0

lerobot/configs/env/aloha_real.yaml

@@ -0,0 +1,10 @@
+# @package _global_
+
+fps: 30
+
+env:
+  name: real_world
+  task: null
+  state_dim: 14
+  action_dim: 14
+  fps: ${fps}

lerobot/configs/env/xarm.yaml

@@ -2,11 +2,6 @@
 
 fps: 15
 
-eval:
-  # `use_async_envs` specifies whether to use asynchronous environments (multiprocessing).
-  # set it to false to avoid some problems of the aloha env
-  use_async_envs: false
-
 env:
   name: xarm
   task: XarmLift-v0

lerobot/configs/policy/act_aloha_real.yaml

@@ -1,16 +1,22 @@
 # @package _global_
 
-# Use `act_real.yaml` to train on real-world Aloha/Aloha2 datasets.
-# Compared to `act.yaml`, it contains 4 cameras (i.e. cam_right_wrist, cam_left_wrist, images,
-# cam_low) instead of 1 camera (i.e. top). Also, `training.eval_freq` is set to -1. This config is used
-# to evaluate checkpoints at a certain frequency of training steps. When it is set to -1, it deactivates evaluation.
-# This is because real-world evaluation is done through [dora-lerobot](https://github.com/dora-rs/dora-lerobot).
-# Look at its README for more information on how to evaluate a checkpoint in the real-world.
+# Use `act_aloha_real.yaml` to train on real-world datasets collected on Aloha or Aloha-2 robots.
+# Compared to `act.yaml`, it contains 4 cameras (i.e. cam_right_wrist, cam_left_wrist, cam_high, cam_low) instead of 1 camera (i.e. top).
+# Also, `training.eval_freq` is set to -1. This config is used to evaluate checkpoints at a certain frequency of training steps.
+# When it is set to -1, it deactivates evaluation. This is because real-world evaluation is done through our `control_robot.py` script.
+# Look at the documentation in header of `control_robot.py` for more information on how to collect data , train and evaluate a policy.
 #
-# Example of usage for training:
+# Example of usage for training and inference with `control_robot.py`:
 # ```bash
 # python lerobot/scripts/train.py \
-#   policy=act_real \
+#   policy=act_aloha_real \
+#   env=aloha_real
+# ```
+#
+# Example of usage for training and inference with [Dora-rs](https://github.com/dora-rs/dora-lerobot):
+# ```bash
+# python lerobot/scripts/train.py \
+#   policy=act_aloha_real \
 #   env=dora_aloha_real
 # ```
 
@@ -36,10 +42,11 @@ override_dataset_stats:
     std: [[[0.229]], [[0.224]], [[0.225]]]  # (c,1,1)
 
 training:
-  offline_steps: 100000
+  offline_steps: 80000
   online_steps: 0
   eval_freq: -1
-  save_freq: 20000
+  save_freq: 10000
+  log_freq: 100
   save_checkpoint: true
 
   batch_size: 8
@@ -62,7 +69,7 @@ policy:
 
   # Input / output structure.
   n_obs_steps: 1
-  chunk_size: 100 # chunk_size
+  chunk_size: 100
   n_action_steps: 100
 
   input_shapes:
@@ -107,7 +114,7 @@ policy:
   n_vae_encoder_layers: 4
 
   # Inference.
-  temporal_ensemble_coeff: null
+  temporal_ensemble_momentum: null
 
   # Training and loss computation.
   dropout: 0.1

lerobot/configs/policy/act_real_no_state.yaml

@@ -1,110 +0,0 @@
-# @package _global_
-
-# Use `act_real_no_state.yaml` to train on real-world Aloha/Aloha2 datasets when cameras are moving (e.g. wrist cameras)
-# Compared to `act_real.yaml`, it is camera only and does not use the state as input which is vector of robot joint positions.
-# We validated experimentaly that not using state reaches better success rate. Our hypothesis is that `act_real.yaml` might
-# overfits to the state, because the images are more complex to learn from since they are moving.
-#
-# Example of usage for training:
-# ```bash
-# python lerobot/scripts/train.py \
-#   policy=act_real_no_state \
-#   env=dora_aloha_real
-# ```
-
-seed: 1000
-dataset_repo_id: lerobot/aloha_static_vinh_cup
-
-override_dataset_stats:
-  observation.images.cam_right_wrist:
-    # stats from imagenet, since we use a pretrained vision model
-    mean: [[[0.485]], [[0.456]], [[0.406]]]  # (c,1,1)
-    std: [[[0.229]], [[0.224]], [[0.225]]]  # (c,1,1)
-  observation.images.cam_left_wrist:
-    # stats from imagenet, since we use a pretrained vision model
-    mean: [[[0.485]], [[0.456]], [[0.406]]]  # (c,1,1)
-    std: [[[0.229]], [[0.224]], [[0.225]]]  # (c,1,1)
-  observation.images.cam_high:
-    # stats from imagenet, since we use a pretrained vision model
-    mean: [[[0.485]], [[0.456]], [[0.406]]]  # (c,1,1)
-    std: [[[0.229]], [[0.224]], [[0.225]]]  # (c,1,1)
-  observation.images.cam_low:
-    # stats from imagenet, since we use a pretrained vision model
-    mean: [[[0.485]], [[0.456]], [[0.406]]]  # (c,1,1)
-    std: [[[0.229]], [[0.224]], [[0.225]]]  # (c,1,1)
-
-training:
-  offline_steps: 100000
-  online_steps: 0
-  eval_freq: -1
-  save_freq: 20000
-  save_checkpoint: true
-
-  batch_size: 8
-  lr: 1e-5
-  lr_backbone: 1e-5
-  weight_decay: 1e-4
-  grad_clip_norm: 10
-  online_steps_between_rollouts: 1
-
-  delta_timestamps:
-    action: "[i / ${fps} for i in range(${policy.chunk_size})]"
-
-eval:
-  n_episodes: 50
-  batch_size: 50
-
-# See `configuration_act.py` for more details.
-policy:
-  name: act
-
-  # Input / output structure.
-  n_obs_steps: 1
-  chunk_size: 100 # chunk_size
-  n_action_steps: 100
-
-  input_shapes:
-    # TODO(rcadene, alexander-soare): add variables for height and width from the dataset/env?
-    observation.images.cam_right_wrist: [3, 480, 640]
-    observation.images.cam_left_wrist: [3, 480, 640]
-    observation.images.cam_high: [3, 480, 640]
-    observation.images.cam_low: [3, 480, 640]
-  output_shapes:
-    action: ["${env.action_dim}"]
-
-  # Normalization / Unnormalization
-  input_normalization_modes:
-    observation.images.cam_right_wrist: mean_std
-    observation.images.cam_left_wrist: mean_std
-    observation.images.cam_high: mean_std
-    observation.images.cam_low: mean_std
-  output_normalization_modes:
-    action: mean_std
-
-  # Architecture.
-  # Vision backbone.
-  vision_backbone: resnet18
-  pretrained_backbone_weights: ResNet18_Weights.IMAGENET1K_V1
-  replace_final_stride_with_dilation: false
-  # Transformer layers.
-  pre_norm: false
-  dim_model: 512
-  n_heads: 8
-  dim_feedforward: 3200
-  feedforward_activation: relu
-  n_encoder_layers: 4
-  # Note: Although the original ACT implementation has 7 for `n_decoder_layers`, there is a bug in the code
-  # that means only the first layer is used. Here we match the original implementation by setting this to 1.
-  # See this issue https://github.com/tonyzhaozh/act/issues/25#issue-2258740521.
-  n_decoder_layers: 1
-  # VAE.
-  use_vae: true
-  latent_dim: 32
-  n_vae_encoder_layers: 4
-
-  # Inference.
-  temporal_ensemble_coeff: null
-
-  # Training and loss computation.
-  dropout: 0.1
-  kl_weight: 10.0

lerobot/configs/policy/tdmpc.yaml

@@ -12,6 +12,7 @@ training:
   grad_clip_norm: 10.0
   lr: 3e-4
 
+  save_freq: 10000
   eval_freq: 5000
   log_freq: 100
 

lerobot/configs/robot/aloha.yaml

@@ -0,0 +1,115 @@
+# Aloha: A Low-Cost Hardware for Bimanual Teleoperation
+# https://aloha-2.github.io
+# https://www.trossenrobotics.com/aloha-stationary
+
+# Requires installing extras packages
+# With pip: `pip install -e ".[dynamixel intelrealsense]"`
+# With poetry: `poetry install --sync --extras "dynamixel intelrealsense"`
+
+_target_: lerobot.common.robot_devices.robots.manipulator.ManipulatorRobot
+robot_type: aloha
+# Specific to Aloha, LeRobot comes with default calibration files. Assuming the motors have been
+# properly assembled, no manual calibration step is expected. If you need to run manual calibration,
+# simply update this path to ".cache/calibration/aloha"
+calibration_dir: .cache/calibration/aloha_default
+
+# /!\ FOR SAFETY, READ THIS /!\
+# `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.
+# For Aloha, for every goal position request, motor rotations are capped at 5 degrees by default.
+# When you feel more confident with teleoperation or running the policy, you can extend
+# this safety limit and even removing it by setting it to `null`.
+# Also, everything is expected to work safely out-of-the-box, but we highly advise to
+# first try to teleoperate the grippers only (by commenting out the rest of the motors in this yaml),
+# then to gradually add more motors (by uncommenting), until you can teleoperate both arms fully
+max_relative_target: 5
+
+leader_arms:
+  left:
+    _target_: lerobot.common.robot_devices.motors.dynamixel.DynamixelMotorsBus
+    port: /dev/ttyDXL_leader_left
+    motors:  # window_x
+      # name: (index, model)
+      waist: [1, xm430-w350]
+      shoulder: [2, xm430-w350]
+      shoulder_shadow: [3, xm430-w350]
+      elbow: [4, xm430-w350]
+      elbow_shadow: [5, xm430-w350]
+      forearm_roll: [6, xm430-w350]
+      wrist_angle: [7, xm430-w350]
+      wrist_rotate: [8, xl430-w250]
+      gripper: [9, xc430-w150]
+  right:
+    _target_: lerobot.common.robot_devices.motors.dynamixel.DynamixelMotorsBus
+    port: /dev/ttyDXL_leader_right
+    motors:  # window_x
+      # name: (index, model)
+      waist: [1, xm430-w350]
+      shoulder: [2, xm430-w350]
+      shoulder_shadow: [3, xm430-w350]
+      elbow: [4, xm430-w350]
+      elbow_shadow: [5, xm430-w350]
+      forearm_roll: [6, xm430-w350]
+      wrist_angle: [7, xm430-w350]
+      wrist_rotate: [8, xl430-w250]
+      gripper: [9, xc430-w150]
+
+follower_arms:
+  left:
+    _target_: lerobot.common.robot_devices.motors.dynamixel.DynamixelMotorsBus
+    port: /dev/ttyDXL_follower_left
+    motors:
+      # name: [index, model]
+      waist: [1, xm540-w270]
+      shoulder: [2, xm540-w270]
+      shoulder_shadow: [3, xm540-w270]
+      elbow: [4, xm540-w270]
+      elbow_shadow: [5, xm540-w270]
+      forearm_roll: [6, xm540-w270]
+      wrist_angle: [7, xm540-w270]
+      wrist_rotate: [8, xm430-w350]
+      gripper: [9, xm430-w350]
+  right:
+    _target_: lerobot.common.robot_devices.motors.dynamixel.DynamixelMotorsBus
+    port: /dev/ttyDXL_follower_right
+    motors:
+      # name: [index, model]
+      waist: [1, xm540-w270]
+      shoulder: [2, xm540-w270]
+      shoulder_shadow: [3, xm540-w270]
+      elbow: [4, xm540-w270]
+      elbow_shadow: [5, xm540-w270]
+      forearm_roll: [6, xm540-w270]
+      wrist_angle: [7, xm540-w270]
+      wrist_rotate: [8, xm430-w350]
+      gripper: [9, xm430-w350]
+
+# Troubleshooting: If one of your IntelRealSense cameras freeze during
+# data recording due to bandwidth limit, you might need to plug the camera
+# on another USB hub or PCIe card.
+cameras:
+  cam_high:
+    _target_: lerobot.common.robot_devices.cameras.intelrealsense.IntelRealSenseCamera
+    serial_number: 128422271347
+    fps: 30
+    width: 640
+    height: 480
+  cam_low:
+    _target_: lerobot.common.robot_devices.cameras.intelrealsense.IntelRealSenseCamera
+    serial_number: 130322270656
+    fps: 30
+    width: 640
+    height: 480
+  cam_left_wrist:
+    _target_: lerobot.common.robot_devices.cameras.intelrealsense.IntelRealSenseCamera
+    serial_number: 218622272670
+    fps: 30
+    width: 640
+    height: 480
+  cam_right_wrist:
+    _target_: lerobot.common.robot_devices.cameras.intelrealsense.IntelRealSenseCamera
+    serial_number: 130322272300
+    fps: 30
+    width: 640
+    height: 480

lerobot/configs/robot/koch.yaml

@@ -1,5 +1,12 @@
-_target_: lerobot.common.robot_devices.robots.koch.KochRobot
-calibration_path: .cache/calibration/koch.pkl
+_target_: lerobot.common.robot_devices.robots.manipulator.ManipulatorRobot
+robot_type: koch
+calibration_dir: .cache/calibration/koch
+
+# `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: null
+
 leader_arms:
   main:
     _target_: lerobot.common.robot_devices.motors.dynamixel.DynamixelMotorsBus
@@ -12,6 +19,7 @@ leader_arms:
       wrist_flex: [4, "xl330-m077"]
       wrist_roll: [5, "xl330-m077"]
       gripper: [6, "xl330-m077"]
+
 follower_arms:
   main:
     _target_: lerobot.common.robot_devices.motors.dynamixel.DynamixelMotorsBus
@@ -24,6 +32,7 @@ follower_arms:
       wrist_flex: [4, "xl330-m288"]
       wrist_roll: [5, "xl330-m288"]
       gripper: [6, "xl330-m288"]
+
 cameras:
   laptop:
     _target_: lerobot.common.robot_devices.cameras.opencv.OpenCVCamera
@@ -37,10 +46,8 @@ cameras:
     fps: 30
     width: 640
     height: 480
-# `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: null
+
+# ~ Koch specific settings ~
 # Sets the leader arm in torque mode with the gripper motor set to this angle. This makes it possible
 # to squeeze the gripper and have it spring back to an open position on its own.
 gripper_open_degree: 35.156

lerobot/configs/robot/koch_bimanual.yaml

@@ -0,0 +1,75 @@
+_target_: lerobot.common.robot_devices.robots.manipulator.ManipulatorRobot
+robot_type: koch
+calibration_dir: .cache/calibration/koch_bimanual
+
+# `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: null
+
+leader_arms:
+  left:
+    _target_: lerobot.common.robot_devices.motors.dynamixel.DynamixelMotorsBus
+    port: /dev/tty.usbmodem585A0085511
+    motors:
+      # name: (index, model)
+      shoulder_pan: [1, "xl330-m077"]
+      shoulder_lift: [2, "xl330-m077"]
+      elbow_flex: [3, "xl330-m077"]
+      wrist_flex: [4, "xl330-m077"]
+      wrist_roll: [5, "xl330-m077"]
+      gripper: [6, "xl330-m077"]
+  right:
+    _target_: lerobot.common.robot_devices.motors.dynamixel.DynamixelMotorsBus
+    port: /dev/tty.usbmodem575E0031751
+    motors:
+      # name: (index, model)
+      shoulder_pan: [1, "xl330-m077"]
+      shoulder_lift: [2, "xl330-m077"]
+      elbow_flex: [3, "xl330-m077"]
+      wrist_flex: [4, "xl330-m077"]
+      wrist_roll: [5, "xl330-m077"]
+      gripper: [6, "xl330-m077"]
+
+follower_arms:
+  left:
+    _target_: lerobot.common.robot_devices.motors.dynamixel.DynamixelMotorsBus
+    port: /dev/tty.usbmodem585A0076891
+    motors:
+      # name: (index, model)
+      shoulder_pan: [1, "xl430-w250"]
+      shoulder_lift: [2, "xl430-w250"]
+      elbow_flex: [3, "xl330-m288"]
+      wrist_flex: [4, "xl330-m288"]
+      wrist_roll: [5, "xl330-m288"]
+      gripper: [6, "xl330-m288"]
+  right:
+    _target_: lerobot.common.robot_devices.motors.dynamixel.DynamixelMotorsBus
+    port: /dev/tty.usbmodem575E0032081
+    motors:
+      # name: (index, model)
+      shoulder_pan: [1, "xl430-w250"]
+      shoulder_lift: [2, "xl430-w250"]
+      elbow_flex: [3, "xl330-m288"]
+      wrist_flex: [4, "xl330-m288"]
+      wrist_roll: [5, "xl330-m288"]
+      gripper: [6, "xl330-m288"]
+
+cameras:
+  laptop:
+    _target_: lerobot.common.robot_devices.cameras.opencv.OpenCVCamera
+    camera_index: 0
+    fps: 30
+    width: 640
+    height: 480
+  phone:
+    _target_: lerobot.common.robot_devices.cameras.opencv.OpenCVCamera
+    camera_index: 1
+    fps: 30
+    width: 640
+    height: 480
+
+# ~ Koch specific settings ~
+# Sets the leader arm in torque mode with the gripper motor set to this angle. This makes it possible
+# to squeeze the gripper and have it spring back to an open position on its own.
+gripper_open_degree: 35.156

lerobot/configs/robot/stretch.yaml

@@ -0,0 +1,24 @@
+_target_: lerobot.common.robot_devices.robots.stretch.StretchRobot
+robot_type: stretch3
+
+cameras:
+  navigation:
+    _target_: lerobot.common.robot_devices.cameras.opencv.OpenCVCamera
+    camera_index: /dev/hello-nav-head-camera
+    fps: 10
+    width: 1280
+    height: 720
+    rotation: -90
+  head:
+    _target_: lerobot.common.robot_devices.cameras.intelrealsense.IntelRealSenseCamera.init_from_name
+    name: Intel RealSense D435I
+    fps: 30
+    width: 640
+    height: 480
+    rotation: 90
+  wrist:
+    _target_: lerobot.common.robot_devices.cameras.intelrealsense.IntelRealSenseCamera.init_from_name
+    name: Intel RealSense D405
+    fps: 30
+    width: 640
+    height: 480

lerobot/scripts/control_robot.py

@@ -102,6 +102,7 @@ import argparse
 import concurrent.futures
 import json
 import logging
+import multiprocessing
 import os
 import platform
 import shutil
@@ -127,7 +128,8 @@ from lerobot.common.datasets.utils import calculate_episode_data_index, create_b
 from lerobot.common.datasets.video_utils import encode_video_frames
 from lerobot.common.policies.factory import make_policy
 from lerobot.common.robot_devices.robots.factory import make_robot
-from lerobot.common.robot_devices.robots.utils import Robot
+from lerobot.common.robot_devices.robots.utils import Robot, get_arm_id
+from lerobot.common.robot_devices.utils import busy_wait, safe_disconnect
 from lerobot.common.utils.utils import get_safe_torch_device, init_hydra_config, init_logging, set_global_seed
 from lerobot.scripts.eval import get_pretrained_policy_path
 from lerobot.scripts.push_dataset_to_hub import (
@@ -169,59 +171,53 @@ def save_image(img_tensor, key, frame_index, episode_index, videos_dir):
     img.save(str(path), quality=100)
 
 
-def busy_wait(seconds):
-    # Significantly more accurate than `time.sleep`, and mendatory for our use case,
-    # but it consumes CPU cycles.
-    # TODO(rcadene): find an alternative: from python 11, time.sleep is precise
-    end_time = time.perf_counter() + seconds
-    while time.perf_counter() < end_time:
-        pass
-
-
 def none_or_int(value):
     if value == "None":
         return None
     return int(value)
 
 
-def log_control_info(robot, dt_s, episode_index=None, frame_index=None, fps=None):
+def log_control_info(robot: Robot, dt_s, episode_index=None, frame_index=None, fps=None):
     log_items = []
     if episode_index is not None:
-        log_items += [f"ep:{episode_index}"]
+        log_items.append(f"ep:{episode_index}")
     if frame_index is not None:
-        log_items += [f"frame:{frame_index}"]
+        log_items.append(f"frame:{frame_index}")
 
     def log_dt(shortname, dt_val_s):
-        nonlocal log_items
-        log_items += [f"{shortname}:{dt_val_s * 1000:5.2f} ({1/ dt_val_s:3.1f}hz)"]
+        nonlocal log_items, fps
+        info_str = f"{shortname}:{dt_val_s * 1000:5.2f} ({1/ dt_val_s:3.1f}hz)"
+        if fps is not None:
+            actual_fps = 1 / dt_val_s
+            if actual_fps < fps - 1:
+                info_str = colored(info_str, "yellow")
+        log_items.append(info_str)
 
     # total step time displayed in milliseconds and its frequency
     log_dt("dt", dt_s)
 
-    for name in robot.leader_arms:
-        key = f"read_leader_{name}_pos_dt_s"
-        if key in robot.logs:
-            log_dt("dtRlead", robot.logs[key])
+    # TODO(aliberts): move robot-specific logs logic in robot.print_logs()
+    if not robot.robot_type.startswith("stretch"):
+        for name in robot.leader_arms:
+            key = f"read_leader_{name}_pos_dt_s"
+            if key in robot.logs:
+                log_dt("dtRlead", robot.logs[key])
 
-    for name in robot.follower_arms:
-        key = f"write_follower_{name}_goal_pos_dt_s"
-        if key in robot.logs:
-            log_dt("dtWfoll", robot.logs[key])
+        for name in robot.follower_arms:
+            key = f"write_follower_{name}_goal_pos_dt_s"
+            if key in robot.logs:
+                log_dt("dtWfoll", robot.logs[key])
 
-        key = f"read_follower_{name}_pos_dt_s"
-        if key in robot.logs:
-            log_dt("dtRfoll", robot.logs[key])
+            key = f"read_follower_{name}_pos_dt_s"
+            if key in robot.logs:
+                log_dt("dtRfoll", robot.logs[key])
 
-    for name in robot.cameras:
-        key = f"read_camera_{name}_dt_s"
-        if key in robot.logs:
-            log_dt(f"dtR{name}", robot.logs[key])
+        for name in robot.cameras:
+            key = f"read_camera_{name}_dt_s"
+            if key in robot.logs:
+                log_dt(f"dtR{name}", robot.logs[key])
 
     info_str = " ".join(log_items)
-    if fps is not None:
-        actual_fps = 1 / dt_s
-        if actual_fps < fps - 1:
-            info_str = colored(info_str, "yellow")
     logging.info(info_str)
 
 
@@ -244,15 +240,102 @@ def is_headless():
         return True
 
 
+def loop_to_save_frame_in_threads(frame_queue, num_image_writers):
+    with concurrent.futures.ThreadPoolExecutor(max_workers=num_image_writers) as executor:
+        futures = []
+        while True:
+            # Blocks until a frame is available
+            frame_data = frame_queue.get()
+
+            # Exit if we send None to stop the worker
+            if frame_data is None:
+                # Wait for all submitted futures to complete before exiting
+                for _ in tqdm.tqdm(
+                    concurrent.futures.as_completed(futures), total=len(futures), desc="Writting images"
+                ):
+                    pass
+                break
+
+            frame, key, frame_index, episode_index, videos_dir = frame_data
+            futures.append(executor.submit(save_image, frame, key, frame_index, episode_index, videos_dir))
+
+
+def start_frame_workers(frame_queue, num_image_writers, num_workers=1):
+    workers = []
+    for _ in range(num_workers):
+        worker = multiprocessing.Process(
+            target=loop_to_save_frame_in_threads,
+            args=(frame_queue, num_image_writers),
+        )
+        worker.start()
+        workers.append(worker)
+    return workers
+
+
+def stop_workers(workers, frame_queue):
+    # Send None to each process to signal it to stop
+    for _ in workers:
+        frame_queue.put(None)
+
+    # Wait for all processes to terminate
+    for process in workers:
+        process.join()
+
+
+def has_method(_object: object, method_name: str):
+    return hasattr(_object, method_name) and callable(getattr(_object, method_name))
+
+
+def get_available_arms(robot):
+    # TODO(rcadene): moves this function in manipulator class?
+    available_arms = []
+    for name in robot.follower_arms:
+        arm_id = get_arm_id(name, "follower")
+        available_arms.append(arm_id)
+    for name in robot.leader_arms:
+        arm_id = get_arm_id(name, "leader")
+        available_arms.append(arm_id)
+    return available_arms
+
+
 ########################################################################################
 # Control modes
 ########################################################################################
 
 
-def calibrate(robot: Robot):
-    if robot.calibration_path.exists():
-        print(f"Removing '{robot.calibration_path}'")
-        robot.calibration_path.unlink()
+@safe_disconnect
+def calibrate(robot: Robot, arms: list[str] | None):
+    # TODO(aliberts): move this code in robots' classes
+    if robot.robot_type.startswith("stretch"):
+        if not robot.is_connected:
+            robot.connect()
+        if not robot.is_homed():
+            robot.home()
+        return
+
+    available_arms = get_available_arms(robot)
+    unknown_arms = [arm_id for arm_id in arms if arm_id not in available_arms]
+    available_arms_str = " ".join(available_arms)
+    unknown_arms_str = " ".join(unknown_arms)
+
+    if arms is None or len(arms) == 0:
+        raise ValueError(
+            "No arm provided. Use `--arms` as argument with one or more available arms.\n"
+            f"For instance, to recalibrate all arms add: `--arms {available_arms_str}`"
+        )
+
+    if len(unknown_arms) > 0:
+        raise ValueError(
+            f"Unknown arms provided ('{unknown_arms_str}'). Available arms are `{available_arms_str}`."
+        )
+
+    for arm_id in arms:
+        arm_calib_path = robot.calibration_dir / f"{arm_id}.json"
+        if arm_calib_path.exists():
+            print(f"Removing '{arm_calib_path}'")
+            arm_calib_path.unlink()
+        else:
+            print(f"Calibration file not found '{arm_calib_path}'")
 
     if robot.is_connected:
         robot.disconnect()
@@ -260,8 +343,11 @@ def calibrate(robot: Robot):
     # Calling `connect` automatically runs calibration
     # when the calibration file is missing
     robot.connect()
+    robot.disconnect()
+    print("Calibration is done! You can now teleoperate and record datasets!")
 
 
+@safe_disconnect
 def teleoperate(robot: Robot, fps: int | None = None, teleop_time_s: float | None = None):
     # TODO(rcadene): Add option to record logs
     if not robot.is_connected:
@@ -283,6 +369,7 @@ def teleoperate(robot: Robot, fps: int | None = None, teleop_time_s: float | Non
             break
 
 
+@safe_disconnect
 def record(
     robot: Robot,
     policy: torch.nn.Module | None = None,
@@ -298,8 +385,9 @@ def record(
     run_compute_stats=True,
     push_to_hub=True,
     tags=None,
-    num_image_writers=8,
+    num_image_writers_per_camera=4,
     force_override=False,
+    display_cameras=True,
 ):
     # TODO(rcadene): Add option to record logs
     # TODO(rcadene): Clean this function via decomposition in higher level functions
@@ -310,9 +398,6 @@ def record(
             f"Your dataset name begins by 'eval_' ({dataset_name}) but no policy is provided ({policy})."
         )
 
-    if not video:
-        raise NotImplementedError()
-
     if not robot.is_connected:
         robot.connect()
 
@@ -336,7 +421,7 @@ def record(
         episode_index = 0
 
     if is_headless():
-        logging.info(
+        logging.warning(
             "Headless environment detected. On-screen cameras display and keyboard inputs will not be available."
         )
 
@@ -404,7 +489,7 @@ def record(
         else:
             observation = robot.capture_observation()
 
-        if not is_headless():
+        if display_cameras and not is_headless():
             image_keys = [key for key in observation if "image" in key]
             for key in image_keys:
                 cv2.imshow(key, cv2.cvtColor(observation[key].numpy(), cv2.COLOR_RGB2BGR))
@@ -418,11 +503,18 @@ def record(
 
         timestamp = time.perf_counter() - start_warmup_t
 
+    if has_method(robot, "teleop_safety_stop"):
+        robot.teleop_safety_stop()
+
     # Save images using threads to reach high fps (30 and more)
     # Using `with` to exist smoothly if an execption is raised.
-    # Using only 4 worker threads to avoid blocking the main thread.
-    futures = []
-    with concurrent.futures.ThreadPoolExecutor(max_workers=num_image_writers) as executor:
+    num_image_writers = num_image_writers_per_camera * len(robot.cameras)
+    num_image_writers = max(num_image_writers, 1)
+    frame_queue = multiprocessing.Queue()
+    frame_workers = start_frame_workers(frame_queue, num_image_writers)
+
+    # Using `try` to exist smoothly if an exception is raised
+    try:
         # Start recording all episodes
         while episode_index < num_episodes:
             logging.info(f"Recording episode {episode_index}")
@@ -443,13 +535,9 @@ def record(
                 not_image_keys = [key for key in observation if "image" not in key]
 
                 for key in image_keys:
-                    futures += [
-                        executor.submit(
-                            save_image, observation[key], key, frame_index, episode_index, videos_dir
-                        )
-                    ]
+                    frame_queue.put((observation[key], key, frame_index, episode_index, videos_dir))
 
-                if not is_headless():
+                if display_cameras and not is_headless():
                     image_keys = [key for key in observation if "image" in key]
                     for key in image_keys:
                         cv2.imshow(key, cv2.cvtColor(observation[key].numpy(), cv2.COLOR_RGB2BGR))
@@ -486,8 +574,11 @@ def record(
                         action = action.to("cpu")
 
                     # Order the robot to move
-                    robot.send_action(action)
-                    action = {"action": action}
+                    action_sent = robot.send_action(action)
+
+                    # Action can eventually be clipped using `max_relative_target`,
+                    # so action actually sent is saved in the dataset.
+                    action = {"action": action_sent}
 
                 for key in action:
                     if key not in ep_dict:
@@ -507,6 +598,10 @@ def record(
                     exit_early = False
                     break
 
+            # TODO(alibets): allow for teleop during reset
+            if has_method(robot, "teleop_safety_stop"):
+                robot.teleop_safety_stop()
+
             if not stop_recording:
                 # Start resetting env while the executor are finishing
                 logging.info("Reset the environment")
@@ -519,15 +614,23 @@ def record(
             num_frames = frame_index
 
             for key in image_keys:
-                tmp_imgs_dir = videos_dir / f"{key}_episode_{episode_index:06d}"
-                fname = f"{key}_episode_{episode_index:06d}.mp4"
-                video_path = local_dir / "videos" / fname
-                if video_path.exists():
-                    video_path.unlink()
-                # Store the reference to the video frame, even tho the videos are not yet encoded
-                ep_dict[key] = []
-                for i in range(num_frames):
-                    ep_dict[key].append({"path": f"videos/{fname}", "timestamp": i / fps})
+                if video:
+                    tmp_imgs_dir = videos_dir / f"{key}_episode_{episode_index:06d}"
+                    fname = f"{key}_episode_{episode_index:06d}.mp4"
+                    video_path = local_dir / "videos" / fname
+                    if video_path.exists():
+                        video_path.unlink()
+                    # Store the reference to the video frame, even tho the videos are not yet encoded
+                    ep_dict[key] = []
+                    for i in range(num_frames):
+                        ep_dict[key].append({"path": f"videos/{fname}", "timestamp": i / fps})
+
+                else:
+                    imgs_dir = videos_dir / f"{key}_episode_{episode_index:06d}"
+                    ep_dict[key] = []
+                    for i in range(num_frames):
+                        img_path = imgs_dir / f"frame_{i:06d}.png"
+                        ep_dict[key].append({"path": str(img_path)})
 
             for key in not_image_keys:
                 ep_dict[key] = torch.stack(ep_dict[key])
@@ -579,33 +682,34 @@ def record(
                     listener.stop()
 
                 logging.info("Waiting for threads writing the images on disk to terminate...")
-                for _ in tqdm.tqdm(
-                    concurrent.futures.as_completed(futures), total=len(futures), desc="Writting images"
-                ):
-                    pass
-                break
+                stop_workers(frame_workers, frame_queue)
+
+    except Exception:
+        traceback.print_exc()
+        stop_workers(frame_workers, frame_queue)
 
     robot.disconnect()
-    if not is_headless():
+    if display_cameras and not is_headless():
         cv2.destroyAllWindows()
 
     num_episodes = episode_index
 
-    logging.info("Encoding videos")
-    say("Encoding videos")
-    # Use ffmpeg to convert frames stored as png into mp4 videos
-    for episode_index in tqdm.tqdm(range(num_episodes)):
-        for key in image_keys:
-            tmp_imgs_dir = videos_dir / f"{key}_episode_{episode_index:06d}"
-            fname = f"{key}_episode_{episode_index:06d}.mp4"
-            video_path = local_dir / "videos" / fname
-            if video_path.exists():
-                # Skip if video is already encoded. Could be the case when resuming data recording.
-                continue
-            # note: `encode_video_frames` is a blocking call. Making it asynchronous shouldn't speedup encoding,
-            # since video encoding with ffmpeg is already using multithreading.
-            encode_video_frames(tmp_imgs_dir, video_path, fps, overwrite=True)
-            shutil.rmtree(tmp_imgs_dir)
+    if video:
+        logging.info("Encoding videos")
+        say("Encoding videos")
+        # Use ffmpeg to convert frames stored as png into mp4 videos
+        for episode_index in tqdm.tqdm(range(num_episodes)):
+            for key in image_keys:
+                tmp_imgs_dir = videos_dir / f"{key}_episode_{episode_index:06d}"
+                fname = f"{key}_episode_{episode_index:06d}.mp4"
+                video_path = local_dir / "videos" / fname
+                if video_path.exists():
+                    # Skip if video is already encoded. Could be the case when resuming data recording.
+                    continue
+                # note: `encode_video_frames` is a blocking call. Making it asynchronous shouldn't speedup encoding,
+                # since video encoding with ffmpeg is already using multithreading.
+                encode_video_frames(tmp_imgs_dir, video_path, fps, overwrite=True)
+                shutil.rmtree(tmp_imgs_dir)
 
     logging.info("Concatenating episodes")
     ep_dicts = []
@@ -712,6 +816,12 @@ if __name__ == "__main__":
     )
 
     parser_calib = subparsers.add_parser("calibrate", parents=[base_parser])
+    parser_calib.add_argument(
+        "--arms",
+        type=str,
+        nargs="*",
+        help="List of arms to calibrate (e.g. `--arms left_follower right_follower left_leader`)",
+    )
 
     parser_teleop = subparsers.add_parser("teleoperate", parents=[base_parser])
     parser_teleop.add_argument(
@@ -772,10 +882,14 @@ if __name__ == "__main__":
         help="Add tags to your dataset on the hub.",
     )
     parser_record.add_argument(
-        "--num-image-writers",
+        "--num-image-writers-per-camera",
         type=int,
-        default=8,
-        help="Number of threads writing the frames as png images on disk. Don't set too much as you might get unstable fps due to main thread being blocked.",
+        default=4,
+        help=(
+            "Number of threads writing the frames as png images on disk, per camera. "
+            "Too much threads might cause unstable teleoperation fps due to main thread being blocked. "
+            "Not enough threads might cause low camera fps."
+        ),
     )
     parser_record.add_argument(
         "--force-override",

lerobot/scripts/eval.py

@@ -57,7 +57,7 @@ import gymnasium as gym
 import numpy as np
 import torch
 from huggingface_hub import snapshot_download
-from huggingface_hub.utils._errors import RepositoryNotFoundError
+from huggingface_hub.errors import RepositoryNotFoundError
 from huggingface_hub.utils._validators import HFValidationError
 from torch import Tensor, nn
 from tqdm import trange

lerobot/scripts/train.py

@@ -93,16 +93,6 @@ def make_optimizer_and_scheduler(cfg, policy):
     elif policy.name == "tdmpc":
         optimizer = torch.optim.Adam(policy.parameters(), cfg.training.lr)
         lr_scheduler = None
-    elif policy.name == "tdmpc2":
-        params_group = [
-            {"params": policy.model._encoder.parameters(), "lr": cfg.training.lr * cfg.training.enc_lr_scale},
-            {"params": policy.model._dynamics.parameters()},
-            {"params": policy.model._reward.parameters()},
-            {"params": policy.model._Qs.parameters()},
-            {"params": policy.model._pi.parameters(), "eps": 1e-5},
-        ]
-        optimizer = torch.optim.Adam(params_group, lr=cfg.training.lr)
-        lr_scheduler = None
     elif cfg.policy.name == "vqbet":
         from lerobot.common.policies.vqbet.modeling_vqbet import VQBeTOptimizer, VQBeTScheduler
 

lerobot/scripts/visualize_dataset_html.py

@@ -57,7 +57,6 @@ import logging
 import shutil
 from pathlib import Path
 
-import torch
 import tqdm
 from flask import Flask, redirect, render_template, url_for
 
@@ -65,19 +64,6 @@ from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.common.utils.utils import init_logging
 
 
-class EpisodeSampler(torch.utils.data.Sampler):
-    def __init__(self, dataset, episode_index):
-        from_idx = dataset.episode_data_index["from"][episode_index].item()
-        to_idx = dataset.episode_data_index["to"][episode_index].item()
-        self.frame_ids = range(from_idx, to_idx)
-
-    def __iter__(self):
-        return iter(self.frame_ids)
-
-    def __len__(self):
-        return len(self.frame_ids)
-
-
 def run_server(
     dataset: LeRobotDataset,
     episodes: list[int],
@@ -112,10 +98,14 @@ def run_server(
             "fps": dataset.fps,
         }
         video_paths = get_episode_video_paths(dataset, episode_id)
+        language_instruction = get_episode_language_instruction(dataset, episode_id)
         videos_info = [
             {"url": url_for("static", filename=video_path), "filename": Path(video_path).name}
             for video_path in video_paths
         ]
+        if language_instruction:
+            videos_info[0]["language_instruction"] = language_instruction
+
         ep_csv_url = url_for("static", filename=get_ep_csv_fname(episode_id))
         return render_template(
             "visualize_dataset_template.html",
@@ -186,6 +176,20 @@ def get_episode_video_paths(dataset: LeRobotDataset, ep_index: int) -> list[str]
     ]
 
 
+def get_episode_language_instruction(dataset: LeRobotDataset, ep_index: int) -> list[str]:
+    # check if the dataset has language instructions
+    if "language_instruction" not in dataset.hf_dataset.features:
+        return None
+
+    # get first frame index
+    first_frame_idx = dataset.episode_data_index["from"][ep_index].item()
+
+    language_instruction = dataset.hf_dataset[first_frame_idx]["language_instruction"]
+    # TODO (michel-aractingi) hack to get the sentence, some strings in openx are badly stored
+    # with the tf.tensor appearing in the string
+    return language_instruction.removeprefix("tf.Tensor(b'").removesuffix("', shape=(), dtype=string)")
+
+
 def visualize_dataset_html(
     repo_id: str,
     root: Path | None = None,

lerobot/templates/visualize_dataset_template.html

@@ -14,7 +14,7 @@
 <!-- Use [Alpin.js](https://alpinejs.dev), a lightweight and easy to learn JS framework -->
 <!-- Use [tailwindcss](https://tailwindcss.com/), CSS classes for styling html -->
 <!-- Use [dygraphs](https://dygraphs.com/), a lightweight JS charting library -->
-<body class="flex h-screen max-h-screen bg-slate-950 text-gray-200" x-data="createAlpineData()" @keydown.window="(e) => {
+<body class="flex flex-col md:flex-row h-screen max-h-screen bg-slate-950 text-gray-200" x-data="createAlpineData()" @keydown.window="(e) => {
     // Use the space bar to play and pause, instead of default action (e.g. scrolling)
     const { keyCode, key } = e;
     if (keyCode === 32 || key === ' ') {
@@ -30,7 +30,7 @@
     }
 }">
     <!-- Sidebar -->
-    <div x-ref="sidebar" class="w-60 bg-slate-900 p-5 break-words max-h-screen overflow-y-auto">
+    <div x-ref="sidebar" class="bg-slate-900 p-5 break-words overflow-y-auto shrink-0 md:shrink md:w-60 md:max-h-screen">
         <h1 class="mb-4 text-xl font-semibold">{{ dataset_info.repo_id }}</h1>
 
         <ul>
@@ -46,7 +46,8 @@
         </ul>
 
         <p>Episodes:</p>
-        <ul class="ml-2">
+        <!-- episodes menu for medium & large screens -->
+        <ul class="ml-2 hidden md:block">
             {% for episode in episodes %}
             <li class="font-mono text-sm mt-0.5">
                 <a href="episode_{{ episode }}" class="underline {% if episode_id == episode %}font-bold -ml-1{% endif %}">
@@ -56,26 +57,47 @@
             {% endfor %}
         </ul>
 
+        <!-- episodes menu for small screens -->
+        <div class="flex overflow-x-auto md:hidden">
+            {% for episode in episodes %}
+            <p class="font-mono text-sm mt-0.5 border-r last:border-r-0 px-2 {% if episode_id == episode %}font-bold{% endif %}">
+                <a href="episode_{{ episode }}" class="">
+                    {{ episode }}
+                </a>
+            </p>
+            {% endfor %}
+        </div>
+
     </div>
 
     <!-- Toggle sidebar button -->
-    <button class="flex items-center opacity-50 hover:opacity-100 mx-1"
+    <button class="flex items-center opacity-50 hover:opacity-100 mx-1 hidden md:block"
         @click="() => ($refs.sidebar.classList.toggle('hidden'))" title="Toggle sidebar">
         <div class="bg-slate-500 w-2 h-10 rounded-full"></div>
     </button>
 
     <!-- Content -->
-    <div class="flex-1 max-h-screen flex flex-col gap-4 overflow-y-auto">
+    <div class="max-h-screen flex flex-col gap-4 overflow-y-auto md:flex-1">
         <h1 class="text-xl font-bold mt-4 font-mono">
             Episode {{ episode_id }}
         </h1>
 
+        <!-- Error message -->
+        <div class="font-medium text-orange-700 hidden" :class="{ 'hidden': !videoCodecError }">
+            <p>Videos could NOT play because <a href="https://en.wikipedia.org/wiki/AV1" target="_blank" class="underline">AV1</a> decoding is not available on your browser.</p>
+            <ul class="list-decimal list-inside">
+                <li>If iPhone: <span class="italic">It is supported with A17 chip or higher.</span></li>
+                <li>If Mac with Safari: <span class="italic">It is supported on most browsers except Safari with M1 chip or higher and on Safari with M3 chip or higher.</span></li>
+                <li>Other: <span class="italic">Contact the maintainers on LeRobot discord channel:</span> <a href="https://discord.com/invite/s3KuuzsPFb" target="_blank" class="underline">https://discord.com/invite/s3KuuzsPFb</a></li>
+            </ul>
+        </div>
+
         <!-- Videos -->
         <div class="flex flex-wrap gap-1">
             {% for video_info in videos_info %}
-            <div class="max-w-96">
+            <div x-show="!videoCodecError" class="max-w-96">
                 <p class="text-sm text-gray-300 bg-gray-800 px-2 rounded-t-xl truncate">{{ video_info.filename }}</p>
-                <video muted loop type="video/mp4" class="min-w-64" @canplay="videoCanPlay" @timeupdate="() => {
+                <video muted loop type="video/mp4" class="object-contain w-full h-full" @canplaythrough="videoCanPlay" @timeupdate="() => {
                     if (video.duration) {
                       const time = video.currentTime;
                       const pc = (100 / video.duration) * time;
@@ -100,6 +122,13 @@
             {% endfor %}
         </div>
 
+        <!-- Language instruction -->
+        {% if videos_info[0].language_instruction %}
+        <p class="font-medium mt-2">
+            Language Instruction: <span class="italic">{{ videos_info[0].language_instruction }}</span>
+        </p>
+        {% endif %}
+
         <!-- Shortcuts info -->
         <div class="text-sm hidden md:block">
             Hotkeys: <span class="font-mono">Space</span> to pause/unpause, <span class="font-mono">Arrow Down</span> to go to next episode, <span class="font-mono">Arrow Up</span> to go to previous episode.
@@ -176,9 +205,9 @@
                             </td>
                             <template x-for="(cell, colIndex) in row">
                                 <td x-show="cell" class="border border-slate-700">
-                                    <div class="flex gap-x-2 w-24 justify-between px-2">
+                                    <div class="flex gap-x-2 w-24 justify-between px-2" :class="{ 'hidden': cell.isNull }">
                                         <input type="checkbox" x-model="cell.checked" @change="updateTableValues()">
-                                        <span x-text="`${cell.value.toFixed(2)}`"
+                                        <span x-text="`${!cell.isNull ? cell.value.toFixed(2) : null}`"
                                             :style="`color: ${cell.color}`"></span>
                                     </div>
                                 </td>
@@ -200,7 +229,9 @@
                 dygraph: null,
                 currentFrameData: null,
                 columnNames: ["state", "action", "pred action"],
-                nColumns: {% if has_policy %}3{% else %}2{% endif %},
+                nColumns: 2,
+                nStates: 0,
+                nActions: 0,
                 checked: [],
                 dygraphTime: 0.0,
                 dygraphIndex: 0,
@@ -209,9 +240,18 @@
                 colors: null,
                 nVideos: {{ videos_info | length }},
                 nVideoReadyToPlay: 0,
+                videoCodecError: false,
 
                 // alpine initialization
                 init() {
+                    // check if videos can play
+                    const dummyVideo = document.createElement('video');
+                    const canPlayVideos = dummyVideo.canPlayType('video/mp4; codecs="av01.0.05M.08"'); // codec source: https://huggingface.co/blog/video-encoding#results
+                    if(!canPlayVideos){
+                        this.videoCodecError = true;
+                    }
+                    
+                    // process CSV data
                     this.videos = document.querySelectorAll('video');
                     this.video = this.videos[0];
                     this.dygraph = new Dygraph(document.getElementById("graph"), '{{ ep_csv_url }}', {
@@ -236,17 +276,19 @@
                                 this.checked = Array(this.colors.length).fill(true);
 
                                 const seriesNames = this.dygraph.getLabels().slice(1);
+                                this.nStates = seriesNames.findIndex(item => item.startsWith('action_'));
+                                this.nActions = seriesNames.length - this.nStates;
                                 const colors = [];
                                 const LIGHTNESS = [30, 65, 85]; // state_lightness, action_lightness, pred_action_lightness
-                                let lightnessIdx = 0;
-                                const chunkSize = Math.ceil(seriesNames.length / this.nColumns);
-                                for (let i = 0; i < seriesNames.length; i += chunkSize) {
-                                    const lightness = LIGHTNESS[lightnessIdx];
-                                    for (let hue = 0; hue < 360; hue += parseInt(360/chunkSize)) {
-                                        const color = `hsl(${hue}, 100%, ${lightness}%)`;
-                                        colors.push(color);
-                                    }
-                                    lightnessIdx += 1;
+                                // colors for "state" lines
+                                for (let hue = 0; hue < 360; hue += parseInt(360/this.nStates)) {
+                                    const color = `hsl(${hue}, 100%, ${LIGHTNESS[0]}%)`;
+                                    colors.push(color);
+                                }
+                                // colors for "action" lines
+                                for (let hue = 0; hue < 360; hue += parseInt(360/this.nActions)) {
+                                    const color = `hsl(${hue}, 100%, ${LIGHTNESS[1]}%)`;
+                                    colors.push(color);
                                 }
                                 this.dygraph.updateOptions({ colors });
                                 this.colors = colors;
@@ -273,37 +315,40 @@
                     if (!this.currentFrameData) {
                         return [];
                     }
-                    const columnSize = Math.ceil(this.currentFrameData.length / this.nColumns);
-                    return Array.from({
-                        length: columnSize
-                    }, (_, rowIndex) => {
-                        const row = [
-                            this.currentFrameData[rowIndex] || null,
-                            this.currentFrameData[rowIndex + columnSize] || null,
-                        ];
-                        if (this.nColumns === 3) {
-                            row.push(this.currentFrameData[rowIndex + 2 * columnSize] || null)
-                        }
-                        return row;
-                    });
+                    const rows = [];
+                    const nRows = Math.max(this.nStates, this.nActions);
+                    let rowIndex = 0;
+                    while(rowIndex < nRows){
+                        const row = [];
+                        // number of states may NOT match number of actions. In this case, we null-pad the 2D array to make a fully rectangular 2d array
+                        const nullCell = { isNull: true };
+                        const stateValueIdx = rowIndex;
+                        const actionValueIdx = stateValueIdx + this.nStates; // because this.currentFrameData = [state0, state1, ..., stateN, action0, action1, ..., actionN]
+                        // row consists of [state value, action value]
+                        row.push(rowIndex < this.nStates ? this.currentFrameData[stateValueIdx] : nullCell); // push "state value" to row
+                        row.push(rowIndex < this.nActions ? this.currentFrameData[actionValueIdx] : nullCell); // push "action value" to row
+                        rowIndex += 1;
+                        rows.push(row);
+                    }
+                    return rows;
                 },
                 isRowChecked(rowIndex) {
-                    return this.rows[rowIndex].every(cell => cell && cell.checked);
+                    return this.rows[rowIndex].every(cell => cell && (cell.isNull || cell.checked));
                 },
                 isColumnChecked(colIndex) {
-                    return this.rows.every(row => row[colIndex] && row[colIndex].checked);
+                    return this.rows.every(row => row[colIndex] && (row[colIndex].isNull || row[colIndex].checked));
                 },
                 toggleRow(rowIndex) {
                     const newState = !this.isRowChecked(rowIndex);
                     this.rows[rowIndex].forEach(cell => {
-                        if (cell) cell.checked = newState;
+                        if (cell && !cell.isNull) cell.checked = newState;
                     });
                     this.updateTableValues();
                 },
                 toggleColumn(colIndex) {
                     const newState = !this.isColumnChecked(colIndex);
                     this.rows.forEach(row => {
-                        if (row[colIndex]) row[colIndex].checked = newState;
+                        if (row[colIndex] && !row[colIndex].isNull) row[colIndex].checked = newState;
                     });
                     this.updateTableValues();
                 },

pyproject.toml

@@ -43,7 +43,7 @@ opencv-python = ">=4.9.0"
 diffusers = ">=0.27.2"
 torchvision = ">=0.17.1"
 h5py = ">=3.10.0"
-huggingface-hub = {extras = ["hf-transfer", "cli"], version = ">=0.23.0"}
+huggingface-hub = {extras = ["hf-transfer", "cli"], version = ">=0.25.0"}
 gymnasium = ">=0.29.1"
 cmake = ">=3.29.0.1"
 gym-dora = { git = "https://github.com/dora-rs/dora-lerobot.git", subdirectory = "gym_dora", optional = true }
@@ -64,9 +64,11 @@ pandas = {version = ">=2.2.2", optional = true}
 scikit-image = {version = ">=0.23.2", optional = true}
 dynamixel-sdk = {version = ">=3.7.31", optional = true}
 pynput = {version = ">=1.7.7", optional = true}
-# TODO(rcadene, salibert): 71.0.1 has a bug
-setuptools = {version = "!=71.0.1", optional = true}
-
+setuptools = {version = "!=71.0.1", optional = true}  # TODO(rcadene, aliberts): 71.0.1 has a bug
+pyrealsense2 = {version = ">=2.55.1.6486", markers = "sys_platform != 'darwin'", optional = true}  # TODO(rcadene, aliberts): Fix on Mac
+pyrender = {git = "https://github.com/mmatl/pyrender.git", markers = "sys_platform == 'linux'", optional = true}
+hello-robot-stretch-body = {version = ">=0.7.27", markers = "sys_platform == 'linux'", optional = true}
+pyserial = {version = ">=3.5", optional = true}
 
 
 [tool.poetry.extras]
@@ -75,10 +77,12 @@ pusht = ["gym-pusht"]
 xarm = ["gym-xarm"]
 aloha = ["gym-aloha"]
 dev = ["pre-commit", "debugpy"]
-test = ["pytest", "pytest-cov"]
+test = ["pytest", "pytest-cov", "pyserial"]
 umi = ["imagecodecs"]
 video_benchmark = ["scikit-image", "pandas"]
-koch = ["dynamixel-sdk", "pynput"]
+dynamixel = ["dynamixel-sdk", "pynput"]
+intelrealsense = ["pyrealsense2"]
+stretch = ["hello-robot-stretch-body", "pyrender", "pyrealsense2", "pynput"]
 
 [tool.ruff]
 line-length = 110

tests/conftest.py

@@ -13,28 +13,104 @@
 # 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 traceback
+
 import pytest
+from serial import SerialException
 
+from lerobot import available_cameras, available_motors, available_robots
 from lerobot.common.utils.utils import init_hydra_config
-
-from .utils import DEVICE, KOCH_ROBOT_CONFIG_PATH
+from tests.utils import DEVICE, ROBOT_CONFIG_PATH_TEMPLATE, make_camera, make_motors_bus
 
 
 def pytest_collection_finish():
     print(f"\nTesting with {DEVICE=}")
 
 
-@pytest.fixture(scope="session")
-def is_koch_available():
+@pytest.fixture
+def is_robot_available(robot_type):
+    if robot_type not in available_robots:
+        raise ValueError(
+            f"The robot type '{robot_type}' is not valid. Expected one of these '{available_robots}"
+        )
+
     try:
         from lerobot.common.robot_devices.robots.factory import make_robot
 
-        robot_cfg = init_hydra_config(KOCH_ROBOT_CONFIG_PATH)
+        config_path = ROBOT_CONFIG_PATH_TEMPLATE.format(robot=robot_type)
+        robot_cfg = init_hydra_config(config_path)
         robot = make_robot(robot_cfg)
         robot.connect()
         del robot
         return True
+
     except Exception as e:
-        print("A koch robot is not available.")
-        print(e)
+        print(f"\nA {robot_type} robot is not available.")
+
+        if isinstance(e, ModuleNotFoundError):
+            print(f"\nInstall module '{e.name}'")
+        elif isinstance(e, SerialException):
+            print("\nNo physical motors bus detected.")
+
+        traceback.print_exc()
         return False
+
+
+@pytest.fixture
+def is_camera_available(camera_type):
+    if camera_type not in available_cameras:
+        raise ValueError(
+            f"The camera type '{camera_type}' is not valid. Expected one of these '{available_cameras}"
+        )
+
+    try:
+        camera = make_camera(camera_type)
+        camera.connect()
+        del camera
+        return True
+
+    except Exception as e:
+        print(f"\nA {camera_type} camera is not available.")
+
+        if isinstance(e, ModuleNotFoundError):
+            print(f"\nInstall module '{e.name}'")
+        elif isinstance(e, ValueError) and "camera_index" in e.args[0]:
+            print("\nNo physical camera detected.")
+
+        traceback.print_exc()
+        return False
+
+
+@pytest.fixture
+def is_motor_available(motor_type):
+    if motor_type not in available_motors:
+        raise ValueError(
+            f"The motor type '{motor_type}' is not valid. Expected one of these '{available_motors}"
+        )
+
+    try:
+        motors_bus = make_motors_bus(motor_type)
+        motors_bus.connect()
+        del motors_bus
+        return True
+
+    except Exception as e:
+        print(f"\nA {motor_type} motor is not available.")
+
+        if isinstance(e, ModuleNotFoundError):
+            print(f"\nInstall module '{e.name}'")
+        elif isinstance(e, SerialException):
+            print("\nNo physical motors bus detected.")
+
+        traceback.print_exc()
+        return False
+
+
+@pytest.fixture
+def patch_builtins_input(monkeypatch):
+    def print_text(text=None):
+        if text is not None:
+            print(text)
+
+    monkeypatch.setattr("builtins.input", print_text)

tests/data/lerobot/aloha_mobile_cabinet/meta_data/info.json

@@ -1,3 +1,11 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:50e40e4c2bb523fca0b54e9a9635281312e9c6f9d757db03c06a0865c5508f29
-size 188
+{
+    "codebase_version": "v1.6",
+    "fps": 50,
+    "video": true,
+    "encoding": {
+        "vcodec": "libsvtav1",
+        "pix_fmt": "yuv420p",
+        "g": 2,
+        "crf": 30
+    }
+}

tests/data/lerobot/aloha_mobile_cabinet/train/dataset_info.json

@@ -1,3 +1,61 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:8b7fbedfdb3d536847bc6fadf2cbabb9f2b5492edf3e2c274a3e8ffb447105e8
-size 1166
+{
+  "citation": "",
+  "description": "",
+  "features": {
+    "observation.images.cam_high": {
+      "_type": "VideoFrame"
+    },
+    "observation.images.cam_left_wrist": {
+      "_type": "VideoFrame"
+    },
+    "observation.images.cam_right_wrist": {
+      "_type": "VideoFrame"
+    },
+    "observation.state": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 14,
+      "_type": "Sequence"
+    },
+    "observation.effort": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 14,
+      "_type": "Sequence"
+    },
+    "action": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 14,
+      "_type": "Sequence"
+    },
+    "episode_index": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "frame_index": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "timestamp": {
+      "dtype": "float32",
+      "_type": "Value"
+    },
+    "next.done": {
+      "dtype": "bool",
+      "_type": "Value"
+    },
+    "index": {
+      "dtype": "int64",
+      "_type": "Value"
+    }
+  },
+  "homepage": "",
+  "license": ""
+}

tests/data/lerobot/aloha_mobile_cabinet/train/state.json

@@ -1,3 +1,13 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:34ece24fb6b302db0b68987858509f31713fb299faa9a9d34b8fd68f10bc3100
-size 247
+{
+  "_data_files": [
+    {
+      "filename": "data-00000-of-00001.arrow"
+    }
+  ],
+  "_fingerprint": "3e76021c95d21c4d",
+  "_format_columns": null,
+  "_format_kwargs": {},
+  "_format_type": null,
+  "_output_all_columns": false,
+  "_split": null
+}

tests/data/lerobot/aloha_mobile_chair/meta_data/info.json

@@ -1,3 +1,11 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:50e40e4c2bb523fca0b54e9a9635281312e9c6f9d757db03c06a0865c5508f29
-size 188
+{
+    "codebase_version": "v1.6",
+    "fps": 50,
+    "video": true,
+    "encoding": {
+        "vcodec": "libsvtav1",
+        "pix_fmt": "yuv420p",
+        "g": 2,
+        "crf": 30
+    }
+}

tests/data/lerobot/aloha_mobile_chair/train/dataset_info.json

@@ -1,3 +1,61 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:8b7fbedfdb3d536847bc6fadf2cbabb9f2b5492edf3e2c274a3e8ffb447105e8
-size 1166
+{
+  "citation": "",
+  "description": "",
+  "features": {
+    "observation.images.cam_high": {
+      "_type": "VideoFrame"
+    },
+    "observation.images.cam_left_wrist": {
+      "_type": "VideoFrame"
+    },
+    "observation.images.cam_right_wrist": {
+      "_type": "VideoFrame"
+    },
+    "observation.state": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 14,
+      "_type": "Sequence"
+    },
+    "observation.effort": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 14,
+      "_type": "Sequence"
+    },
+    "action": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 14,
+      "_type": "Sequence"
+    },
+    "episode_index": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "frame_index": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "timestamp": {
+      "dtype": "float32",
+      "_type": "Value"
+    },
+    "next.done": {
+      "dtype": "bool",
+      "_type": "Value"
+    },
+    "index": {
+      "dtype": "int64",
+      "_type": "Value"
+    }
+  },
+  "homepage": "",
+  "license": ""
+}

tests/data/lerobot/aloha_mobile_chair/train/state.json

@@ -1,3 +1,13 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:38cf4116a65cb92a5c43f9b9da7a7b81cfa9168b17605c8c456f7d3a3a23b77a
-size 247
+{
+  "_data_files": [
+    {
+      "filename": "data-00000-of-00001.arrow"
+    }
+  ],
+  "_fingerprint": "872117944c4ecdff",
+  "_format_columns": null,
+  "_format_kwargs": {},
+  "_format_type": null,
+  "_output_all_columns": false,
+  "_split": null
+}

tests/data/lerobot/aloha_mobile_elevator/meta_data/info.json

@@ -1,3 +1,11 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:50e40e4c2bb523fca0b54e9a9635281312e9c6f9d757db03c06a0865c5508f29
-size 188
+{
+    "codebase_version": "v1.6",
+    "fps": 50,
+    "video": true,
+    "encoding": {
+        "vcodec": "libsvtav1",
+        "pix_fmt": "yuv420p",
+        "g": 2,
+        "crf": 30
+    }
+}

tests/data/lerobot/aloha_mobile_elevator/train/dataset_info.json

@@ -1,3 +1,61 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:8b7fbedfdb3d536847bc6fadf2cbabb9f2b5492edf3e2c274a3e8ffb447105e8
-size 1166
+{
+  "citation": "",
+  "description": "",
+  "features": {
+    "observation.images.cam_high": {
+      "_type": "VideoFrame"
+    },
+    "observation.images.cam_left_wrist": {
+      "_type": "VideoFrame"
+    },
+    "observation.images.cam_right_wrist": {
+      "_type": "VideoFrame"
+    },
+    "observation.state": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 14,
+      "_type": "Sequence"
+    },
+    "observation.effort": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 14,
+      "_type": "Sequence"
+    },
+    "action": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 14,
+      "_type": "Sequence"
+    },
+    "episode_index": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "frame_index": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "timestamp": {
+      "dtype": "float32",
+      "_type": "Value"
+    },
+    "next.done": {
+      "dtype": "bool",
+      "_type": "Value"
+    },
+    "index": {
+      "dtype": "int64",
+      "_type": "Value"
+    }
+  },
+  "homepage": "",
+  "license": ""
+}

tests/data/lerobot/aloha_mobile_elevator/train/state.json

@@ -1,3 +1,13 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:d394d451929b805f2d94f9fc5b12d15c31cfc494df76d7d642b63378b8ba0131
-size 247
+{
+  "_data_files": [
+    {
+      "filename": "data-00000-of-00001.arrow"
+    }
+  ],
+  "_fingerprint": "ba1d9dc6ea5a9717",
+  "_format_columns": null,
+  "_format_kwargs": {},
+  "_format_type": null,
+  "_output_all_columns": false,
+  "_split": null
+}

tests/data/lerobot/aloha_mobile_shrimp/meta_data/info.json

@@ -1,3 +1,11 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:50e40e4c2bb523fca0b54e9a9635281312e9c6f9d757db03c06a0865c5508f29
-size 188
+{
+    "codebase_version": "v1.6",
+    "fps": 50,
+    "video": true,
+    "encoding": {
+        "vcodec": "libsvtav1",
+        "pix_fmt": "yuv420p",
+        "g": 2,
+        "crf": 30
+    }
+}

tests/data/lerobot/aloha_mobile_shrimp/train/dataset_info.json

@@ -1,3 +1,61 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:8b7fbedfdb3d536847bc6fadf2cbabb9f2b5492edf3e2c274a3e8ffb447105e8
-size 1166
+{
+  "citation": "",
+  "description": "",
+  "features": {
+    "observation.images.cam_high": {
+      "_type": "VideoFrame"
+    },
+    "observation.images.cam_left_wrist": {
+      "_type": "VideoFrame"
+    },
+    "observation.images.cam_right_wrist": {
+      "_type": "VideoFrame"
+    },
+    "observation.state": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 14,
+      "_type": "Sequence"
+    },
+    "observation.effort": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 14,
+      "_type": "Sequence"
+    },
+    "action": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 14,
+      "_type": "Sequence"
+    },
+    "episode_index": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "frame_index": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "timestamp": {
+      "dtype": "float32",
+      "_type": "Value"
+    },
+    "next.done": {
+      "dtype": "bool",
+      "_type": "Value"
+    },
+    "index": {
+      "dtype": "int64",
+      "_type": "Value"
+    }
+  },
+  "homepage": "",
+  "license": ""
+}

tests/data/lerobot/aloha_mobile_shrimp/train/state.json

@@ -1,3 +1,13 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:69435f30146a309c8d7d0eb01216555bf0547095db1fc9c20218d481d6fe62c8
-size 247
+{
+  "_data_files": [
+    {
+      "filename": "data-00000-of-00001.arrow"
+    }
+  ],
+  "_fingerprint": "d95a3a7eae59566d",
+  "_format_columns": null,
+  "_format_kwargs": {},
+  "_format_type": null,
+  "_output_all_columns": false,
+  "_split": null
+}

tests/data/lerobot/aloha_mobile_wash_pan/meta_data/info.json

@@ -1,3 +1,11 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:50e40e4c2bb523fca0b54e9a9635281312e9c6f9d757db03c06a0865c5508f29
-size 188
+{
+    "codebase_version": "v1.6",
+    "fps": 50,
+    "video": true,
+    "encoding": {
+        "vcodec": "libsvtav1",
+        "pix_fmt": "yuv420p",
+        "g": 2,
+        "crf": 30
+    }
+}

tests/data/lerobot/aloha_mobile_wash_pan/train/dataset_info.json

@@ -1,3 +1,61 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:8b7fbedfdb3d536847bc6fadf2cbabb9f2b5492edf3e2c274a3e8ffb447105e8
-size 1166
+{
+  "citation": "",
+  "description": "",
+  "features": {
+    "observation.images.cam_high": {
+      "_type": "VideoFrame"
+    },
+    "observation.images.cam_left_wrist": {
+      "_type": "VideoFrame"
+    },
+    "observation.images.cam_right_wrist": {
+      "_type": "VideoFrame"
+    },
+    "observation.state": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 14,
+      "_type": "Sequence"
+    },
+    "observation.effort": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 14,
+      "_type": "Sequence"
+    },
+    "action": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 14,
+      "_type": "Sequence"
+    },
+    "episode_index": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "frame_index": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "timestamp": {
+      "dtype": "float32",
+      "_type": "Value"
+    },
+    "next.done": {
+      "dtype": "bool",
+      "_type": "Value"
+    },
+    "index": {
+      "dtype": "int64",
+      "_type": "Value"
+    }
+  },
+  "homepage": "",
+  "license": ""
+}

tests/data/lerobot/aloha_mobile_wash_pan/train/state.json

@@ -1,3 +1,13 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:921505133c62906bd53034a613a827996994875d84c8b26d69d188df9a7ffeba
-size 247
+{
+  "_data_files": [
+    {
+      "filename": "data-00000-of-00001.arrow"
+    }
+  ],
+  "_fingerprint": "09ba9b66b7f468bc",
+  "_format_columns": null,
+  "_format_kwargs": {},
+  "_format_type": null,
+  "_output_all_columns": false,
+  "_split": null
+}

tests/data/lerobot/aloha_mobile_wipe_wine/meta_data/info.json

@@ -1,3 +1,11 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:50e40e4c2bb523fca0b54e9a9635281312e9c6f9d757db03c06a0865c5508f29
-size 188
+{
+    "codebase_version": "v1.6",
+    "fps": 50,
+    "video": true,
+    "encoding": {
+        "vcodec": "libsvtav1",
+        "pix_fmt": "yuv420p",
+        "g": 2,
+        "crf": 30
+    }
+}

tests/data/lerobot/aloha_mobile_wipe_wine/train/dataset_info.json

@@ -1,3 +1,61 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:8b7fbedfdb3d536847bc6fadf2cbabb9f2b5492edf3e2c274a3e8ffb447105e8
-size 1166
+{
+  "citation": "",
+  "description": "",
+  "features": {
+    "observation.images.cam_high": {
+      "_type": "VideoFrame"
+    },
+    "observation.images.cam_left_wrist": {
+      "_type": "VideoFrame"
+    },
+    "observation.images.cam_right_wrist": {
+      "_type": "VideoFrame"
+    },
+    "observation.state": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 14,
+      "_type": "Sequence"
+    },
+    "observation.effort": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 14,
+      "_type": "Sequence"
+    },
+    "action": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 14,
+      "_type": "Sequence"
+    },
+    "episode_index": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "frame_index": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "timestamp": {
+      "dtype": "float32",
+      "_type": "Value"
+    },
+    "next.done": {
+      "dtype": "bool",
+      "_type": "Value"
+    },
+    "index": {
+      "dtype": "int64",
+      "_type": "Value"
+    }
+  },
+  "homepage": "",
+  "license": ""
+}

tests/data/lerobot/aloha_mobile_wipe_wine/train/state.json

@@ -1,3 +1,13 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:505a42c408d56c8a7d3e2367280b41e27667b58334f32e84c937c44c38217bd6
-size 247
+{
+  "_data_files": [
+    {
+      "filename": "data-00000-of-00001.arrow"
+    }
+  ],
+  "_fingerprint": "5a33376149b4e966",
+  "_format_columns": null,
+  "_format_kwargs": {},
+  "_format_type": null,
+  "_output_all_columns": false,
+  "_split": null
+}

tests/data/lerobot/aloha_sim_insertion_human/meta_data/info.json

@@ -1,3 +1,11 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:50e40e4c2bb523fca0b54e9a9635281312e9c6f9d757db03c06a0865c5508f29
-size 188
+{
+    "codebase_version": "v1.6",
+    "fps": 50,
+    "video": true,
+    "encoding": {
+        "vcodec": "libsvtav1",
+        "pix_fmt": "yuv420p",
+        "g": 2,
+        "crf": 30
+    }
+}

tests/data/lerobot/aloha_sim_insertion_human/train/dataset_info.json

@@ -1,3 +1,47 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:3f44d13de5d5a417263bbd4984942ed42ed3fa0633405aa14d9a969a45274944
-size 842
+{
+  "citation": "",
+  "description": "",
+  "features": {
+    "observation.images.top": {
+      "_type": "VideoFrame"
+    },
+    "observation.state": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 14,
+      "_type": "Sequence"
+    },
+    "action": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 14,
+      "_type": "Sequence"
+    },
+    "episode_index": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "frame_index": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "timestamp": {
+      "dtype": "float32",
+      "_type": "Value"
+    },
+    "next.done": {
+      "dtype": "bool",
+      "_type": "Value"
+    },
+    "index": {
+      "dtype": "int64",
+      "_type": "Value"
+    }
+  },
+  "homepage": "",
+  "license": ""
+}

tests/data/lerobot/aloha_sim_insertion_human/train/state.json

@@ -1,3 +1,13 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:c0b18566cbf59e399ea40f1630df12ffbbb9f73bbc733d1d4eba62d675b1fda5
-size 247
+{
+  "_data_files": [
+    {
+      "filename": "data-00000-of-00001.arrow"
+    }
+  ],
+  "_fingerprint": "3aa08798f073758b",
+  "_format_columns": null,
+  "_format_kwargs": {},
+  "_format_type": null,
+  "_output_all_columns": false,
+  "_split": null
+}

tests/data/lerobot/aloha_sim_insertion_human_image/meta_data/info.json

@@ -1,3 +1,5 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:e7ab5c2bd7d176d4d7902a600240318c2828b7d75f4a888d0887327e4eff089d
-size 65
+{
+    "codebase_version": "v1.6",
+    "fps": 50,
+    "video": 0
+}

tests/data/lerobot/aloha_sim_insertion_human_image/train/dataset_info.json

@@ -1,3 +1,47 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:4c9545525dc1f4d550591bd5efb63b55c15b983ae0510fefda5a16d77c78b6ef
-size 837
+{
+  "citation": "",
+  "description": "",
+  "features": {
+    "observation.images.top": {
+      "_type": "Image"
+    },
+    "observation.state": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 14,
+      "_type": "Sequence"
+    },
+    "action": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 14,
+      "_type": "Sequence"
+    },
+    "episode_index": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "frame_index": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "timestamp": {
+      "dtype": "float32",
+      "_type": "Value"
+    },
+    "next.done": {
+      "dtype": "bool",
+      "_type": "Value"
+    },
+    "index": {
+      "dtype": "int64",
+      "_type": "Value"
+    }
+  },
+  "homepage": "",
+  "license": ""
+}

tests/data/lerobot/aloha_sim_insertion_human_image/train/state.json

@@ -1,3 +1,13 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:d7aa033603dc90582516dbcdf3e71e4d3113b70ad49098535def0b282135b5f3
-size 247
+{
+  "_data_files": [
+    {
+      "filename": "data-00000-of-00001.arrow"
+    }
+  ],
+  "_fingerprint": "9fe3a4bf575a8a67",
+  "_format_columns": null,
+  "_format_kwargs": {},
+  "_format_type": null,
+  "_output_all_columns": false,
+  "_split": null
+}

tests/data/lerobot/aloha_sim_insertion_scripted/meta_data/info.json

@@ -1,3 +1,11 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:50e40e4c2bb523fca0b54e9a9635281312e9c6f9d757db03c06a0865c5508f29
-size 188
+{
+    "codebase_version": "v1.6",
+    "fps": 50,
+    "video": true,
+    "encoding": {
+        "vcodec": "libsvtav1",
+        "pix_fmt": "yuv420p",
+        "g": 2,
+        "crf": 30
+    }
+}

tests/data/lerobot/aloha_sim_insertion_scripted/train/dataset_info.json

@@ -1,3 +1,47 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:3f44d13de5d5a417263bbd4984942ed42ed3fa0633405aa14d9a969a45274944
-size 842
+{
+  "citation": "",
+  "description": "",
+  "features": {
+    "observation.images.top": {
+      "_type": "VideoFrame"
+    },
+    "observation.state": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 14,
+      "_type": "Sequence"
+    },
+    "action": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 14,
+      "_type": "Sequence"
+    },
+    "episode_index": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "frame_index": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "timestamp": {
+      "dtype": "float32",
+      "_type": "Value"
+    },
+    "next.done": {
+      "dtype": "bool",
+      "_type": "Value"
+    },
+    "index": {
+      "dtype": "int64",
+      "_type": "Value"
+    }
+  },
+  "homepage": "",
+  "license": ""
+}

tests/data/lerobot/aloha_sim_insertion_scripted/train/state.json

@@ -1,3 +1,13 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:5fd5fe80657788d044cdc8a1baf1456c7695cc951049347a469165002a83c6c7
-size 247
+{
+  "_data_files": [
+    {
+      "filename": "data-00000-of-00001.arrow"
+    }
+  ],
+  "_fingerprint": "e7db591f2ec3eeec",
+  "_format_columns": null,
+  "_format_kwargs": {},
+  "_format_type": null,
+  "_output_all_columns": false,
+  "_split": null
+}

tests/data/lerobot/aloha_sim_insertion_scripted_image/meta_data/info.json

@@ -1,3 +1,5 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:e7ab5c2bd7d176d4d7902a600240318c2828b7d75f4a888d0887327e4eff089d
-size 65
+{
+    "codebase_version": "v1.6",
+    "fps": 50,
+    "video": 0
+}

tests/data/lerobot/aloha_sim_insertion_scripted_image/train/dataset_info.json

@@ -1,3 +1,47 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:4c9545525dc1f4d550591bd5efb63b55c15b983ae0510fefda5a16d77c78b6ef
-size 837
+{
+  "citation": "",
+  "description": "",
+  "features": {
+    "observation.images.top": {
+      "_type": "Image"
+    },
+    "observation.state": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 14,
+      "_type": "Sequence"
+    },
+    "action": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 14,
+      "_type": "Sequence"
+    },
+    "episode_index": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "frame_index": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "timestamp": {
+      "dtype": "float32",
+      "_type": "Value"
+    },
+    "next.done": {
+      "dtype": "bool",
+      "_type": "Value"
+    },
+    "index": {
+      "dtype": "int64",
+      "_type": "Value"
+    }
+  },
+  "homepage": "",
+  "license": ""
+}

tests/data/lerobot/aloha_sim_insertion_scripted_image/train/state.json

@@ -1,3 +1,13 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:752660d8fd884b33b7302a4a42ec7c680de2a3e5022d7d007586f4c6337ce08a
-size 247
+{
+  "_data_files": [
+    {
+      "filename": "data-00000-of-00001.arrow"
+    }
+  ],
+  "_fingerprint": "9be74ce70441e055",
+  "_format_columns": null,
+  "_format_kwargs": {},
+  "_format_type": null,
+  "_output_all_columns": false,
+  "_split": null
+}

tests/data/lerobot/aloha_sim_transfer_cube_human/meta_data/info.json

@@ -1,3 +1,11 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:50e40e4c2bb523fca0b54e9a9635281312e9c6f9d757db03c06a0865c5508f29
-size 188
+{
+    "codebase_version": "v1.6",
+    "fps": 50,
+    "video": true,
+    "encoding": {
+        "vcodec": "libsvtav1",
+        "pix_fmt": "yuv420p",
+        "g": 2,
+        "crf": 30
+    }
+}

tests/data/lerobot/aloha_sim_transfer_cube_human/train/dataset_info.json

@@ -1,3 +1,47 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:3f44d13de5d5a417263bbd4984942ed42ed3fa0633405aa14d9a969a45274944
-size 842
+{
+  "citation": "",
+  "description": "",
+  "features": {
+    "observation.images.top": {
+      "_type": "VideoFrame"
+    },
+    "observation.state": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 14,
+      "_type": "Sequence"
+    },
+    "action": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 14,
+      "_type": "Sequence"
+    },
+    "episode_index": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "frame_index": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "timestamp": {
+      "dtype": "float32",
+      "_type": "Value"
+    },
+    "next.done": {
+      "dtype": "bool",
+      "_type": "Value"
+    },
+    "index": {
+      "dtype": "int64",
+      "_type": "Value"
+    }
+  },
+  "homepage": "",
+  "license": ""
+}

tests/data/lerobot/aloha_sim_transfer_cube_human/train/state.json

@@ -1,3 +1,13 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:bf10e41f2df8c5dc1c19ba8e2d2256ec9d81bdbe28b95079d28039fe2a28504c
-size 247
+{
+  "_data_files": [
+    {
+      "filename": "data-00000-of-00001.arrow"
+    }
+  ],
+  "_fingerprint": "fe4ace534342bfad",
+  "_format_columns": null,
+  "_format_kwargs": {},
+  "_format_type": null,
+  "_output_all_columns": false,
+  "_split": null
+}

tests/data/lerobot/aloha_sim_transfer_cube_human_image/meta_data/info.json

@@ -1,3 +1,5 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:e7ab5c2bd7d176d4d7902a600240318c2828b7d75f4a888d0887327e4eff089d
-size 65
+{
+    "codebase_version": "v1.6",
+    "fps": 50,
+    "video": 0
+}

tests/data/lerobot/aloha_sim_transfer_cube_human_image/train/dataset_info.json

@@ -1,3 +1,47 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:4c9545525dc1f4d550591bd5efb63b55c15b983ae0510fefda5a16d77c78b6ef
-size 837
+{
+  "citation": "",
+  "description": "",
+  "features": {
+    "observation.images.top": {
+      "_type": "Image"
+    },
+    "observation.state": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 14,
+      "_type": "Sequence"
+    },
+    "action": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 14,
+      "_type": "Sequence"
+    },
+    "episode_index": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "frame_index": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "timestamp": {
+      "dtype": "float32",
+      "_type": "Value"
+    },
+    "next.done": {
+      "dtype": "bool",
+      "_type": "Value"
+    },
+    "index": {
+      "dtype": "int64",
+      "_type": "Value"
+    }
+  },
+  "homepage": "",
+  "license": ""
+}

tests/data/lerobot/aloha_sim_transfer_cube_human_image/train/state.json

@@ -1,3 +1,13 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:ec153065a4e52f7d55a7f026d804c57a3ce05dc1faa255a1947369f83c70f1e7
-size 247
+{
+  "_data_files": [
+    {
+      "filename": "data-00000-of-00001.arrow"
+    }
+  ],
+  "_fingerprint": "d7e89e41d66a244c",
+  "_format_columns": null,
+  "_format_kwargs": {},
+  "_format_type": null,
+  "_output_all_columns": false,
+  "_split": null
+}

tests/data/lerobot/aloha_sim_transfer_cube_scripted/meta_data/info.json

@@ -1,3 +1,11 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:50e40e4c2bb523fca0b54e9a9635281312e9c6f9d757db03c06a0865c5508f29
-size 188
+{
+    "codebase_version": "v1.6",
+    "fps": 50,
+    "video": true,
+    "encoding": {
+        "vcodec": "libsvtav1",
+        "pix_fmt": "yuv420p",
+        "g": 2,
+        "crf": 30
+    }
+}

tests/data/lerobot/aloha_sim_transfer_cube_scripted/train/dataset_info.json

@@ -1,3 +1,47 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:3f44d13de5d5a417263bbd4984942ed42ed3fa0633405aa14d9a969a45274944
-size 842
+{
+  "citation": "",
+  "description": "",
+  "features": {
+    "observation.images.top": {
+      "_type": "VideoFrame"
+    },
+    "observation.state": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 14,
+      "_type": "Sequence"
+    },
+    "action": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 14,
+      "_type": "Sequence"
+    },
+    "episode_index": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "frame_index": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "timestamp": {
+      "dtype": "float32",
+      "_type": "Value"
+    },
+    "next.done": {
+      "dtype": "bool",
+      "_type": "Value"
+    },
+    "index": {
+      "dtype": "int64",
+      "_type": "Value"
+    }
+  },
+  "homepage": "",
+  "license": ""
+}

tests/data/lerobot/aloha_sim_transfer_cube_scripted/train/state.json

@@ -1,3 +1,13 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:8f3dc6a1e423b8cd007396a0261664f364d330d350e569a6a0958b834ee1619b
-size 247
+{
+  "_data_files": [
+    {
+      "filename": "data-00000-of-00001.arrow"
+    }
+  ],
+  "_fingerprint": "88443018205d3133",
+  "_format_columns": null,
+  "_format_kwargs": {},
+  "_format_type": null,
+  "_output_all_columns": false,
+  "_split": null
+}

tests/data/lerobot/aloha_sim_transfer_cube_scripted_image/meta_data/info.json

@@ -1,3 +1,5 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:e7ab5c2bd7d176d4d7902a600240318c2828b7d75f4a888d0887327e4eff089d
-size 65
+{
+    "codebase_version": "v1.6",
+    "fps": 50,
+    "video": 0
+}