fix

examples/7_get_started_with_real_robot.md

@@ -0,0 +1,394 @@
+# Robots in the real-world
+
+This tutorial explains how to get started with real robots and train a neural network to control them autonomously.
+
+It covers how to:
+1. order and assemble your robot,
+2. connect your robot, configure it and calibrate it,
+3. record your dataset and visualize it,
+4. train a policy on your data and make sure it's ready for evaluation,
+5. evaluate your policy and visualize the result afterwards.
+
+Following these steps, you should be able to reproduce behaviors like picking a lego block and placing it in a bin with a relatively high success rate.
+
+While this tutorial is general and easily extendable to any type of robots by changing a configuration, it is based on the [Koch v1.1](https://github.com/jess-moss/koch-v1-1) affordable robot. Koch v1.1 is composed of a leader arm and a follower arm with 6 motors each. In addition, various cameras can record the scene and serve as visual sensors for the robot.
+
+During data collection, you will control the follower arm by moving the leader arm. This is called "teleoperation". More specifically, the present position of the motors of the leader arm is read at high frequency and sent as a goal position for the motors of the follower arm, which effectively "follow" the movements of the leader arm. While you teleoperate the robot, a few modalities are recorded:
+- the present position of the follower arm, called the "state",
+- the goal position sent to the follower arm, called the "action",
+- the video stream from the cameras.
+
+Finally, you will train a neural network to predict the future actions given the state and camera frames as input ; and deploy it to autonomously control the robot via the high frequency communication of goal positions to the follower arm.
+
+
+## 1. Order and Assemble your Koch v1.1
+
+Follow the bill of materials on the [Koch v1.1 github page](https://github.com/jess-moss/koch-v1-1) to order a leader and a follower arm. Some parts and prices are a bit different with respect to the geo location.
+
+Once the parts are received, follow this video to guide you through the assembly:
+
+## 2. Connect, Configure, and Calibrate your Koch v1.1
+
+Connect the leader arm (the smaller one) with the 5V alimentation and the follower arm with the 12V alimentation. Then connect both arms to your computer with USB.
+
+### Control your motors with DynamixelMotorsBus
+
+[`DynamixelMotorsBus`](lerobot/common/robot_devices/motors/dynamixel.py) allows to efficiently read from and write to the motors connected as a chain to the corresponding usb bus. Underneath, it relies on the python [dynamixel sdk](https://emanual.robotis.com/docs/en/software/dynamixel/dynamixel_sdk/sample_code/python_read_write_protocol_2_0/#python-read-write-protocol-20).
+
+**Instantiate**
+
+Each `DynamixelMotorsBus` requires its corresponding usb port (e.g. `DynamixelMotorsBus(port="/dev/tty.usbmodem575E0031751"`). Run our utility script for each arm to find their ports. Here is an example of what it looks like:
+```bash
+python lerobot/common/robot_devices/motors/dynamixel.py
+>>> Finding all available ports for the DynamixelMotorsBus.
+>>> ['/dev/tty.usbmodem575E0032081', '/dev/tty.usbmodem575E0031751']
+>>> Remove the usb cable from your DynamixelMotorsBus and press Enter when done.
+... **Disconnect leader arm and press Enter**
+>>> The port of this DynamixelMotorsBus is /dev/tty.usbmodem575E0031751.
+>>> Reconnect the usb cable.
+
+python lerobot/common/robot_devices/motors/dynamixel.py
+>>> Finding all available ports for the DynamixelMotorsBus.
+>>> ['/dev/tty.usbmodem575E0032081', '/dev/tty.usbmodem575E0031751']
+>>> Remove the usb cable from your DynamixelMotorsBus and press Enter when done.
+... **Disconnect follower arm and press Enter**
+>>> The port of this DynamixelMotorsBus is /dev/tty.usbmodem575E0032081.
+>>> Reconnect the usb cable.
+```
+
+Then you can instantiate each arm by listing their motors with their name, motor index, and model. The initial motor index from factory for every motors is `1`. However, unique indices are required for these motors to function in a chain on a common bus. To this end, we set different indices and follow the ascendant convention starting from index `1` (e.g. "1, 2, 3, 4, 5, 6" ). These indices will be written inside the persisting memory of each motor during the first connection. Here is an example of what the instantiation looks like:
+```python
+from lerobot.common.robot_devices.motors.dynamixel import DynamixelMotorsBus
+
+leader_arm = 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_arm = 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"),
+    },
+)
+```
+
+**Configure and Connect**
+
+During the first connection of the motors, `DynamixelMotorsBus` automatically detects a mismatch between the present motor indices (all `1` by default) and the specified motor indices (e.g. "1, 2, 3, 4, 5, 6"). This triggers the configuration procedure which requires to unplug the power cord and motors, and to sequentially plug each motor again, starting from the closest to the bus. Because it is quite involved, we provide a youtube video for help. The output of the procedure looks like that:
+```python
+leader_arm.connect()
+
+TODO
+
+follower_arm.connect()
+
+TODO
+```
+
+Congrats! Now both arms are well configured and connected. Of course, next time you connect the arms, you won't have to follow configuration procedure ever again. For instance, let's try to disconnect and connect again like that:
+```python
+leader_arm.disconnect()
+leader_arm.connect()
+
+follower_arm.disconnect()
+follower_arm.connect()
+```
+
+**Read and Write**
+
+Just to get familiar with how `DynamixelMotorsBus` is used to command the motors, let's try to read from them. You should have something like:
+```python
+values = leader_arm.read("Present_Position")
+print(values)
+>>> TODO
+
+values = follower_arm.read("Present_Position")
+print(values)
+>>> TODO
+```
+
+The full address is `X_SERIES_CONTROL_TABLE`. TODO
+
+Now let's try to enable torque in the follower arm:
+```python
+from lerobot.common.robot_devices.motors.dynamixel import TorqueMode
+
+follower_arm.write("Torque_Enable", TorqueMode.ENABLED.value)
+
+values = follower_arm.read("Present_Position")
+
+values[0] += 10  # Try with positive or negative numbers
+follower_arm.write("Goal_Position", values)
+
+follower_arm.write("Goal_Position", values[0], "shoulder_pan")
+```
+
+
+### Teleoperate your Koch v1.1 with KochRobot
+
+**Instantiate**
+```python
+robot = KochRobot(
+    leader_arms={"main": leader_arm},
+    follower_arms={"main": follower_arm},
+    calibration_path=".cache/calibration/koch.pkl",
+)
+```
+
+**Calibrate and Connect**
+
+```
+robot.connect()
+>>>
+```
+
+```python
+degrees = leader_arms.read("Present_Position)
+print(degrees)
+>>>
+
+degrees = follower_arms.read("Present_Position)
+print(degrees)
+>>>
+```
+
+**Teleoperate**
+
+TODO: explain in pseudo code what the teleop is doing
+
+```python
+# Teleoperate for 60 seconds if running at 200 hz
+for _ in range(60*200):
+    robot.teleop_step()
+```
+
+TODO: explain in pseudo code what the teleop(record_data=True) is doing
+
+```python
+observation, action = robot.teleop_step(record_data=True)
+print(observation)
+>>>
+print(action)
+>>>
+```
+
+### Add your cameras with OpenCVCamera
+
+**Instantiate**
+
+The `OpenCVCamera` class allows to efficiently record images from cameras. It relies on opencv2 to communicate with the cameras. Most cameras are compatible. For more info, see the [Video I/O with OpenCV Overview](https://docs.opencv.org/4.x/d0/da7/videoio_overview.html).
+
+An `OpenCVCamera` instance requires a camera index (e.g. `OpenCVCamera(camera_index=0)`). When you only have one camera like a webcam of a laptop, the camera index is expected to be 0, but it might also be very different, and the camera index might change if you reboot your computer or re-plug your camera. This behavior depends on your operation system.
+
+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/opencv.py --images-dir outputs/images_from_opencv_cameras
+>>> TODO
+```
+
+When an `OpenCVCamera` 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 of the class:
+```python
+camera = OpenCVCamera(camera_index=0)
+camera.connect()
+color_image = camera.read()
+# when done using the camera, consider disconnecting
+camera.disconnect()
+```
+
+**Add to robot**
+
+TODO: explain that the cameras run asynchronously.
+
+```python
+del robot
+robot = KochRobot(
+    leader_arms={"main": leader_arm},
+    follower_arms={"main": follower_arm},
+    calibration_path=".cache/calibration/koch.pkl",
+    cameras={
+        "laptop": OpenCVCamera(0, fps=30, width=640, height=480),
+        "phone": OpenCVCamera(1, fps=30, width=640, height=480),
+    },
+)
+
+robot.connect()
+
+observation, action = robot.teleop_step(record_data=True)
+print(observation)
+>>>
+print(action)
+>>>
+```
+
+### Use `koch.yaml` and our `teleoperate` function
+
+See: `lerobot/configs/robot/koch.yaml`
+
+```bash
+python lerobot/scripts/control_robot.py teleoperate \
+  --robot-path lerobot/configs/robot/koch.yaml
+
+>>>
+```
+
+```bash
+python lerobot/scripts/control_robot.py teleoperate \
+  --robot-path lerobot/configs/robot/koch.yaml \
+  --robot-overrides \
+    leader_arms.main.port=/dev/tty.usbmodem575E0031751 \
+    follower_arms.main.port=/dev/tty.usbmodem575E0032081
+
+>>>
+```
+
+```bash
+python lerobot/scripts/control_robot.py teleoperate \
+  --robot-path lerobot/configs/robot/koch.yaml \
+  --robot-overrides \
+    leader_arms.main.port=/dev/tty.usbmodem575E0031751 \
+    follower_arms.main.port=/dev/tty.usbmodem575E0032081
+    '~cameras'
+```
+
+
+## 3. Record your Dataset and Visualize it
+
+TODO: ideally we could only do this
+
+```python
+from lerobot.scripts.control_robot import busy_wait
+
+fps = 30
+record_time_s = 60
+for _ in range(fps * record_time_s):
+    start_time = time.perf_counter()
+
+    observation, action = robot.teleop_step(record_data=True)
+
+    dt_s = time.perf_counter() - start_time
+    busy_wait(1 / fps - dt_s)
+```
+
+### Use `koch.yaml` and the `record` function
+
+TODO: We added ways to write the frames to disk in multiple thread
+We added warmap, reset time between episodes
+At the end we encode the frames into videos
+control
+if fail, re-record episode
+checkpointing
+We consolidate the data into a LeRobotDataset and upload on the hub.
+
+Here is an example for 1 episode
+```bash
+python lerobot/scripts/control_robot.py record \
+    --fps 30 \
+    --root /tmp/data \
+    --repo-id $USER/koch_test \
+    --num-episodes 10 \
+    --run-compute-stats 1
+```
+
+TODO: USER HF, make sure you can push
+
+### Replay episode on your robot with the `replay` function
+
+```bash
+python lerobot/scripts/control_robot.py replay \
+    --fps 30 \
+    --root /tmp/data \
+    --repo-id $USER/koch_test \
+    --episode 0
+```
+
+Note: TODO
+```bash
+export DATA_DIR=data
+```
+
+### Visualize all episodes
+
+```bash
+python lerobot/scripts/visualize_dataset.py \
+    --repo-id $USER/koch_test
+```
+
+
+## 4. Train a policy on your data
+
+### Use our `train` script
+
+```bash
+python lerobot/scripts/train.py \
+    policy=act_koch_real \
+    env=koch_real \
+    dataset_repo_id=$USER/koch_pick_place_lego \
+    hydra.run.dir=outputs/train/act_koch_real
+```
+
+TODO: image and plots of wandb
+
+```bash
+ckpt=100000
+huggingface-cli upload cadene/2024_07_27_act_koch_pick_place_1_lego_raph_nightly_${ckpt} \
+  outputs/train/2024_07_27_act_koch_pick_place_1_lego_raph_nightly/checkpoints/${ckpt}/pretrained_model
+```
+
+### Visualize predictions on training set
+
+```bash
+python lerobot/scripts/visualize_dataset_html.py \
+    --repo-id lerobot/koch_pick_place_1_lego \
+    --episodes 0 1 2 \
+    -p ../lerobot/outputs/train/2024_07_29_act_koch_pick_place_1_lego_mps/checkpoints/006000/pretrained_model
+```
+
+## 5. Evaluate your policy
+
+### Use our `record` function
+
+```bash
+python lerobot/scripts/control_robot.py record \
+    --fps 30 \
+    --root /tmp/data \
+    --repo-id $USER/eval_koch_test \
+    --num-episodes 10 \
+    --run-compute-stats 1
+    -p ../lerobot/outputs/train/2024_07_29_act_koch_pick_place_1_lego_mps/checkpoints/006000/pretrained_model
+```
+
+### Visualize evaluation afterwards
+
+```bash
+python lerobot/scripts/visualize_dataset.py \
+    --repo-id $USER/koch_test
+```
+
+
+## What's next?
+
+### More datasets
+
+Collect a slightly more difficult dataset, like grasping 5 lego blocks in a row, and co-train on it
+
+###
+
+
+- Improve the dataset

lerobot/common/robot_devices/cameras/opencv.py

@@ -5,6 +5,7 @@ This file contains utilities for recording frames from cameras. For more info lo
 import argparse
 import concurrent.futures
 import math
+import platform
 import shutil
 import threading
 import time
@@ -33,8 +34,18 @@ MAX_OPENCV_INDEX = 60
 
 
 def find_camera_indices(raise_when_empty=False, max_index_search_range=MAX_OPENCV_INDEX):
+    if platform.system() == "Linux":
+        # Linux uses camera ports
+        print("Linux detected. Finding available camera indices through scanning '/dev/video*' ports")
+        possible_camera_ids = [int(port.replace("/dev/video", "")) for port in Path("/dev").glob("video*")]
+    else:
+        print(
+            f"Mac or Windows dtected. Finding available camera indices through scanning all indices from 0 to {MAX_OPENCV_INDEX}"
+        )
+        possible_camera_ids = range(max_index_search_range)
+
     camera_ids = []
-    for camera_idx in range(max_index_search_range):
+    for camera_idx in possible_camera_ids:
         camera = cv2.VideoCapture(camera_idx)
         is_open = camera.isOpened()
         camera.release()
@@ -59,10 +70,9 @@ def save_image(img_array, camera_index, frame_index, images_dir):
 
 
 def save_images_from_cameras(
-    images_dir: Path, camera_ids=None, fps=None, width=None, height=None, record_time_s=2
+    images_dir: Path, camera_ids: list[int] | None = None, fps=None, width=None, height=None, record_time_s=2
 ):
     if camera_ids is None:
-        print("Finding available camera indices")
         camera_ids = find_camera_indices()
 
     print("Connecting cameras")
@@ -75,9 +85,7 @@ def save_images_from_cameras(
         )
         cameras.append(camera)
 
-    images_dir = Path(
-        images_dir,
-    )
+    images_dir = Path(images_dir)
     if images_dir.exists():
         shutil.rmtree(
             images_dir,
@@ -160,7 +168,7 @@ class OpenCVCamera:
     When an OpenCVCamera 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 of the class:
+    Example of usage:
     ```python
     camera = OpenCVCamera(camera_index=0)
     camera.connect()
@@ -194,11 +202,6 @@ class OpenCVCamera:
         self.height = config.height
         self.color_mode = config.color_mode
 
-        if not isinstance(self.camera_index, int):
-            raise ValueError(
-                f"Camera index must be provided as an int, but {self.camera_index} was given instead."
-            )
-
         self.camera = None
         self.is_connected = False
         self.thread = None
@@ -212,7 +215,13 @@ class OpenCVCamera:
 
         # First create a temporary camera trying to access `camera_index`,
         # and verify it is a valid camera by calling `isOpened`.
-        tmp_camera = cv2.VideoCapture(self.camera_index)
+
+        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)
+
         is_camera_open = tmp_camera.isOpened()
         # Release camera to make it accessible for `find_camera_indices`
         del tmp_camera
@@ -232,7 +241,10 @@ class OpenCVCamera:
         # Secondly, create the camera that will be used downstream.
         # Note: For some unknown reason, calling `isOpened` blocks the camera which then
         # needs to be re-created.
-        self.camera = cv2.VideoCapture(self.camera_index)
+        if platform.system() == "Linux":
+            self.camera = cv2.VideoCapture(f"/dev/video{self.camera_index}")
+        else:
+            self.camera = cv2.VideoCapture(self.camera_index)
 
         if self.fps is not None:
             self.camera.set(cv2.CAP_PROP_FPS, self.fps)

lerobot/common/robot_devices/cameras/utils.py

@@ -2,6 +2,7 @@ from pathlib import Path
 from typing import Protocol
 
 import cv2
+import einops
 import numpy as np
 
 
@@ -39,6 +40,16 @@ def save_depth_image(depth, path, write_shape=False):
     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

@@ -5,6 +5,7 @@ from copy import deepcopy
 from pathlib import Path
 
 import numpy as np
+import tqdm
 from dynamixel_sdk import (
     COMM_SUCCESS,
     DXL_HIBYTE,
@@ -21,9 +22,11 @@ from lerobot.common.robot_devices.utils import RobotDeviceAlreadyConnectedError,
 from lerobot.common.utils.utils import capture_timestamp_utc
 
 PROTOCOL_VERSION = 2.0
-BAUD_RATE = 1_000_000
+BAUDRATE = 1_000_000
 TIMEOUT_MS = 1000
 
+MAX_ID_RANGE = 252
+
 # 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
@@ -86,6 +89,16 @@ X_SERIES_CONTROL_TABLE = {
     "Present_Temperature": (146, 1),
 }
 
+X_SERIES_BAUDRATE_TABLE = {
+    0: 9_600,
+    1: 57_600,
+    2: 115_200,
+    3: 1_000_000,
+    4: 2_000_000,
+    5: 3_000_000,
+    6: 4_000_000,
+}
+
 CALIBRATION_REQUIRED = ["Goal_Position", "Present_Position"]
 CONVERT_UINT32_TO_INT32_REQUIRED = ["Goal_Position", "Present_Position"]
 
@@ -96,9 +109,78 @@ 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 = {
+    "x_series": 4096,
+    "xl330-m077": 4096,
+    "xl330-m288": 4096,
+    "xl430-w250": 4096,
+    "xm430-w350": 4096,
+    "xm540-w270": 4096,
+    "xc430-w150": 4096,
+}
+
+MODEL_BAUDRATE_TABLE = {
+    "x_series": X_SERIES_BAUDRATE_TABLE,
+    "xl330-m077": X_SERIES_BAUDRATE_TABLE,
+    "xl330-m288": X_SERIES_BAUDRATE_TABLE,
+    "xl430-w250": X_SERIES_BAUDRATE_TABLE,
+    "xm430-w350": X_SERIES_BAUDRATE_TABLE,
+    "xm540-w270": X_SERIES_BAUDRATE_TABLE,
+    "xc430-w150": X_SERIES_BAUDRATE_TABLE,
 }
 
 NUM_READ_RETRY = 10
+NUM_WRITE_RETRY = 10
+
+
+def convert_indices_to_baudrates(values: np.ndarray | list[int], models: list[str]):
+    assert len(values) == len(models)
+    for i in range(len(values)):
+        model = models[i]
+        index = values[i]
+        values[i] = MODEL_BAUDRATE_TABLE[model][index]
+    return values
+
+
+def convert_baudrates_to_indices(values: np.ndarray | list[int], models: list[str]):
+    assert len(values) == len(models)
+    for i in range(len(values)):
+        model = models[i]
+        brate = values[i]
+        table_values = list(MODEL_BAUDRATE_TABLE[model].values())
+        table_keys = list(MODEL_BAUDRATE_TABLE[model].keys())
+        values[i] = table_keys[table_values.index(brate)]
+    return values
+
+
+def convert_to_bytes(value, bytes):
+    # 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)),
+        ]
+    elif bytes == 2:
+        data = [
+            DXL_LOBYTE(DXL_LOWORD(value)),
+            DXL_HIBYTE(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)),
+        ]
+    else:
+        raise NotImplementedError(
+            f"Value of the number of bytes to be sent is expected to be in [1, 2, 4], but "
+            f"{bytes} is provided instead."
+        )
+    return data
 
 
 def get_group_sync_key(data_name, motor_names):
@@ -207,7 +289,6 @@ class DynamixelMotorsBus:
     >>> The port of this DynamixelMotorsBus is /dev/tty.usbmodem575E0031751.
     >>> Reconnect the usb cable.
     ```
-    To find the motor indices, use [DynamixelWizzard2](https://emanual.robotis.com/docs/en/software/dynamixel/dynamixel_wizard2).
 
     Example of usage for 1 motor connected to the bus:
     ```python
@@ -221,7 +302,9 @@ class DynamixelMotorsBus:
     )
     motors_bus.connect()
 
-    motors_bus.teleop_step()
+    degrees = motors_bus.read("Present_Position")
+
+    motors_bus.write("Goal_Position", degrees + 30)
 
     # when done, consider disconnecting
     motors_bus.disconnect()
@@ -233,6 +316,7 @@ class DynamixelMotorsBus:
         port: str,
         motors: dict[str, tuple[int, str]],
         extra_model_control_table: dict[str, list[tuple]] | None = None,
+        extra_model_resolution: dict[str, int] | None = None,
     ):
         self.port = port
         self.motors = motors
