Control aloha robot natively (#316)

Co-authored-by: Simon Alibert <[email protected]>

.cache/calibration/aloha_default/left_follower.json

@@ -0,0 +1 @@
+{"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 @@
+{"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 @@
+{"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 @@
+{"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

@@ -1,3 +1,4 @@
+.cache/calibration/aloha_default/*.json -filter -diff -merge text
 *.memmap filter=lfs diff=lfs merge=lfs -text
 *.stl filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text

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/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,16 @@ 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"
 ```
 
 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 +147,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 +175,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 +301,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.
+
+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`.
+
+**Calibrate and Connect the ManipulatorRobot**
 
-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.
+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 [`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.
+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 +362,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 +421,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 +614,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 +932,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

lerobot/__init__.py

@@ -27,6 +27,7 @@
         print(lerobot.available_real_world_datasets)
         print(lerobot.available_policies)
         print(lerobot.available_policies_per_env)
+        print(lerobot.available_robots)
     ```
 
 When implementing a new dataset loadable with LeRobotDataset follow these steps:
@@ -182,14 +183,21 @@
     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",
     "tdmpc",
     "vqbet",
 ]
 
+# lists all available robots from `lerobot/common/robot_devices/robots`
+available_robots = [
+    "koch",
+    "koch_bimanual",
+    "aloha",
+]
+
 # keys and values refer to yaml files
 available_policies_per_env = {
     "aloha": ["act"],

lerobot/common/robot_devices/cameras/opencv.py

@@ -17,9 +17,12 @@
 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.