wangzixiang
/
lerobot
mirror of https://github.com/huggingface/lerobot.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

format

This commit is contained in:
Remi Cadene 2024-07-10 14:11:53 +02:00
parent 68a561570c
commit dd7bce7563
11 changed files with 148 additions and 89 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

37
lerobot/common/robot_devices/cameras/opencv.py
View File

@ -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
@ -199,7 +201,9 @@ class OpenCVCamera:
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()
@ -248,15 +256,18 @@ class OpenCVCamera:
while self.color_image is None:
num_tries += 1
time.sleep(1 / self.fps)
if num_tries > self.fps:
if self.thread.ident is None or not self.thread.is_alive():
raise Exception("The thread responsible for `self.async_read()` took too much time to start. There might be an issue. Verify that `self.thread.start()` has been called.")
if num_tries > self.fps and (self.thread.ident is None or not self.thread.is_alive()):
raise Exception(
"The thread responsible for `self.async_read()` took too much time to start. There might be an issue. Verify that `self.thread.start()` has been called."
)
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

38
lerobot/common/robot_devices/motors/dynamixel.py
View File

@ -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_TRIES):
for _ in range(NUM_READ_RETRY):
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]

44
lerobot/common/robot_devices/robots/koch.py
View File

@ -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
@ -153,7 +152,9 @@ def run_arm_calibration(arm: MotorsBus, name: str, arm_type: str):
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(
default_factory=lambda: {}
)
follower_arms: dict[str, MotorsBus] = field(
default_factory=lambda: {}
)
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: {})
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()

11
lerobot/common/robot_devices/utils.py
View File

@ -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)

14
lerobot/scripts/control_robot.py
View File

@ -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:

13
tests/test_cameras.py
View File

@ -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
@ -116,4 +116,3 @@ def test_camera():
with pytest.raises(OSError):
camera.connect()
del camera

10
tests/test_control_robot.py
View File

@ -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
@ -26,7 +25,9 @@ def test_record_dataset_and_replay_episode_and_run_policy(tmpdir):
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

4
tests/test_motors.py
View File

@ -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()

24
tests/test_robots.py
View File

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