@@ -241,6 +325,10 @@ class DynamixelMotorsBus:
         if extra_model_control_table:
             self.model_ctrl_table.update(extra_model_control_table)
 
+        self.model_resolution = deepcopy(MODEL_RESOLUTION)
+        if extra_model_resolution:
+            self.model_resolution.update(extra_model_resolution)
+
         self.port_handler = None
         self.packet_handler = None
         self.calibration = None
@@ -268,52 +356,279 @@ class DynamixelMotorsBus:
             )
             raise
 
-        self.port_handler.setBaudRate(BAUD_RATE)
-        self.port_handler.setPacketTimeoutMillis(TIMEOUT_MS)
-
+        # Allow to read and write
         self.is_connected = True
 
+        self.port_handler.setPacketTimeoutMillis(TIMEOUT_MS)
+
+        # Set expected baudrate for the bus
+        self.set_bus_baudrate(BAUDRATE)
+
+        if not self.are_motors_configured():
+            print(
+                r"/!\ First, verify that all the cables are connected the proper way. If you detect an issue, before making any modification, unplug the power cord to not damage the motors. Rewire correctly. Then plug the power again and relaunch the script."
+            )
+            print(
+                r"/!\ Secondly, if the cables connection look correct and it is the first time that you use these motors, follow these manual steps to configure them."
+            )
+            input("Press Enter to configure your motors...")
+            print()
+            self.configure_motors()
+
+    def reconnect(self):
+        self.port_handler = PortHandler(self.port)
+        self.packet_handler = 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):
+        try:
+            return (self.motor_indices == self.read("ID")).all()
+        except ConnectionError as e:
+            print(e)
+            return False
+
+    def configure_motors(self):
+        # TODO(rcadene): This script assumes motors follow the X_SERIES baudrates
+
+        print("Scanning all baudrates and motor indices")
+        all_baudrates = set(X_SERIES_BAUDRATE_TABLE.values())
+        ids_per_baudrate = {}
+        for baudrate in all_baudrates:
+            self.set_bus_baudrate(baudrate)
+            present_ids = self.find_motor_indices()
+            if len(present_ids) > 0:
+                ids_per_baudrate[baudrate] = present_ids
+        print(f"Motor indices detected: {ids_per_baudrate}")
+        print()
+
+        possible_baudrates = list(ids_per_baudrate.keys())
+        possible_ids = list({idx for sublist in ids_per_baudrate.values() for idx in sublist})
+        untaken_ids = list(set(range(MAX_ID_RANGE)) - set(possible_ids) - set(self.motor_indices))
+
+        # Connect successively one motor to the chain and write a unique random index for each
+        for i in range(len(self.motors)):
+            self.disconnect()
+            print("1. Unplug the power cord")
+            print(
+                f"2. Plug/unplug minimal number of cables to only have the first {i+1} motor(s) ({self.motor_names[:i+1]}) connected."
+            )
+            print("3. Re-plug the power cord.")
+            input("Press Enter to continue...")
+            print()
+            self.reconnect()
+
+            if i > 0:
+                try:
+                    self._read_with_motor_ids(self.motor_models, untaken_ids[:i], "ID")
+                except ConnectionError:
+                    print(f"Failed to read from {untaken_ids[:i+1]}. Make sure the power cord is plugged in.")
+                    input("Press Enter to continue...")
+                    print()
+                    self.reconnect()
+
+            print("Scanning possible baudrates and motor indices")
+            motor_found = False
+            for baudrate in possible_baudrates:
+                self.set_bus_baudrate(baudrate)
+                present_ids = self.find_motor_indices(possible_ids)
+                if len(present_ids) == 1:
+                    present_idx = present_ids[0]
+                    print(f"Detected motor with index {present_idx}")
+
+                    if baudrate != BAUDRATE:
+                        print(f"Setting its baudrate to {BAUDRATE}")
+                        baudrate_idx = list(X_SERIES_BAUDRATE_TABLE.values()).index(BAUDRATE)
+
+                        # The write can fail, so we allow retries
+                        for _ in range(NUM_WRITE_RETRY):
+                            self._write_with_motor_ids(
+                                self.motor_models, present_idx, "Baud_Rate", baudrate_idx
+                            )
+                            time.sleep(0.5)
+                            self.set_bus_baudrate(BAUDRATE)
+                            try:
+                                present_baudrate_idx = self._read_with_motor_ids(
+                                    self.motor_models, present_idx, "Baud_Rate"
+                                )
+                            except ConnectionError:
+                                print("Failed to write baudrate. Retrying.")
+                                self.set_bus_baudrate(baudrate)
+                                continue
+                            break
+                        else:
+                            raise
+
+                        if present_baudrate_idx != baudrate_idx:
+                            raise OSError("Failed to write baudrate.")
+
+                    print(f"Setting its index to a temporary untaken index ({untaken_ids[i]})")
+                    self._write_with_motor_ids(self.motor_models, present_idx, "ID", untaken_ids[i])
+
+                    present_idx = self._read_with_motor_ids(self.motor_models, untaken_ids[i], "ID")
+                    if present_idx != untaken_ids[i]:
+                        raise OSError("Failed to write index.")
+
+                    motor_found = True
+                    break
+                elif len(present_ids) > 1:
+                    raise OSError(f"More than one motor detected ({present_ids}), but only one was expected.")
+
+            if not motor_found:
+                raise OSError(
+                    "No motor found, but one new motor expected. Verify power cord is plugged in and retry."
+                )
+            print()
+
+        print(f"Setting expected motor indices: {self.motor_indices}")
+        self.set_bus_baudrate(BAUDRATE)
+        self._write_with_motor_ids(
+            self.motor_models, untaken_ids[: len(self.motors)], "ID", self.motor_indices
+        )
+        print()
+
+        if (self.read("ID") != self.motor_indices).any():
+            raise OSError("Failed to write motors indices.")
+
+        print("Configuration is done!")
+
+    def find_motor_indices(self, possible_ids=None):
+        if possible_ids is None:
+            possible_ids = range(MAX_ID_RANGE)
+
+        indices = []
+        for idx in tqdm.tqdm(possible_ids):
+            try:
+                present_idx = self._read_with_motor_ids(self.motor_models, [idx], "ID")[0]
+            except ConnectionError:
+                continue
+
+            if idx != present_idx:
+                # sanity check
+                raise OSError(
+                    "Motor index used to communicate through the bus is not the same as the one present in the motor memory. The motor memory might be damaged."
+                )
+            indices.append(idx)
+
+        return indices
+
+    def set_bus_baudrate(self, baudrate):
+        present_bus_baudrate = self.port_handler.getBaudRate()
+        if present_bus_baudrate != baudrate:
+            print(f"Setting bus baud rate to {baudrate}. Previously {present_bus_baudrate}.")
+            self.port_handler.setBaudRate(baudrate)
+
+            if self.port_handler.getBaudRate() != baudrate:
+                raise OSError("Failed to write bus baud rate.")
+
     @property
