format

lerobot/common/robot_devices/cameras/opencv.py

@@ -7,16 +7,16 @@ from pathlib import Path
 from threading import Thread
 
 import cv2
-
-from lerobot.common.robot_devices.utils import RobotDeviceAlreadyConnectedError, RobotDeviceNotConnectedError
-# Use 1 thread to avoid blocking the main thread. Especially useful during data collection
-# when other threads are used to save the images.
-cv2.setNumThreads(1)
 import numpy as np
 
 from lerobot.common.robot_devices.cameras.utils import save_color_image
+from lerobot.common.robot_devices.utils import RobotDeviceAlreadyConnectedError, RobotDeviceNotConnectedError
 from lerobot.common.utils.utils import capture_timestamp_utc
 
+# Use 1 thread to avoid blocking the main thread. Especially useful during data collection
+# when other threads are used to save the images.
+cv2.setNumThreads(1)
+
 
 def find_camera_indices(raise_when_empty=False, max_index_search_range=60):
     camera_ids = []
@@ -124,7 +124,9 @@ class OpenCVCamera:
         self.color = config.color
 
         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.")
+            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
@@ -193,13 +195,15 @@ class OpenCVCamera:
 
     def read(self, temporary_color: str | None = None) -> np.ndarray:
         """Read a frame from the camera returned in the format (height, width, channels)
-           (e.g. (640, 480, 3)), contrarily to the pytorch format which is channel first.
+        (e.g. (640, 480, 3)), contrarily to the pytorch format which is channel first.
 
-           Note: Reading a frame is done every `camera.fps` times per second, and it is blocking.
+        Note: Reading a frame is done every `camera.fps` times per second, and it is blocking.
-           If you are reading data from other sensors, we advise to use `camera.async_read()` which is non blocking version of `camera.read()`.
+        If you are reading data from other sensors, we advise to use `camera.async_read()` which is non blocking version of `camera.read()`.
         """
         if not self.is_connected:
-            raise RobotDeviceNotConnectedError(f"OpenCVCamera({self.camera_index}) is not connected. Try running `camera.connect()` first.")
+            raise RobotDeviceNotConnectedError(
+                f"OpenCVCamera({self.camera_index}) is not connected. Try running `camera.connect()` first."
+            )
 
         start_time = time.perf_counter()
 
@@ -220,7 +224,9 @@ class OpenCVCamera:
 
         h, w, _ = color_image.shape
         if h != self.height or w != self.width:
-            raise OSError(f"Can't capture color image with expected height and width ({self.height} x {self.width}). ({h} x {w}) returned instead.")
+            raise OSError(
+                f"Can't capture color image with expected height and width ({self.height} x {self.width}). ({h} x {w}) returned instead."
+            )
 
         # log the number of seconds it took to read the image
         self.logs["delta_timestamp_s"] = time.perf_counter() - start_time
@@ -236,7 +242,9 @@ class OpenCVCamera:
 
     def async_read(self):
         if not self.is_connected:
-            raise RobotDeviceNotConnectedError(f"OpenCVCamera({self.camera_index}) is not connected. Try running `camera.connect()` first.")
+            raise RobotDeviceNotConnectedError(
+                f"OpenCVCamera({self.camera_index}) is not connected. Try running `camera.connect()` first."
+            )
 
         if self.thread is None:
             self.stop_event = threading.Event()
@@ -247,16 +255,19 @@ class OpenCVCamera:
         num_tries = 0
         while self.color_image is None:
             num_tries += 1
-            time.sleep(1/self.fps)
+            time.sleep(1 / self.fps)
-            if num_tries > self.fps:
+            if num_tries > self.fps and (self.thread.ident is None or not self.thread.is_alive()):
-                if self.thread.ident is None or not self.thread.is_alive():
+                raise Exception(
-                    raise Exception("The thread responsible for `self.async_read()` took too much time to start. There might be an issue. Verify that `self.thread.start()` has been called.")
+                    "The thread responsible for `self.async_read()` took too much time to start. There might be an issue. Verify that `self.thread.start()` has been called."
+                )
 
         return self.color_image
 
     def disconnect(self):
         if not self.is_connected:
