diff --git a/lerobot/calibrate.py b/lerobot/calibrate.py index 088aee299..6780577ff 100644 --- a/lerobot/calibrate.py +++ b/lerobot/calibrate.py @@ -37,18 +37,21 @@ from lerobot.common.robots import ( # noqa: F401 Robot, RobotConfig, koch_follower, + lekiwi, make_robot_from_config, so100_follower, + so101_follower, ) from lerobot.common.teleoperators import ( # noqa: F401 Teleoperator, TeleoperatorConfig, + koch_leader, make_teleoperator_from_config, + so100_leader, + so101_leader, ) from lerobot.common.utils.utils import init_logging -from .common.teleoperators import koch_leader, so100_leader # noqa: F401 - @dataclass class CalibrateConfig: diff --git a/lerobot/common/datasets/utils.py b/lerobot/common/datasets/utils.py index a16fe8acd..542daf2aa 100644 --- a/lerobot/common/datasets/utils.py +++ b/lerobot/common/datasets/utils.py @@ -400,15 +400,22 @@ def hw_to_dataset_features( joint_fts = {key: ftype for key, ftype in hw_features.items() if ftype is float} cam_fts = {key: shape for key, shape in hw_features.items() if isinstance(shape, tuple)} - if joint_fts: - features[f"{prefix}.joints"] = { + if joint_fts and prefix == "action": + features[prefix] = { + "dtype": "float32", + "shape": (len(joint_fts),), + "names": list(joint_fts), + } + + if joint_fts and prefix == "observation": + features[f"{prefix}.state"] = { "dtype": "float32", "shape": (len(joint_fts),), "names": list(joint_fts), } for key, shape in cam_fts.items(): - features[f"{prefix}.cameras.{key}"] = { + features[f"{prefix}.images.{key}"] = { "dtype": "video" if use_video else "image", "shape": shape, "names": ["height", "width", "channels"], @@ -428,7 +435,7 @@ def build_dataset_frame( elif ft["dtype"] == "float32" and len(ft["shape"]) == 1: frame[key] = np.array([values[name] for name in ft["names"]], dtype=np.float32) elif ft["dtype"] in ["image", "video"]: - frame[key] = values[key.removeprefix(f"{prefix}.cameras.")] + frame[key] = values[key.removeprefix(f"{prefix}.images.")] return frame diff --git a/lerobot/common/robots/lekiwi/config_lekiwi.py b/lerobot/common/robots/lekiwi/config_lekiwi.py index 946718a15..97c8060d5 100644 --- a/lerobot/common/robots/lekiwi/config_lekiwi.py +++ b/lerobot/common/robots/lekiwi/config_lekiwi.py @@ -14,7 +14,7 @@ from dataclasses import dataclass, field -from lerobot.common.cameras.configs import CameraConfig +from lerobot.common.cameras.configs import CameraConfig, Cv2Rotation from lerobot.common.cameras.opencv.configuration_opencv import OpenCVCameraConfig from ..config import RobotConfig @@ -34,11 +34,9 @@ class LeKiwiConfig(RobotConfig): cameras: dict[str, CameraConfig] = field( default_factory=lambda: { - "front": OpenCVCameraConfig( - camera_index="/dev/video0", fps=30, width=640, height=480, rotation=None - ), + "front": OpenCVCameraConfig(index_or_path="/dev/video0", fps=30, width=640, height=480), "wrist": OpenCVCameraConfig( - camera_index="/dev/video2", fps=30, width=640, height=480, rotation=180 + index_or_path="/dev/video2", fps=30, width=640, height=480, rotation=Cv2Rotation.ROTATE_180 ), } ) diff --git a/lerobot/common/teleoperators/koch_leader/koch_leader.py b/lerobot/common/teleoperators/koch_leader/koch_leader.py index 6821e0200..820acc87c 100644 --- a/lerobot/common/teleoperators/koch_leader/koch_leader.py +++ b/lerobot/common/teleoperators/koch_leader/koch_leader.py @@ -140,7 +140,8 @@ class KochLeader(Teleoperator): self.bus.write("Operating_Mode", "gripper", OperatingMode.CURRENT_POSITION.value) # Set gripper's goal pos in current position mode so that we can use it as a trigger. self.bus.enable_torque("gripper") - self.bus.write("Goal_Position", "gripper", self.config.gripper_open_pos) + if self.is_calibrated: + self.bus.write("Goal_Position", "gripper", self.config.gripper_open_pos) def setup_motors(self) -> None: for motor in reversed(self.bus.motors): diff --git a/lerobot/record.py b/lerobot/record.py index 6b1fdeb78..348d43567 100644 --- a/lerobot/record.py +++ b/lerobot/record.py @@ -57,6 +57,7 @@ from lerobot.common.robots import ( # noqa: F401 koch_follower, make_robot_from_config, so100_follower, + so101_follower, ) from lerobot.common.teleoperators import ( # noqa: F401 Teleoperator, @@ -80,7 +81,7 @@ from lerobot.common.utils.visualization_utils import _init_rerun from lerobot.configs import parser from lerobot.configs.policies import PreTrainedConfig -from .common.teleoperators import koch_leader, so100_leader # noqa: F401 +from .common.teleoperators import koch_leader, so100_leader, so101_leader # noqa: F401 @dataclass diff --git a/lerobot/replay.py b/lerobot/replay.py index 5b8f86190..36eb0864d 100644 --- a/lerobot/replay.py +++ b/lerobot/replay.py @@ -42,6 +42,7 @@ from lerobot.common.robots import ( # noqa: F401 koch_follower, make_robot_from_config, so100_follower, + so101_follower, ) from lerobot.common.utils.robot_utils import busy_wait from lerobot.common.utils.utils import ( @@ -77,16 +78,16 @@ def replay(cfg: ReplayConfig): robot = make_robot_from_config(cfg.robot) dataset = LeRobotDataset(cfg.dataset.repo_id, root=cfg.dataset.root, episodes=[cfg.dataset.episode]) - actions = dataset.hf_dataset.select_columns("action.joints") + actions = dataset.hf_dataset.select_columns("action") robot.connect() log_say("Replaying episode", cfg.play_sounds, blocking=True) for idx in range(dataset.num_frames): start_episode_t = time.perf_counter() - action_array = actions[idx]["action.joints"] + action_array = actions[idx]["action"] action = {} - for i, name in enumerate(dataset.features["action.joints"]["names"]): + for i, name in enumerate(dataset.features["action"]["names"]): action[name] = action_array[i] robot.send_action(action) diff --git a/lerobot/setup_motors.py b/lerobot/setup_motors.py index a42b9cac1..7909dc68d 100644 --- a/lerobot/setup_motors.py +++ b/lerobot/setup_motors.py @@ -28,12 +28,20 @@ from dataclasses import dataclass import draccus -from .common.robots import RobotConfig, koch_follower, make_robot_from_config, so100_follower # noqa: F401 +from .common.robots import ( # noqa: F401 + RobotConfig, + koch_follower, + lekiwi, + make_robot_from_config, + so100_follower, + so101_follower, +) from .common.teleoperators import ( # noqa: F401 TeleoperatorConfig, koch_leader, make_teleoperator_from_config, so100_leader, + so101_leader, ) COMPATIBLE_DEVICES = [ @@ -41,6 +49,9 @@ COMPATIBLE_DEVICES = [ "koch_leader", "so100_follower", "so100_leader", + "so101_follower", + "so101_leader", + "lekiwi", ] diff --git a/lerobot/teleoperate.py b/lerobot/teleoperate.py index f3496a69c..84acfdc5b 100644 --- a/lerobot/teleoperate.py +++ b/lerobot/teleoperate.py @@ -19,13 +19,14 @@ Example: ```shell python -m lerobot.teleoperate \ - --robot.type=so100_follower \ + --robot.type=so101_follower \ --robot.port=/dev/tty.usbmodem58760431541 \ - --robot.cameras="{laptop: {type: opencv, camera_index: 0}}" \ + --robot.cameras="{ front: {type: opencv, index_or_path: 0, width: 1920, height: 1080, fps: 30}}" \ --robot.id=black \ - --teleop.type=so100_leader \ + --teleop.type=so101_leader \ --teleop.port=/dev/tty.usbmodem58760431551 \ - --teleop.id=blue + --teleop.id=blue \ + --display_data=true ``` """ @@ -46,6 +47,7 @@ from lerobot.common.robots import ( # noqa: F401 koch_follower, make_robot_from_config, so100_follower, + so101_follower, ) from lerobot.common.teleoperators import ( Teleoperator, @@ -55,7 +57,7 @@ from lerobot.common.teleoperators import ( from lerobot.common.utils.utils import init_logging, move_cursor_up from lerobot.common.utils.visualization_utils import _init_rerun -from .common.teleoperators import koch_leader, so100_leader # noqa: F401 +from .common.teleoperators import koch_leader, so100_leader, so101_leader # noqa: F401 @dataclass