-    def motor_names(self) -> list[int]:
+    def motor_names(self) -> list[str]:
         return list(self.motors.keys())
 
+    @property
+    def motor_models(self) -> list[str]:
+        return [model for _, model in self.motors.values()]
+
+    @property
+    def motor_indices(self) -> list[int]:
+        return [idx for idx, _ in self.motors.values()]
+
     def set_calibration(self, calibration: dict[str, tuple[int, bool]]):
         self.calibration = calibration
 
     def apply_calibration(self, values: np.ndarray | list, motor_names: list[str] | None):
-        if not self.calibration:
-            return values
+        """Convert from unsigned int32 joint position range [0, 2**32[ to the universal float32 centered degree range [-180.0, 180.0[
 
+        Joints values are original in [0, 2**32[ (unsigned int32). Each motor are expected to complete a full rotation
+        when given a goal position that is + or - their resolution. For instance, dynamixel xl330-m077 have a resolution of 4096, and
+        at any position in their original range, let's say the position 56734, they complete a full rotation clockwise by moving to 60830,
+        or anticlockwise by moving to 42638. The position in the original range is arbitrary and might change a lot between each motor.
+        To harmonize between motors of the same model, different robots, or even models of different brands, we propose to work
+        in the centered degree range [-180, 180[. This function first applies the pre-computed calibration to convert
+        from [0, 2**32[ to [-2048, 2048[, then divide by 2048.
+        """
         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]
 