-            raise RobotDeviceNotConnectedError(f"OpenCVCamera({self.camera_index}) is not connected. Try running `camera.connect()` first.")
+            raise RobotDeviceNotConnectedError(
+                f"OpenCVCamera({self.camera_index}) is not connected. Try running `camera.connect()` first."
+            )
 
         if self.thread is not None and self.thread.is_alive():
             # wait for the thread to finish

lerobot/common/robot_devices/motors/dynamixel.py

@@ -1,8 +1,6 @@
 import enum
-from copy import deepcopy
-from queue import Queue
-from threading import Thread
 import time
+from copy import deepcopy
 
 import numpy as np
 from dynamixel_sdk import (
@@ -100,6 +98,8 @@ MODEL_CONTROL_TABLE = {
     "xm540-w270": X_SERIES_CONTROL_TABLE,
 }
 
+NUM_READ_RETRY = 10
+
 
 # TODO(rcadene): find better namming for these functions
 def uint32_to_int32(values: np.ndarray):
@@ -121,6 +121,7 @@ def int32_to_uint32(values: np.ndarray):
             values[i] = values[i] + 4294967296
     return values
 
+
 # def motor_position_to_angle(position: np.ndarray) -> np.ndarray:
 #     """
 #     Convert from motor position in [-2048, 2048] to radian in [-pi, pi]
@@ -139,21 +140,25 @@ def get_group_sync_key(data_name, motor_names):
     group_key = f"{data_name}_" + "_".join(motor_names)
     return group_key
 
+
 def get_result_name(fn_name, data_name, motor_names):
     group_key = get_group_sync_key(data_name, motor_names)
     rslt_name = f"{fn_name}_{group_key}"
     return rslt_name
 
+
 def get_queue_name(fn_name, data_name, motor_names):
     group_key = get_group_sync_key(data_name, motor_names)
     queue_name = f"{fn_name}_{group_key}"
     return queue_name
 
+
 def get_log_name(var_name, fn_name, data_name, motor_names):
     group_key = get_group_sync_key(data_name, motor_names)
     log_name = f"{var_name}_{fn_name}_{group_key}"
     return log_name
 
+
 class TorqueMode(enum.Enum):
     ENABLED = 1
     DISABLED = 0
@@ -192,6 +197,7 @@ class DynamixelMotorsBus:
     motors_bus.disconnect()
     ```
     """
+
     def __init__(
         self,
         port: str,
@@ -215,7 +221,9 @@ class DynamixelMotorsBus:
 
     def connect(self):
         if self.is_connected:
-            raise RobotDeviceAlreadyConnectedError(f"DynamixelMotorsBus({self.port}) is already connected. Do not call `motors_bus.connect()` twice.")
+            raise RobotDeviceAlreadyConnectedError(
+                f"DynamixelMotorsBus({self.port}) is already connected. Do not call `motors_bus.connect()` twice."
+            )
 
         self.port_handler = PortHandler(self.port)
         self.packet_handler = PacketHandler(PROTOCOL_VERSION)
@@ -271,7 +279,9 @@ class DynamixelMotorsBus:
 
     def read(self, data_name, motor_names: str | list[str] | None = None):
         if not self.is_connected:
-            raise RobotDeviceNotConnectedError(f"DynamixelMotorsBus({self.port}) is not connected. You need to run `motors_bus.connect()`.")
+            raise RobotDeviceNotConnectedError(
+                f"DynamixelMotorsBus({self.port}) is not connected. You need to run `motors_bus.connect()`."
+            )
 
         start_time = time.perf_counter()
 
@@ -298,8 +308,7 @@ class DynamixelMotorsBus:
             for idx in motor_ids:
                 self.group_readers[group_key].addParam(idx)
 
-        NUM_TRIES = 10
+        for _ in range(NUM_READ_RETRY):
-        for _ in range(NUM_TRIES):
             comm = self.group_readers[group_key].txRxPacket()
             if comm == COMM_SUCCESS:
                 break
@@ -324,9 +333,6 @@ class DynamixelMotorsBus:
         if data_name in CALIBRATION_REQUIRED:
             values = self.apply_calibration(values, motor_names)
 
-        if data_name in CONVERT_POSITION_TO_ANGLE_REQUIRED:
-            values = motor_position_to_angle(values)
-
         # log the number of seconds it took to read the data from the motors
         delta_ts_name = get_log_name("delta_timestamp_s", "read", data_name, motor_names)
         self.logs[delta_ts_name] = time.perf_counter() - start_time
@@ -339,7 +345,9 @@ class DynamixelMotorsBus:
 
     def write(self, data_name, values: int | float | np.ndarray, motor_names: str | list[str] | None = None):
         if not self.is_connected:
-            raise RobotDeviceNotConnectedError(f"DynamixelMotorsBus({self.port}) is not connected. You need to run `motors_bus.connect()`.")
+            raise RobotDeviceNotConnectedError(
+                f"DynamixelMotorsBus({self.port}) is not connected. You need to run `motors_bus.connect()`."
+            )
 
         start_time = time.perf_counter()
 
@@ -361,9 +369,6 @@ class DynamixelMotorsBus:
             motor_ids.append(motor_idx)
             models.append(model)
 
-        if data_name in CONVERT_POSITION_TO_ANGLE_REQUIRED:
-            values = motor_angle_to_position(values)
-
         if data_name in CALIBRATION_REQUIRED:
             values = self.revert_calibration(values, motor_names)
 
@@ -429,7 +434,9 @@ class DynamixelMotorsBus:
 
     def disconnect(self):
         if not self.is_connected:
-            raise RobotDeviceNotConnectedError(f"DynamixelMotorsBus({self.port}) is not connected. Try running `motors_bus.connect()` first.")
+            raise RobotDeviceNotConnectedError(
+                f"DynamixelMotorsBus({self.port}) is not connected. Try running `motors_bus.connect()` first."
+            )
 
         if self.port_handler is not None:
             self.port_handler.closePort()
@@ -444,7 +451,6 @@ class DynamixelMotorsBus:
         if getattr(self, "is_connected", False):
             self.disconnect()
 
-
     # def read(self, data_name, motor_name: str):
     #     motor_idx, model = self.motors[motor_name]
     #     addr, bytes = self.model_ctrl_table[model][data_name]

lerobot/common/robot_devices/robots/koch.py

@@ -1,9 +1,8 @@
 import pickle
+import time
 from dataclasses import dataclass, field, replace
 from pathlib import Path
-import time
 
-import einops
 import numpy as np
 import torch
 
@@ -145,15 +144,17 @@ def reset_arm(arm: MotorsBus):
 
 
 def run_arm_calibration(arm: MotorsBus, name: str, arm_type: str):
-    """ Example of usage:
+    """Example of usage:
-        ```python
+    ```python
-        run_arm_calibration(arm, "left", "follower")
+    run_arm_calibration(arm, "left", "follower")
-        ```
+    ```
     """
     reset_arm(arm)
 
     # 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(
+        f"Please move the '{name} {arm_type}' arm to the horizontal position (gripper fully closed, see {URL_HORIZONTAL_POSITION[arm_type]})"
+    )
     input("Press Enter to continue...")
 
     horizontal_homing_offset = compute_homing_offset(
@@ -161,7 +162,9 @@ def run_arm_calibration(arm: MotorsBus, name: str, arm_type: str):
     )
 
     # 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(
+        f"Please move the '{name} {arm_type}' arm to the 90 degree position (gripper fully open, see {URL_90_DEGREE_POSITION[arm_type]})"
+    )
     input("Press Enter to continue...")
 
     drive_mode = compute_drive_mode(arm, horizontal_homing_offset)
@@ -197,14 +200,11 @@ class KochRobotConfig:
     """
 
     # Define all components of the robot
-    leader_arms: dict[str, MotorsBus] = field(
+    leader_arms: dict[str, MotorsBus] = field(default_factory=lambda: {})
-        default_factory=lambda: {}
+    follower_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: {})
 
+
 class KochRobot:
     """Tau Robotics: https://tau-robotics.com
 
@@ -318,10 +318,14 @@ class KochRobot:
 
     def connect(self):
         if self.is_connected:
-            raise RobotDeviceAlreadyConnectedError(f"KochRobot is already connected. Do not run `robot.connect()` twice.")
+            raise RobotDeviceAlreadyConnectedError(
+                "KochRobot is already connected. Do not run `robot.connect()` twice."
+            )
 
         if not self.leader_arms and not self.follower_arms and not self.cameras:
-            raise ValueError("KochRobot doesn't have any device to connect. See example of usage in docstring of the class.")
+            raise ValueError(
+                "KochRobot doesn't have any device to connect. See example of usage in docstring of the class."
+            )
 
         # Connect the arms
         for name in self.follower_arms:
@@ -391,7 +395,9 @@ class KochRobot:
         self, record_data=False
     ) -> None | tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
         if not self.is_connected:
-            raise RobotDeviceNotConnectedError(f"KochRobot is not connected. You need to run `robot.connect()`.")
+            raise RobotDeviceNotConnectedError(
+                "KochRobot is not connected. You need to run `robot.connect()`."
+            )
 
         # Prepare to assign the positions of the leader to the follower
         leader_pos = {}
@@ -455,7 +461,9 @@ class KochRobot:
     def capture_observation(self):
         """The returned observations do not have a batch dimension."""
         if not self.is_connected:
-            raise RobotDeviceNotConnectedError(f"KochRobot is not connected. You need to run `robot.connect()`.")
+            raise RobotDeviceNotConnectedError(
+                "KochRobot is not connected. You need to run `robot.connect()`."
+            )
 
         # Read follower position
         follower_pos = {}
@@ -488,7 +496,9 @@ class KochRobot:
     def send_action(self, action: torch.Tensor):
         """The provided action is expected to be a vector."""
         if not self.is_connected:
-            raise RobotDeviceNotConnectedError(f"KochRobot is not connected. You need to run `robot.connect()`.")
+            raise RobotDeviceNotConnectedError(
+                "KochRobot is not connected. You need to run `robot.connect()`."
+            )
 
         from_idx = 0
         to_idx = 0
@@ -504,7 +514,9 @@ class KochRobot:
 
     def disconnect(self):
         if not self.is_connected:
-            raise RobotDeviceNotConnectedError(f"KochRobot is not connected. You need to run `robot.connect()` before disconnecting.")
+            raise RobotDeviceNotConnectedError(
+                "KochRobot is not connected. You need to run `robot.connect()` before disconnecting."
+            )
 
         for name in self.follower_arms:
             self.follower_arms[name].disconnect()

lerobot/common/robot_devices/utils.py

@@ -1,12 +1,19 @@
 class RobotDeviceNotConnectedError(Exception):
     """Exception raised when the robot device is not connected."""
-    def __init__(self, message="This robot device is not connected. Try calling `robot_device.connect()` first."):
+
+    def __init__(
+        self, message="This robot device is not connected. Try calling `robot_device.connect()` first."
+    ):
         self.message = message
         super().__init__(self.message)
 
 
 class RobotDeviceAlreadyConnectedError(Exception):
     """Exception raised when the robot device is already connected."""
-    def __init__(self, message="This robot device is already connected. Try not calling `robot_device.connect()` twice."):
+
+    def __init__(
+        self,
+        message="This robot device is already connected. Try not calling `robot_device.connect()` twice.",
+    ):
         self.message = message
         super().__init__(self.message)

lerobot/scripts/control_robot.py

@@ -102,6 +102,7 @@ def save_image(img_tensor, key, frame_index, episode_index, videos_dir):
     path.parent.mkdir(parents=True, exist_ok=True)
     img.save(str(path), quality=100)
 
+
 def busy_wait(seconds):
     # Significantly more accurate than `time.sleep`, and mendatory for our use case,
     # but it consumes CPU cycles.
@@ -110,11 +111,13 @@ def busy_wait(seconds):
     while time.perf_counter() < end_time:
         pass
 
+
 def none_or_int(value):
     if value == "None":
         return None
     return int(value)
 
+
 def log_control_info(robot, dt_s, episode_index=None, frame_index=None):
     log_items = []
     if episode_index is not None:
@@ -130,16 +133,16 @@ def log_control_info(robot, dt_s, episode_index=None, frame_index=None):
     log_dt("dt", dt_s)
 
     for name in robot.leader_arms:
-        key = f'read_leader_{name}_pos_dt_s'
+        key = f"read_leader_{name}_pos_dt_s"
         if key in robot.logs:
             log_dt("dtRlead", robot.logs[key])
 
     for name in robot.follower_arms:
-        key = f'write_follower_{name}_goal_pos_dt_s'
+        key = f"write_follower_{name}_goal_pos_dt_s"
         if key in robot.logs:
             log_dt("dtRfoll", robot.logs[key])
 
-        key = f'read_follower_{name}_pos_dt_s'
+        key = f"read_follower_{name}_pos_dt_s"
         if key in robot.logs:
             log_dt("dtWfoll", robot.logs[key])
 
@@ -154,6 +157,7 @@ def log_control_info(robot, dt_s, episode_index=None, frame_index=None):
 
     logging.info(" ".join(log_items))
 
+
 ########################################################################################
 # Control modes
 ########################################################################################
@@ -256,7 +260,9 @@ def record_dataset(
                 not_image_keys = [key for key in observation if "image" not in key]
 
                 for key in image_keys:
-                    future = executor.submit(save_image, observation[key], key, frame_index, episode_index, videos_dir)
+                    future = executor.submit(
+                        save_image, observation[key], key, frame_index, episode_index, videos_dir
+                    )
                     futures.append(future)
 
                 for key in not_image_keys:

tests/test_cameras.py

@@ -1,18 +1,15 @@
-
-
-from pathlib import Path
 import numpy as np
 import pytest
 
 from lerobot.common.robot_devices.cameras.opencv import OpenCVCamera
-from lerobot.common.robot_devices.utils import RobotDeviceNotConnectedError, RobotDeviceAlreadyConnectedError
+from lerobot.common.robot_devices.utils import RobotDeviceAlreadyConnectedError, RobotDeviceNotConnectedError
-
 
 CAMERA_INDEX = 2
 # Maximum absolute difference between two consecutive images recored by a camera.
 # This value differs with respect to the camera.
 MAX_PIXEL_DIFFERENCE = 25
 
+
 def compute_max_pixel_difference(first_image, second_image):
     return np.abs(first_image.astype(float) - second_image.astype(float)).max()
 
@@ -68,7 +65,10 @@ def test_camera():
         camera.read()
     color_image = camera.read()
     async_color_image = camera.async_read()
-    print("max_pixel_difference between read() and async_read()", compute_max_pixel_difference(color_image, async_color_image))
+    print(
+        "max_pixel_difference between read() and async_read()",
+        compute_max_pixel_difference(color_image, async_color_image),
+    )
     assert np.allclose(color_image, async_color_image, rtol=1e-5, atol=MAX_PIXEL_DIFFERENCE)
 
     # Test disconnecting
@@ -86,7 +86,7 @@ def test_camera():
     camera.connect()
     assert camera.color == "bgr"
     bgr_color_image = camera.read()
-    assert np.allclose(color_image, bgr_color_image[:, :, [2,1,0]], rtol=1e-5, atol=MAX_PIXEL_DIFFERENCE)
+    assert np.allclose(color_image, bgr_color_image[:, :, [2, 1, 0]], rtol=1e-5, atol=MAX_PIXEL_DIFFERENCE)
     del camera
 
     # TODO(rcadene): Add a test for a camera that doesnt support fps=60 and raises an OSError
@@ -116,4 +116,3 @@ def test_camera():
     with pytest.raises(OSError):
         camera.connect()
     del camera
-

tests/test_control_robot.py

@@ -1,6 +1,5 @@
-
-
 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
@@ -21,12 +20,14 @@ def test_record_dataset_and_replay_episode_and_run_policy(tmpdir):
     env_name = "koch_real"
     policy_name = "act_koch_real"
 
-    #root = Path(tmpdir)
+    # root = Path(tmpdir)
     root = Path("tmp/data")
     repo_id = "lerobot/debug"
 
     robot = make_robot(robot_name)
-    dataset = record_dataset(robot, fps=30, root=root, repo_id=repo_id, warmup_time_s=1, episode_time_s=1, num_episodes=2)
+    dataset = record_dataset(
+        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)
 
@@ -36,7 +37,7 @@ def test_record_dataset_and_replay_episode_and_run_policy(tmpdir):
             f"env={env_name}",
             f"policy={policy_name}",
             f"device={DEVICE}",
-        ]
+        ],
     )
 
     policy = make_policy(hydra_cfg=cfg, dataset_stats=dataset.stats)
@@ -44,4 +45,3 @@ def test_record_dataset_and_replay_episode_and_run_policy(tmpdir):
     run_policy(robot, policy, cfg, run_time_s=1)
 
     del robot
-

tests/test_motors.py

@@ -1,7 +1,8 @@
-
 import time
+
 import numpy as np
 import pytest
+
 from lerobot.common.robot_devices.motors.dynamixel import DynamixelMotorsBus
 from lerobot.common.robot_devices.utils import RobotDeviceAlreadyConnectedError, RobotDeviceNotConnectedError
 
@@ -80,4 +81,3 @@ def test_motors_bus():
     time.sleep(1)
     new_values = motors_bus.read("Present_Position")
     assert (new_values == values).all()
-

tests/test_robots.py

@@ -1,7 +1,9 @@
-from pathlib import Path
 import pickle
+from pathlib import Path
+
 import pytest
 import torch
+
 from lerobot.common.robot_devices.robots.factory import make_robot
 from lerobot.common.robot_devices.robots.koch import KochRobot
 from lerobot.common.robot_devices.utils import RobotDeviceAlreadyConnectedError, RobotDeviceNotConnectedError
@@ -13,7 +15,24 @@ def test_robot(tmpdir):
     # TODO(rcadene): add compatibility with other robots
 
     # Save calibration preset
-    calibration = {'follower_main': {'shoulder_pan': (-2048, False), 'shoulder_lift': (2048, True), 'elbow_flex': (-1024, False), 'wrist_flex': (2048, True), 'wrist_roll': (2048, True), 'gripper': (2048, True)}, 'leader_main': {'shoulder_pan': (-2048, False), 'shoulder_lift': (1024, True), 'elbow_flex': (2048, True), 'wrist_flex': (-2048, False), 'wrist_roll': (2048, True), 'gripper': (2048, True)}}
+    calibration = {
+        "follower_main": {
+            "shoulder_pan": (-2048, False),
+            "shoulder_lift": (2048, True),
+            "elbow_flex": (-1024, False),
+            "wrist_flex": (2048, True),
+            "wrist_roll": (2048, True),
+            "gripper": (2048, True),
+        },
+        "leader_main": {
+            "shoulder_pan": (-2048, False),
+            "shoulder_lift": (1024, True),
+            "elbow_flex": (2048, True),
+            "wrist_flex": (-2048, False),
+            "wrist_roll": (2048, True),
+            "gripper": (2048, True),
+        },
+    }
     tmpdir = Path(tmpdir)
     calibration_path = tmpdir / "calibration.pkl"
     calibration_path.parent.mkdir(parents=True, exist_ok=True)
@@ -105,4 +124,3 @@ def test_robot(tmpdir):
     for name in robot.cameras:
         assert not robot.cameras[name].is_connected
     del robot
-