-            if values[i] is not None:
-                if drive_mode:
-                    values[i] *= -1
-                values[i] += homing_offset
+            # 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 centered resolution 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[
+        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
 
         return values
 
     def revert_calibration(self, values: np.ndarray | list, motor_names: list[str] | None):
-        if not self.calibration:
-            return values
-
         if motor_names is None:
             motor_names = self.motor_names
 
+        # Convert from the universal float32 centered degree range [-180, 180[ to centered 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)
+
+        values = np.round(values).astype(np.int32)
+
+        # Convert from 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
 
-            if values[i] is not None:
-                values[i] -= homing_offset
-                if drive_mode:
-                    values[i] *= -1
+            # 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):
+        return_list = True
+        if not isinstance(motor_ids, list):
+            return_list = False
+            motor_ids = [motor_ids]
+
+        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)
+        for idx in motor_ids:
+            group.addParam(idx)
+
+        comm = group.txRxPacket()
+        if comm != 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)}"
+            )
+
+        values = []
+        for idx in motor_ids:
+            value = group.getData(idx, addr, bytes)
+            values.append(value)
+
+        if return_list:
+            return values
+        else:
+            return values[0]
+
     def read(self, data_name, motor_names: str | list[str] | None = None):
         if not self.is_connected:
             raise RobotDeviceNotConnectedError(
@@ -367,7 +682,7 @@ class DynamixelMotorsBus:
         if data_name in CONVERT_UINT32_TO_INT32_REQUIRED:
             values = values.astype(np.int32)
 
-        if data_name in CALIBRATION_REQUIRED:
+        if data_name in CALIBRATION_REQUIRED and self.calibration is not None:
             values = self.apply_calibration(values, motor_names)
 
         # log the number of seconds it took to read the data from the motors
@@ -380,6 +695,26 @@ class DynamixelMotorsBus:
 
         return values
 
+    def _write_with_motor_ids(self, motor_models, motor_ids, data_name, values):
+        if not isinstance(motor_ids, list):
+            motor_ids = [motor_ids]
+        if not isinstance(values, list):
+            values = [values]
+
+        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)
+        for idx, value in zip(motor_ids, values, strict=True):
+            data = convert_to_bytes(value, bytes)
+            group.addParam(idx, data)
+
+        comm = group.txPacket()
+        if comm != 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)}"
+            )
+
     def write(self, data_name, values: int | float | np.ndarray, motor_names: str | list[str] | None = None):
         if not self.is_connected:
             raise RobotDeviceNotConnectedError(
@@ -406,7 +741,7 @@ class DynamixelMotorsBus:
             motor_ids.append(motor_idx)
             models.append(model)
 
-        if data_name in CALIBRATION_REQUIRED:
+        if data_name in CALIBRATION_REQUIRED and self.calibration is not None:
             values = self.revert_calibration(values, motor_names)
 
         values = values.tolist()
@@ -422,30 +757,7 @@ class DynamixelMotorsBus:
             )
 
         for idx, value in zip(motor_ids, values, strict=True):
-            # 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)),
-                ]
-            elif bytes == 2:
-                data = [
-                    DXL_LOBYTE(DXL_LOWORD(value)),
-                    DXL_HIBYTE(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)),
-                ]
-            else:
-                raise NotImplementedError(
-                    f"Value of the number of bytes to be sent is expected to be in [1, 2, 4], but "
-                    f"{bytes} is provided instead."
-                )
-
+            data = convert_to_bytes(value, bytes)
             if init_group:
                 self.group_writers[group_key].addParam(idx, data)
             else:

lerobot/common/robot_devices/robots/factory.py

@@ -1,46 +1,7 @@
-def make_robot(name):
-    if name == "koch":
-        # TODO(rcadene): Add configurable robot from command line and yaml config
-        # TODO(rcadene): Add example with and without cameras
-        from lerobot.common.robot_devices.cameras.opencv import OpenCVCamera
-        from lerobot.common.robot_devices.motors.dynamixel import DynamixelMotorsBus
-        from lerobot.common.robot_devices.robots.koch import KochRobot
+import hydra
+from omegaconf import DictConfig
 
-        robot = KochRobot(
-            leader_arms={
-                "main": 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={
-                "main": 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": OpenCVCamera(0, fps=30, width=640, height=480),
-                "phone": OpenCVCamera(1, fps=30, width=640, height=480),
-            },
-        )
-    else:
-        raise ValueError(f"Robot '{name}' not found.")
 
+def make_robot(cfg: DictConfig):
+    robot = hydra.utils.instantiate(cfg)
     return robot

lerobot/common/robot_devices/robots/koch.py

@@ -8,119 +8,86 @@ import torch
 
 from lerobot.common.robot_devices.cameras.utils import Camera
 from lerobot.common.robot_devices.motors.dynamixel import (
-    DriveMode,
-    DynamixelMotorsBus,
     OperatingMode,
     TorqueMode,
 )
 from lerobot.common.robot_devices.motors.utils import MotorsBus
 from lerobot.common.robot_devices.utils import RobotDeviceAlreadyConnectedError, RobotDeviceNotConnectedError
 
-URL_HORIZONTAL_POSITION = {
-    "follower": "https://raw.githubusercontent.com/huggingface/lerobot/main/media/koch/follower_horizontal.png",
-    "leader": "https://raw.githubusercontent.com/huggingface/lerobot/main/media/koch/leader_horizontal.png",
-}
-URL_90_DEGREE_POSITION = {
-    "follower": "https://raw.githubusercontent.com/huggingface/lerobot/main/media/koch/follower_90_degree.png",
-    "leader": "https://raw.githubusercontent.com/huggingface/lerobot/main/media/koch/leader_90_degree.png",
-}
-
 ########################################################################
 # Calibration logic
 ########################################################################
 
-TARGET_HORIZONTAL_POSITION = np.array([0, -1024, 1024, 0, -1024, 0])
-TARGET_90_DEGREE_POSITION = np.array([1024, 0, 0, 1024, 0, -1024])
-GRIPPER_OPEN = np.array([-400])
+AVAILABLE_ROBOT_TYPES = ["koch", "aloha"]
+
+URL_TEMPLATE = "https://raw.githubusercontent.com/huggingface/lerobot/main/media/{robot}/{arm}_{position}.png"
+
+# In nominal range ]-2048, 2048[
+# First target position consists in moving koch arm to a straight horizontal position with gripper closed.
+KOCH_FIRST_POSITION = np.array([0, 0, 0, 0, 0, 0], dtype=np.int32)
+# Second target position consists in moving koch arm from the first target position by rotating every motor
+# by 90 degree. When the direction is ambiguous, always rotate on the right. Gripper is open, directed towards you.
+# TODO(rcadene): Take motor resolution into account instead of assuming 4096
+KOCH_SECOND_POSITION = np.array([1024, 1024, 1024, 1024, 1024, 1024], dtype=np.int32)
+
+# In nominal range ]-180, 180[
+KOCH_GRIPPER_OPEN = 35.156
+KOCH_REST_POSITION = np.array([0, 135, 90, 0, 0, KOCH_GRIPPER_OPEN])
+
+# In nominal range ]-2048, 2048[
+ALOHA_FIRST_POSITION = np.array([0, 0, 0, 0, 0, 0, 0, 0, 0], dtype=np.int32)
+ALOHA_SECOND_POSITION = np.array([1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 512], dtype=np.int32)
+# In nominal range ]-180, 180[
+ALOHA_GRIPPER_OPEN = 30
+ALOHA_REST_POSITION = np.array([0, 0, 0, 0, 0, 0, 0, 0, 0], dtype=np.int32)
 
 
-def apply_homing_offset(values: np.array, homing_offset: np.array) -> np.array:
-    for i in range(len(values)):
-        if values[i] is not None:
-            values[i] += homing_offset[i]
-    return values
+def assert_robot_type(robot_type):
+    if robot_type not in AVAILABLE_ROBOT_TYPES:
+        raise ValueError(robot_type)
 
 
-def apply_drive_mode(values: np.array, drive_mode: np.array) -> np.array:
-    for i in range(len(values)):
-        if values[i] is not None and drive_mode[i]:
-            values[i] = -values[i]
-    return values
+def get_first_position(robot_type):
+    if robot_type == "koch":
+        return KOCH_FIRST_POSITION
+    elif robot_type == "aloha":
+        return ALOHA_FIRST_POSITION
 
 
-def apply_calibration(values: np.array, homing_offset: np.array, drive_mode: np.array) -> np.array:
-    values = apply_drive_mode(values, drive_mode)
-    values = apply_homing_offset(values, homing_offset)
-    return values
+def get_second_position(robot_type):
+    if robot_type == "koch":
+        return KOCH_SECOND_POSITION
+    elif robot_type == "aloha":
+        return ALOHA_SECOND_POSITION
 
 
-def revert_calibration(values: np.array, homing_offset: np.array, drive_mode: np.array) -> np.array:
-    """
-    Transform working position into real position for the robot.
-    """
-    values = apply_homing_offset(
-        values,
-        np.array([-homing_offset if homing_offset is not None else None for homing_offset in homing_offset]),
-    )
-    values = apply_drive_mode(values, drive_mode)
-    return values
+def get_rest_position(robot_type):
+    if robot_type == "koch":
+        return KOCH_REST_POSITION
+    elif robot_type == "aloha":
+        return ALOHA_REST_POSITION
 
 
-def revert_appropriate_positions(positions: np.array, drive_mode: list[bool]) -> np.array:
-    for i, revert in enumerate(drive_mode):
-        if not revert and positions[i] is not None:
-            positions[i] = -positions[i]
-    return positions
+def assert_drive_mode(drive_mode):
+    # `drive_mode` is in [0,1] with 0 means original rotation direction for the motor, and 1 means inverted.
+    if not np.all(np.isin(drive_mode, [0, 1])):
+        raise ValueError(f"`drive_mode` contains values other than 0 or 1: ({drive_mode})")
 
 
-def compute_corrections(positions: np.array, drive_mode: list[bool], target_position: np.array) -> np.array:
-    correction = revert_appropriate_positions(positions, drive_mode)
-
-    for i in range(len(positions)):
-        if correction[i] is not None:
-            if drive_mode[i]:
-                correction[i] -= target_position[i]
-            else:
-                correction[i] += target_position[i]
-
-    return correction
+def apply_drive_mode(position, drive_mode):
+    assert_drive_mode(drive_mode)
+    # Convert `drive_mode` from [0, 1] with 0 indicates original rotation direction and 1 inverted,
+    # to [-1, 1] with 1 indicates original rotation direction and -1 inverted.
+    signed_drive_mode = -(drive_mode * 2 - 1)
+    position *= signed_drive_mode
+    return position
 
 
-def compute_nearest_rounded_positions(positions: np.array) -> np.array:
-    return np.array(
-        [
-            round(positions[i] / 1024) * 1024 if positions[i] is not None else None
-            for i in range(len(positions))
-        ]
-    )
-
-
-def compute_homing_offset(
-    arm: DynamixelMotorsBus, drive_mode: list[bool], target_position: np.array
-) -> np.array:
-    # Get the present positions of the servos
-    present_positions = apply_calibration(
-        arm.read("Present_Position"), np.array([0, 0, 0, 0, 0, 0]), drive_mode
-    )
-
-    nearest_positions = compute_nearest_rounded_positions(present_positions)
-    correction = compute_corrections(nearest_positions, drive_mode, target_position)
-    return correction
-
-
-def compute_drive_mode(arm: DynamixelMotorsBus, offset: np.array):
-    # Get current positions
-    present_positions = apply_calibration(
-        arm.read("Present_Position"), offset, np.array([False, False, False, False, False, False])
-    )
-
-    nearest_positions = compute_nearest_rounded_positions(present_positions)
-
-    # construct 'drive_mode' list comparing nearest_positions and TARGET_90_DEGREE_POSITION
-    drive_mode = []
-    for i in range(len(nearest_positions)):
-        drive_mode.append(nearest_positions[i] != TARGET_90_DEGREE_POSITION[i])
-    return drive_mode
+def compute_nearest_rounded_position(position, second_position):
+    # TODO(rcadene): Take motor resolution into account instead of assuming 4096
+    # Assumes 4096 steps for a full revolution for the motors
+    # Hence 90 degree is 1024 steps
+    return np.round(position / second_position).astype(position.dtype) * second_position
 
 
 def reset_arm(arm: MotorsBus):
@@ -132,55 +99,65 @@ def reset_arm(arm: MotorsBus):
     # 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"]
-    arm.write("Operating_Mode", OperatingMode.EXTENDED_POSITION.value, all_motors_except_gripper)
+    if len(all_motors_except_gripper) > 0:
+        arm.write("Operating_Mode", OperatingMode.EXTENDED_POSITION.value, all_motors_except_gripper)
 
     # TODO(rcadene): why?
     # Use 'position control current based' for gripper
     arm.write("Operating_Mode", OperatingMode.CURRENT_CONTROLLED_POSITION.value, "gripper")
 
-    # Make sure the native calibration (homing offset abd drive mode) is disabled, since we use our own calibration layer to be more generic
-    arm.write("Homing_Offset", 0)
-    arm.write("Drive_Mode", DriveMode.NON_INVERTED.value)
 
-
-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):
     """Example of usage:
     ```python
-    run_arm_calibration(arm, "left", "follower")
+    run_arm_calibration(arm, "aloha", "left", "follower")
     ```
     """
     reset_arm(arm)
 
+    print(f"\nRunning calibration of {robot_type} {arm_name} {arm_type}...")
+
     # TODO(rcadene): document what position 1 mean
-    print(
-        f"Please move the '{name} {arm_type}' arm to the horizontal position (gripper fully closed, see {URL_HORIZONTAL_POSITION[arm_type]})"
-    )
+    print("\nMove arm to first target position")
+    print("See: " + URL_TEMPLATE.format(robot=robot_type, arm=arm_type, position="first"))
     input("Press Enter to continue...")
 
-    horizontal_homing_offset = compute_homing_offset(
-        arm, [False, False, False, False, False, False], TARGET_HORIZONTAL_POSITION
-    )
+    # The first position zeros all motors, i.e. after calibration, if write goal position to be all 0,
+    # the robot will move to this first position.
+    first_position = get_first_position(robot_type)
+    # The second position rotates all motors with 90 degrees angle clock-wise from the perspective of the first motor or the preceeding motor in the chain.
+    # Note: if 90 degree rotation cannot be achieved (e.g. gripper of Aloha), then it will rotate to 45 degrees.
+    second_position = get_second_position(robot_type)
+
+    # Compute homing offset so that `present_position + homing_offset ~= target_position`
+    position = arm.read("Present_Position")
+    position = compute_nearest_rounded_position(position, second_position)
+    homing_offset = first_position - position
 
     # TODO(rcadene): document what position 2 mean
-    print(
-        f"Please move the '{name} {arm_type}' arm to the 90 degree position (gripper fully open, see {URL_90_DEGREE_POSITION[arm_type]})"
-    )
+    print("\nMove arm to second target position")
+    print("See: " + URL_TEMPLATE.format(robot=robot_type, arm=arm_type, position="second"))
     input("Press Enter to continue...")
 
-    drive_mode = compute_drive_mode(arm, horizontal_homing_offset)
-    homing_offset = compute_homing_offset(arm, drive_mode, TARGET_90_DEGREE_POSITION)
+    # Find drive mode by rotating each motor by 90 degree.
+    # After applying homing offset, if position equals target position, then drive mode is 0,
+    # to indicate an original rotation direction for the motor ; else, drive mode is 1,
+    # to indicate an inverted rotation direction.
+    position = arm.read("Present_Position")
+    position += homing_offset
+    position = compute_nearest_rounded_position(position, second_position)
+    drive_mode = (position != second_position).astype(np.int32)
 
-    # Invert offset for all drive_mode servos
-    for i in range(len(drive_mode)):
-        if drive_mode[i]:
-            homing_offset[i] = -homing_offset[i]
+    # 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, second_position)
+    homing_offset = second_position - position
 
-    print("Calibration is done!")
-
-    print("=====================================")
-    print("      HOMING_OFFSET: ", " ".join([str(i) for i in homing_offset]))
-    print("      DRIVE_MODE: ", " ".join([str(i) for i in drive_mode]))
-    print("=====================================")
+    print("\nMove arm to rest position")
+    print("See: " + URL_TEMPLATE.format(robot=robot_type, arm=arm_type, position="rest"))
+    input("Press Enter to continue...")
+    print()
 
     return homing_offset, drive_mode
 
@@ -199,11 +176,15 @@ class KochRobotConfig:
     ```
     """
 
+    robot_type: str = "koch"
     # Define all components of the robot
     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: {})
 
+    def __post_init__(self):
+        assert_robot_type(self.robot_type)
+
 
 class KochRobot:
     # TODO(rcadene): Implement force feedback
@@ -261,12 +242,12 @@ class KochRobot:
     Example of highest frequency data collection with cameras:
     ```python
     # Defines how to communicate with 2 cameras connected to the computer.
-    # Here, the webcam of the mackbookpro and the iphone (connected in USB to the macbookpro)
+    # Here, the webcam of the laptop and the phone (connected in USB to the laptop)
     # can be reached respectively using the camera indices 0 and 1. These indices can be
     # arbitrary. See the documentation of `OpenCVCamera` to find your own camera indices.
     cameras = {
-        "macbookpro": OpenCVCamera(camera_index=0, fps=30, width=640, height=480),
-        "iphone": OpenCVCamera(camera_index=1, fps=30, width=640, height=480),
+        "laptop": OpenCVCamera(camera_index=0, fps=30, width=640, height=480),
+        "phone": OpenCVCamera(camera_index=1, fps=30, width=640, height=480),
     }
 
     # Assumes leader and follower arms have been instantiated already (see first example)
@@ -311,6 +292,7 @@ class KochRobot:
         self.config = replace(config, **kwargs)
         self.calibration_path = Path(calibration_path)
 
+        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
@@ -330,7 +312,9 @@ class KochRobot:
 
         # Connect the arms
         for name in self.follower_arms:
+            print(f"Connecting {name} follower arm.")
             self.follower_arms[name].connect()
+            print(f"Connecting {name} leader arm.")
             self.leader_arms[name].connect()
 
         # Reset the arms and load or run calibration
@@ -357,22 +341,36 @@ class KochRobot:
         for name in self.leader_arms:
             self.leader_arms[name].set_calibration(calibration[f"leader_{name}"])
 
-        # 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")
+        for name in self.leader_arms:
+            values = self.leader_arms[name].read("Present_Position")
+            # if (values < -180).any() or (values >= 180).any():
+            #     raise ValueError(
+            #         f"At least one of the motor of the {name} leader arm has a joint value outside of its centered degree range of ]-180, 180[."
+            #         'This "jump of range" can be caused by a hardware issue, or you might have unexpectedly completed a full rotation of the motor '
+            #         "during manipulation or transportation of your robot. "
+            #         f"The values and motors: {values} {self.leader_arms[name].motor_names}.\n"
+            #         "Rotate the arm to fit the range ]-180, 180[ and relaunch the script, or recalibrate all motors by setting a different "
+            #         "calibration path during the instatiation of your robot (e.g. `--robot-overrides calibration_path=.cache/calibration/koch_v2.pkl`)"
+            #     )
 
         # Enable torque on all motors of the follower arms
         for name in self.follower_arms:
             self.follower_arms[name].write("Torque_Enable", 1)
 
-        # 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.
-        for name in self.leader_arms:
-            self.leader_arms[name].write("Torque_Enable", 1, "gripper")
-            self.leader_arms[name].write("Goal_Position", GRIPPER_OPEN, "gripper")
+        # Custom setup for each robot type
+        if self.robot_type == "koch":
+            # 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.
+            for name in self.leader_arms:
+                self.leader_arms[name].write("Torque_Enable", 1, "gripper")
+                self.leader_arms[name].write("Goal_Position", KOCH_GRIPPER_OPEN, "gripper")
+
+            # 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
+            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")
 
         # Connect the cameras
         for name in self.cameras:
@@ -384,14 +382,18 @@ class KochRobot:
         calibration = {}
 
         for name in self.follower_arms:
-            homing_offset, drive_mode = run_arm_calibration(self.follower_arms[name], name, "follower")
+            homing_offset, drive_mode = run_arm_calibration(
+                self.follower_arms[name], self.robot_type, name, "follower"
+            )
 
             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])
 
         for name in self.leader_arms:
-            homing_offset, drive_mode = run_arm_calibration(self.leader_arms[name], name, "leader")
+            homing_offset, drive_mode = run_arm_calibration(
+                self.leader_arms[name], self.robot_type, name, "leader"
+            )
 
             calibration[f"leader_{name}"] = {}
             for idx, motor_name in enumerate(self.leader_arms[name].motor_names):
@@ -407,7 +409,7 @@ class KochRobot:
                 "KochRobot is not connected. You need to run `robot.connect()`."
             )
 
-        # Prepare to assign the positions of the leader to the follower
+        # Prepare to assign the position of the leader to the follower
         leader_pos = {}
         for name in self.leader_arms:
             now = time.perf_counter()
@@ -500,11 +502,7 @@ class KochRobot:
         obs_dict = {}
         obs_dict["observation.state"] = torch.from_numpy(state)
         for name in self.cameras:
-            # Convert to pytorch format: channel first and float32 in [0,1]
-            img = torch.from_numpy(images[name])
-            img = img.type(torch.float32) / 255
-            img = img.permute(2, 0, 1).contiguous()
-            obs_dict[f"observation.images.{name}"] = img
+            obs_dict[f"observation.images.{name}"] = torch.from_numpy(images[name])
         return obs_dict
 
     def send_action(self, action: torch.Tensor):

lerobot/common/utils/utils.py

@@ -158,6 +158,7 @@ def init_hydra_config(config_path: str, overrides: list[str] | None = None) -> D
         version_base="1.2",
     )
     cfg = hydra.compose(Path(config_path).stem, overrides)
+
     return cfg
 
 

lerobot/configs/robot/aloha.yaml

@@ -0,0 +1,91 @@
+_target_: lerobot.common.robot_devices.robots.koch.KochRobot
+robot_type: aloha
+calibration_path: .cache/calibration/aloha.pkl
+leader_arms:
+  left:
+    _target_: lerobot.common.robot_devices.motors.dynamixel.DynamixelMotorsBus
+    port: /dev/ttyDXL_master_left
+    # port: /dev/tty.usbserial-FT89FM77
+    motors:  # window_x
+      # name: (index, model)
+      waist: [1, xm430-w350]
+      shoulder_a: [2, xm430-w350]
+      shoulder_b: [3, xm430-w350]
+      elbow_a: [4, xm430-w350]
+      elbow_b: [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_master_right
+    # port: /dev/tty.usbserial-FT891KPN
+    motors:  # window_x
+      # name: (index, model)
+      waist: [1, xm430-w350]
+      shoulder_a: [2, xm430-w350]
+      shoulder_b: [3, xm430-w350]
+      elbow_a: [4, xm430-w350]
+      elbow_b: [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_puppet_left
+    # port: /dev/tty.usbserial-FT891KBG
+    motors:
+      # name: [index, model]
+      waist: [1, xm540-w270]
+      shoulder_a: [2, xm540-w270]
+      shoulder_b: [3, xm540-w270]
+      elbow_a: [4, xm540-w270]
+      elbow_b: [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_puppet_right
+    # port: /dev/tty.usbserial-FT89FM09
+    motors:
+      # name: [index, model]
+      waist: [1, xm540-w270]
+      shoulder_a: [2, xm540-w270]
+      shoulder_b: [3, xm540-w270]
+      elbow_a: [4, xm540-w270]
+      elbow_b: [5, xm540-w270]
+      forearm_roll: [6, xm540-w270]
+      wrist_angle: [7, xm540-w270]
+      wrist_rotate: [8, xm430-w350]
+      gripper: [9, xm430-w350]
+cameras:
+  cam_high:
+    _target_: lerobot.common.robot_devices.cameras.opencv.OpenCVCamera
+    camera_index: 0
+    fps: 30
+    width: 640
+    height: 480
+  cam_low:
+    _target_: lerobot.common.robot_devices.cameras.opencv.OpenCVCamera
+    camera_index: 1
+    fps: 30
+    width: 640
+    height: 480
+  cam_left_wrist:
+    _target_: lerobot.common.robot_devices.cameras.opencv.OpenCVCamera
+    camera_index: 2
+    fps: 30
+    width: 640
+    height: 480
+  cam_right_wrist:
+    _target_: lerobot.common.robot_devices.cameras.opencv.OpenCVCamera
+    camera_index: 3
+    fps: 30
+    width: 640
+    height: 480

lerobot/configs/robot/koch.yaml

@@ -0,0 +1,40 @@
+_target_: lerobot.common.robot_devices.robots.koch.KochRobot
+calibration_path: .cache/calibration/koch.pkl
+robot_type: koch
+leader_arms:
+  main:
+    _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:
+  main:
+    _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

lerobot/scripts/control_robot.py

@@ -4,6 +4,9 @@ Examples of usage:
 - Unlimited teleoperation at highest frequency (~200 Hz is expected), to exit with CTRL+C:
 ```bash
 python lerobot/scripts/control_robot.py teleoperate
+
+# Remove the cameras from the robot definition. They are not used in 'teleoperate' anyway.
+python lerobot/scripts/control_robot.py teleoperate --robot-overrides '~cameras'
 ```
 
 - Unlimited teleoperation at a limited frequency of 30 Hz, to simulate data recording frequency:
@@ -14,7 +17,7 @@ python lerobot/scripts/control_robot.py teleoperate \
 
 - Record one episode in order to test replay:
 ```bash
-python lerobot/scripts/control_robot.py record_dataset \
+python lerobot/scripts/control_robot.py record \
     --fps 30 \
     --root tmp/data \
     --repo-id $USER/koch_test \
@@ -32,7 +35,7 @@ python lerobot/scripts/visualize_dataset.py \
 
 - Replay this test episode:
 ```bash
-python lerobot/scripts/control_robot.py replay_episode \
+python lerobot/scripts/control_robot.py replay \
     --fps 30 \
     --root tmp/data \
     --repo-id $USER/koch_test \
@@ -42,12 +45,11 @@ python lerobot/scripts/control_robot.py replay_episode \
 - Record a full dataset in order to train a policy, with 2 seconds of warmup,
 30 seconds of recording for each episode, and 10 seconds to reset the environment in between episodes:
 ```bash
-python lerobot/scripts/control_robot.py record_dataset \
+python lerobot/scripts/control_robot.py record \
     --fps 30 \
     --root data \
     --repo-id $USER/koch_pick_place_lego \
     --num-episodes 50 \
-    --run-compute-stats 1 \
     --warmup-time-s 2 \
     --episode-time-s 30 \
     --reset-time-s 10
@@ -74,7 +76,14 @@ DATA_DIR=data python lerobot/scripts/train.py \
 
 - Run the pretrained policy on the robot:
 ```bash
-python lerobot/scripts/control_robot.py run_policy \
+python lerobot/scripts/control_robot.py record \
+    --fps 30 \
+    --root data \
+    --repo-id $USER/eval_act_koch_real \
+    --num-episodes 10 \
+    --warmup-time-s 2 \
+    --episode-time-s 30 \
+    --reset-time-s 10
     -p outputs/train/act_koch_real/checkpoints/080000/pretrained_model
 ```
 """
@@ -88,8 +97,10 @@ import platform
 import shutil
 import time
 from contextlib import nullcontext
+from functools import cache
 from pathlib import Path
 
+import cv2
 import torch
 import tqdm
 from huggingface_hub import create_branch
@@ -179,7 +190,8 @@ def log_control_info(robot, dt_s, episode_index=None, frame_index=None, fps=None
     logging.info(info_str)
 
 
-def get_is_headless():
+@cache
+def is_headless():
     if platform.system() == "Linux":
         display = os.environ.get("DISPLAY")
         if display is None or display == "":
@@ -213,8 +225,10 @@ def teleoperate(robot: Robot, fps: int | None = None, teleop_time_s: float | Non
             break
 
 
-def record_dataset(
+def record(
     robot: Robot,
+    policy: torch.nn.Module | None = None,
+    hydra_cfg: DictConfig | None = None,
     fps: int | None = None,
     root="data",
     repo_id="lerobot/debug",
@@ -255,32 +269,10 @@ def record_dataset(
     else:
         episode_index = 0
 
-    is_headless = get_is_headless()
-
-    # Execute a few seconds without recording data, to give times
-    # to the robot devices to connect and start synchronizing.
-    timestamp = 0
-    start_time = time.perf_counter()
-    is_warmup_print = False
-    while timestamp < warmup_time_s:
-        if not is_warmup_print:
-            logging.info("Warming up (no data recording)")
-            os.system('say "Warmup" &')
-            is_warmup_print = True
-
-        now = time.perf_counter()
-        observation, action = robot.teleop_step(record_data=True)
-
-        if not is_headless:
-            image_keys = [key for key in observation if "image" in key]
-
-        dt_s = time.perf_counter() - now
-        busy_wait(1 / fps - dt_s)
-
-        dt_s = time.perf_counter() - now
-        log_control_info(robot, dt_s, fps=fps)
-
-        timestamp = time.perf_counter() - start_time
+    if is_headless():
+        logging.info(
+            "Headless environment detected. Display cameras on screen and keyboard inputs will not be available."
+        )
 
     # Allow to exit early while recording an episode or resetting the environment,
     # by tapping the right arrow key '->'. This might require a sudo permission
@@ -290,9 +282,7 @@ def record_dataset(
     stop_recording = False
 
     # Only import pynput if not in a headless environment
-    if is_headless:
-        logging.info("Headless environment detected. Keyboard input will not be available.")
-    else:
+    if not is_headless():
         from pynput import keyboard
 
         def on_press(key):
@@ -315,6 +305,53 @@ def record_dataset(
         listener = keyboard.Listener(on_press=on_press)
         listener.start()
 
+    # Load policy if any
+    if policy is not None:
+        # Check device is available
+        device = get_safe_torch_device(hydra_cfg.device, log=True)
+
+        policy.eval()
+        policy.to(device)
+
+        torch.backends.cudnn.benchmark = True
+        torch.backends.cuda.matmul.allow_tf32 = True
+        set_global_seed(hydra_cfg.seed)
+
+        # override fps using policy fps
+        fps = hydra_cfg.env.fps
+
+    # Execute a few seconds without recording data, to give times
+    # to the robot devices to connect and start synchronizing.
+    timestamp = 0
+    start_time = time.perf_counter()
+    is_warmup_print = False
+    while timestamp < warmup_time_s:
+        if not is_warmup_print:
+            logging.info("Warming up (no data recording)")
+            os.system('say "Warmup" &')
+            is_warmup_print = True
+
+        now = time.perf_counter()
+
+        if policy is None:
+            observation, action = robot.teleop_step(record_data=True)
+        else:
+            observation = robot.capture_observation()
+
+        if 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))
+            cv2.waitKey(1)
+
+        dt_s = time.perf_counter() - now
+        busy_wait(1 / fps - dt_s)
+
+        dt_s = time.perf_counter() - now
+        log_control_info(robot, dt_s, fps=fps)
+
+        timestamp = time.perf_counter() - start_time
+
     # 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.
@@ -330,7 +367,11 @@ def record_dataset(
             start_time = time.perf_counter()
             while timestamp < episode_time_s:
                 now = time.perf_counter()
-                observation, action = robot.teleop_step(record_data=True)
+
+                if policy is None:
+                    observation, action = robot.teleop_step(record_data=True)
+                else:
+                    observation = robot.capture_observation()
 
                 image_keys = [key for key in observation if "image" in key]
                 not_image_keys = [key for key in observation if "image" not in key]
@@ -342,11 +383,44 @@ def record_dataset(
                         )
                     ]
 
+                if 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))
+                    cv2.waitKey(1)
+
                 for key in not_image_keys:
                     if key not in ep_dict:
                         ep_dict[key] = []
                     ep_dict[key].append(observation[key])
 
+                if policy is not None:
+                    with (
+                        torch.inference_mode(),
+                        torch.autocast(device_type=device.type)
+                        if device.type == "cuda" and hydra_cfg.use_amp
+                        else nullcontext(),
+                    ):
+                        # Convert to pytorch format: channel first and float32 in [0,1] with batch dimension
+                        for name in observation:
+                            if "image" in name:
+                                observation[name] = observation[name].type(torch.float32) / 255
+                                observation[name] = observation[name].permute(2, 0, 1).contiguous()
+                            observation[name] = observation[name].unsqueeze(0)
+
+                        if device.type == "mps":
+                            for name in observation:
+                                observation[name] = observation[name].to(device)
+
+                        action = policy.select_action(observation)
+
+                        # remove batch dimension
+                        action = action.squeeze(0)
+                        action = action.to("cpu")
+
+                    robot.send_action(action)
+                    action = {"action": action}
+
                 for key in action:
                     if key not in ep_dict:
                         ep_dict[key] = []
@@ -434,7 +508,7 @@ def record_dataset(
             if is_last_episode:
                 logging.info("Done recording")
                 os.system('say "Done recording"')
-                if not is_headless:
+                if not is_headless():
                     listener.stop()
 
                 logging.info("Waiting for threads writing the images on disk to terminate...")
@@ -444,6 +518,10 @@ def record_dataset(
                     pass
                 break
 
+    robot.disconnect()
+    if not is_headless():
+        cv2.destroyAllWindows()
+
     num_episodes = episode_index
 
     logging.info("Encoding videos")
@@ -495,6 +573,7 @@ def record_dataset(
         stats = compute_stats(lerobot_dataset)
         lerobot_dataset.stats = stats
     else:
+        stats = {}
         logging.info("Skipping computation of the dataset statistrics")
 
     hf_dataset = hf_dataset.with_format(None)  # to remove transforms that cant be saved
@@ -516,7 +595,7 @@ def record_dataset(
     return lerobot_dataset
 
 
-def replay_episode(robot: Robot, episode: int, fps: int | None = None, root="data", repo_id="lerobot/debug"):
+def replay(robot: Robot, episode: int, fps: int | None = None, root="data", repo_id="lerobot/debug"):
     # TODO(rcadene): Add option to record logs
     local_dir = Path(root) / repo_id
     if not local_dir.exists():
@@ -546,61 +625,6 @@ def replay_episode(robot: Robot, episode: int, fps: int | None = None, root="dat
         log_control_info(robot, dt_s, fps=fps)
 
 
-def run_policy(robot: Robot, policy: torch.nn.Module, hydra_cfg: DictConfig, run_time_s: float | None = None):
-    # TODO(rcadene): Add option to record eval dataset and logs
-
-    # Check device is available
-    device = get_safe_torch_device(hydra_cfg.device, log=True)
-
-    policy.eval()
-    policy.to(device)
-
-    torch.backends.cudnn.benchmark = True
-    torch.backends.cuda.matmul.allow_tf32 = True
-    set_global_seed(hydra_cfg.seed)
-
-    fps = hydra_cfg.env.fps
-
-    if not robot.is_connected:
-        robot.connect()
-
-    start_time = time.perf_counter()
-    while True:
-        now = time.perf_counter()
-
-        observation = robot.capture_observation()
-
-        with (
-            torch.inference_mode(),
-            torch.autocast(device_type=device.type)
-            if device.type == "cuda" and hydra_cfg.use_amp
-            else nullcontext(),
-        ):
-            # add batch dimension to 1
-            for name in observation:
-                observation[name] = observation[name].unsqueeze(0)
-
-            if device.type == "mps":
-                for name in observation:
-                    observation[name] = observation[name].to(device)
-
-            action = policy.select_action(observation)
-
-            # remove batch dimension
-            action = action.squeeze(0)
-
-        robot.send_action(action.to("cpu"))
-
-        dt_s = time.perf_counter() - now
-        busy_wait(1 / fps - dt_s)
-
-        dt_s = time.perf_counter() - now
-        log_control_info(robot, dt_s, fps=fps)
-
-        if run_time_s is not None and time.perf_counter() - start_time > run_time_s:
-            break
-
-
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     subparsers = parser.add_subparsers(dest="mode", required=True)
@@ -608,10 +632,16 @@ if __name__ == "__main__":
     # Set common options for all the subparsers
     base_parser = argparse.ArgumentParser(add_help=False)
     base_parser.add_argument(
-        "--robot",
+        "--robot-path",
         type=str,
-        default="koch",
-        help="Name of the robot provided to the `make_robot(name)` factory function.",
+        default="lerobot/configs/robot/koch.yaml",
+        help="Path to robot yaml file used to instantiate the robot using `make_robot` factory function.",
+    )
+    base_parser.add_argument(
+        "--robot-overrides",
+        type=str,
+        nargs="*",
+        help="Any key=value arguments to override config values (use dots for.nested=overrides)",
     )
 
     parser_teleop = subparsers.add_parser("teleoperate", parents=[base_parser])
@@ -619,7 +649,7 @@ if __name__ == "__main__":
         "--fps", type=none_or_int, default=None, help="Frames per second (set to None to disable)"
     )
 
-    parser_record = subparsers.add_parser("record_dataset", parents=[base_parser])
+    parser_record = subparsers.add_parser("record", parents=[base_parser])
     parser_record.add_argument(
         "--fps", type=none_or_int, default=None, help="Frames per second (set to None to disable)"
     )
@@ -638,19 +668,19 @@ if __name__ == "__main__":
     parser_record.add_argument(
         "--warmup-time-s",
         type=int,
-        default=2,
+        default=10,
         help="Number of seconds before starting data collection. It allows the robot devices to warmup and synchronize.",
     )
     parser_record.add_argument(
         "--episode-time-s",
         type=int,
-        default=10,
+        default=60,
         help="Number of seconds for data recording for each episode.",
     )
     parser_record.add_argument(
         "--reset-time-s",
         type=int,
-        default=5,
+        default=60,
         help="Number of seconds for resetting the environment after each episode.",
     )
     parser_record.add_argument("--num-episodes", type=int, default=50, help="Number of episodes to record.")
@@ -678,8 +708,23 @@ if __name__ == "__main__":
         default=0,
         help="By default, data recording is resumed. When set to 1, delete the local directory and start data recording from scratch.",
     )
+    parser_record.add_argument(
+        "-p",
+        "--pretrained-policy-name-or-path",
+        type=str,
+        help=(
+            "Either the repo ID of a model hosted on the Hub or a path to a directory containing weights "
+            "saved using `Policy.save_pretrained`."
+        ),
+    )
+    parser_record.add_argument(
+        "--policy-overrides",
+        type=str,
+        nargs="*",
+        help="Any key=value arguments to override config values (use dots for.nested=overrides)",
+    )
 
-    parser_replay = subparsers.add_parser("replay_episode", parents=[base_parser])
+    parser_replay = subparsers.add_parser("replay", parents=[base_parser])
     parser_replay.add_argument(
         "--fps", type=none_or_int, default=None, help="Frames per second (set to None to disable)"
     )
@@ -697,41 +742,38 @@ if __name__ == "__main__":
     )
     parser_replay.add_argument("--episode", type=int, default=0, help="Index of the episode to replay.")
 
-    parser_policy = subparsers.add_parser("run_policy", parents=[base_parser])
-    parser_policy.add_argument(
-        "-p",
-        "--pretrained-policy-name-or-path",
-        type=str,
-        help=(
-            "Either the repo ID of a model hosted on the Hub or a path to a directory containing weights "
-            "saved using `Policy.save_pretrained`."
-        ),
-    )
-    parser_policy.add_argument(
-        "overrides",
-        nargs="*",
-        help="Any key=value arguments to override config values (use dots for.nested=overrides)",
-    )
     args = parser.parse_args()
 
     init_logging()
 
     control_mode = args.mode
-    robot_name = args.robot
+    robot_path = args.robot_path
+    robot_overrides = args.robot_overrides
     kwargs = vars(args)
     del kwargs["mode"]
-    del kwargs["robot"]
+    del kwargs["robot_path"]
+    del kwargs["robot_overrides"]
+
+    robot_cfg = init_hydra_config(robot_path, robot_overrides)
+    robot = make_robot(robot_cfg)
 
-    robot = make_robot(robot_name)
     if control_mode == "teleoperate":
         teleoperate(robot, **kwargs)
-    elif control_mode == "record_dataset":
-        record_dataset(robot, **kwargs)
-    elif control_mode == "replay_episode":
-        replay_episode(robot, **kwargs)
 
-    elif control_mode == "run_policy":
-        pretrained_policy_path = get_pretrained_policy_path(args.pretrained_policy_name_or_path)
-        hydra_cfg = init_hydra_config(pretrained_policy_path / "config.yaml", args.overrides)
-        policy = make_policy(hydra_cfg=hydra_cfg, pretrained_policy_name_or_path=pretrained_policy_path)
-        run_policy(robot, policy, hydra_cfg)
+    elif control_mode == "record":
+        pretrained_policy_name_or_path = args.pretrained_policy_name_or_path
+        policy_overrides = args.policy_overrides
+        del kwargs["pretrained_policy_name_or_path"]
+        del kwargs["policy_overrides"]
+
+        policy_cfg = None
+        if pretrained_policy_name_or_path is not None:
+            pretrained_policy_path = get_pretrained_policy_path(pretrained_policy_name_or_path)
+            policy_cfg = init_hydra_config(pretrained_policy_path / "config.yaml", policy_overrides)
+            policy = make_policy(hydra_cfg=policy_cfg, pretrained_policy_name_or_path=pretrained_policy_path)
+            record(robot, policy, policy_cfg, **kwargs)
+        else:
+            record(robot, **kwargs)
+
+    elif control_mode == "replay":
+        replay(robot, **kwargs)

tests/conftest.py

@@ -15,7 +15,9 @@
 # limitations under the License.
 import pytest
 
-from .utils import DEVICE
+from lerobot.common.utils.utils import init_hydra_config
+
+from .utils import DEVICE, KOCH_ROBOT_CONFIG_PATH
 
 
 def pytest_collection_finish():
@@ -27,11 +29,12 @@ def is_koch_available():
     try:
         from lerobot.common.robot_devices.robots.factory import make_robot
 
-        robot = make_robot("koch")
+        robot_cfg = init_hydra_config(KOCH_ROBOT_CONFIG_PATH)
+        robot = make_robot(robot_cfg)
         robot.connect()
         del robot
         return True
     except Exception as e:
-        print("An alexander koch robot is not available.")
+        print("A koch robot is not available.")
         print(e)
         return False

tests/test_control_robot.py

@@ -3,13 +3,19 @@ from pathlib import Path
 from lerobot.common.policies.factory import make_policy
 from lerobot.common.robot_devices.robots.factory import make_robot
 from lerobot.common.utils.utils import init_hydra_config
-from lerobot.scripts.control_robot import record_dataset, replay_episode, run_policy, teleoperate
-from tests.utils import DEFAULT_CONFIG_PATH, DEVICE, require_koch
+from lerobot.scripts.control_robot import record, replay, teleoperate
+from tests.utils import DEFAULT_CONFIG_PATH, DEVICE, KOCH_ROBOT_CONFIG_PATH, require_koch
+
+
+def make_robot_():
+    robot_cfg = init_hydra_config(KOCH_ROBOT_CONFIG_PATH)
+    robot = make_robot(robot_cfg)
+    return robot
 
 
 @require_koch
 def test_teleoperate(request):
-    robot = make_robot("koch")
+    robot = make_robot_()
     teleoperate(robot, teleop_time_s=1)
     teleoperate(robot, fps=30, teleop_time_s=1)
     teleoperate(robot, fps=60, teleop_time_s=1)
@@ -17,20 +23,19 @@ def test_teleoperate(request):
 
 
 @require_koch
-def test_record_dataset_and_replay_episode_and_run_policy(tmpdir, request):
-    robot_name = "koch"
+def test_record_and_replay_and_policy(tmpdir, request):
     env_name = "koch_real"
     policy_name = "act_koch_real"
 
     root = Path(tmpdir)
     repo_id = "lerobot/debug"
 
-    robot = make_robot(robot_name)
-    dataset = record_dataset(
+    robot = make_robot_()
+    dataset = record(
         robot, fps=30, root=root, repo_id=repo_id, warmup_time_s=1, episode_time_s=1, num_episodes=2
     )
 
-    replay_episode(robot, episode=0, fps=30, root=root, repo_id=repo_id)
+    replay(robot, episode=0, fps=30, root=root, repo_id=repo_id)
 
     cfg = init_hydra_config(
         DEFAULT_CONFIG_PATH,
@@ -43,6 +48,6 @@ def test_record_dataset_and_replay_episode_and_run_policy(tmpdir, request):
 
     policy = make_policy(hydra_cfg=cfg, dataset_stats=dataset.stats)
 
-    run_policy(robot, policy, cfg, run_time_s=1)
+    record(robot, policy, cfg, run_time_s=1)
 
     del robot

tests/test_motors.py

@@ -1,33 +1,54 @@
+# TODO(rcadene): measure fps in nightly?
+# TODO(rcadene): test logs
+# TODO(rcadene): test calibration
+# TODO(rcadene): add compatibility with other motors bus
+
 import time
 
+import hydra
 import numpy as np
 import pytest
 
 from lerobot.common.robot_devices.utils import RobotDeviceAlreadyConnectedError, RobotDeviceNotConnectedError
-from tests.utils import require_koch
+from lerobot.common.utils.utils import init_hydra_config
+from tests.utils import KOCH_ROBOT_CONFIG_PATH, require_koch
+
+
+def make_motors_bus():
+    robot_cfg = init_hydra_config(KOCH_ROBOT_CONFIG_PATH)
+    # Instantiating a common motors structure.
+    # Here the one from Alexander Koch follower arm.
+    motors_bus = hydra.utils.instantiate(robot_cfg.leader_arms.main)
+    return motors_bus
+
+
+@require_koch
+def test_find_port(request):
+    from lerobot.common.robot_devices.motors.dynamixel import find_port
+
+    find_port()
+
+
+@require_koch
+def test_configure_motors_all_ids_1(request):
+    # This test expect the configuration was already correct.
+    motors_bus = make_motors_bus()
+    motors_bus.connect()
+    motors_bus.write("Baud_Rate", [0] * len(motors_bus.motors))
+    motors_bus.set_bus_baudrate(9_600)
+    motors_bus.write("ID", [1] * len(motors_bus.motors))
+    del motors_bus
+
+    # Test configure
+    motors_bus = make_motors_bus()
+    motors_bus.connect()
+    assert motors_bus.are_motors_configured()
+    del motors_bus
 
 
 @require_koch
 def test_motors_bus(request):
-    # TODO(rcadene): measure fps in nightly?
-    # TODO(rcadene): test logs
-    # TODO(rcadene): test calibration
-    # TODO(rcadene): add compatibility with other motors bus
-    from lerobot.common.robot_devices.motors.dynamixel import DynamixelMotorsBus
-
-    # Test instantiating a common motors structure.
-    # Here the one from Alexander Koch follower arm.
-    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"),
-    }
-    motors_bus = DynamixelMotorsBus(port, motors)
+    motors_bus = make_motors_bus()
 
     # Test reading and writting before connecting raises an error
     with pytest.raises(RobotDeviceNotConnectedError):
@@ -41,7 +62,7 @@ def test_motors_bus(request):
     del motors_bus
 
     # Test connecting
-    motors_bus = DynamixelMotorsBus(port, motors)
+    motors_bus = make_motors_bus()
     motors_bus.connect()
 
     # Test connecting twice raises an error
@@ -52,7 +73,7 @@ def test_motors_bus(request):
     motors_bus.write("Torque_Enable", 0)
     values = motors_bus.read("Torque_Enable")
     assert isinstance(values, np.ndarray)
-    assert len(values) == len(motors)
+    assert len(values) == len(motors_bus.motors)
     assert (values == 0).all()
 
     # Test writing torque on a specific motor
@@ -83,10 +104,3 @@ def test_motors_bus(request):
     time.sleep(1)
     new_values = motors_bus.read("Present_Position")
     assert (new_values == values).all()
-
-
-@require_koch
-def test_find_port(request):
-    from lerobot.common.robot_devices.motors.dynamixel import find_port
-
-    find_port()

tests/utils.py

@@ -23,6 +23,7 @@ from lerobot.common.utils.import_utils import is_package_available
 
 # Pass this as the first argument to init_hydra_config.
 DEFAULT_CONFIG_PATH = "lerobot/configs/default.yaml"
+KOCH_ROBOT_CONFIG_PATH = "lerobot/configs/robot/koch.yaml"
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -161,6 +162,7 @@ def require_koch(func):
         if request is None:
             raise ValueError("The 'request' fixture must be passed to the test function as a parameter.")
 
+        # The function `is_koch_available` is defined in `tests/conftest.py`
         if not request.getfixturevalue("is_koch_available"):
             pytest.skip("An alexander koch robot is not available.")
         return func(*args, **kwargs)