add aloha setup

examples/aloha_record.py

@@ -1,85 +1,118 @@
-import time
-
 from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig
-from lerobot.robots.viperx import ViperX, ViperXConfig
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.teleoperators.widowx import WidowX, WidowXConfig
+from lerobot.datasets.utils import hw_to_dataset_features
+from lerobot.record import record_loop
+from lerobot.robots.aloha import Aloha, AlohaConfig
+from lerobot.teleoperators.aloha_teleop import AlohaTeleop, AlohaTeleopConfig
+from lerobot.utils.control_utils import init_keyboard_listener
+from lerobot.utils.utils import log_say
+from lerobot.utils.visualization_utils import _init_rerun
 
-# TODO: pepijn create a Aloha robot with has the two robots integrated with the 3 camera's (teleoprators can still be separate for now)
+# Recording configuration
-# TODO: pepijn add record loop and dataset etc...
+NUM_EPISODES = 5
+FPS = 30
+EPISODE_TIME_SEC = 60
+RESET_TIME_SEC = 10
+TASK_DESCRIPTION = "ALOHA bimanual manipulation task"
+REPO_ID = "pepijn223/aloha-record-test"
 
+# Create camera configuration
 camera_config = {
-    "front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=30),
+    "front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS),
-    "wrist_right": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=30),
+    "wrist_right": OpenCVCameraConfig(index_or_path=1, width=640, height=480, fps=FPS),
-    "wrist_left": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=30),
+    "wrist_left": OpenCVCameraConfig(index_or_path=2, width=640, height=480, fps=FPS),
 }
 
-config_follower_right = ViperXConfig(
+# ALOHA Robot Configuration (dual ViperX followers)
-    port="/dev/tty.usbserial-FT891KBG",
+aloha_robot_config = AlohaConfig(
-    id="viperx_right",
+    id="aloha",
-    max_relative_target=20.0,  # increased from default 5.0
+    left_arm_port="/dev/tty.usbserial-FT89FM09",
-    use_degrees=True,
+    right_arm_port="/dev/tty.usbserial-FT891KBG",
+    left_arm_max_relative_target=20.0,
+    right_arm_max_relative_target=20.0,
+    left_arm_use_degrees=True,
+    right_arm_use_degrees=True,
     cameras=camera_config,
 )
 
-config_leader_right = WidowXConfig(
+# ALOHA Teleoperator Configuration (dual WidowX leaders)
-    port="/dev/tty.usbserial-FT89FM77",
+aloha_teleop_config = AlohaTeleopConfig(
-    id="widowx_right",
+    id="aloha_teleop",
-    gripper_motor="xc430-w150",
+    left_arm_port="/dev/tty.usbserial-FT891KPN",
-    use_degrees=True,
+    right_arm_port="/dev/tty.usbserial-FT89FM77",
+    left_arm_gripper_motor="xl430-w250",
+    right_arm_gripper_motor="xc430-w150",
+    left_arm_use_degrees=True,
+    right_arm_use_degrees=True,
 )
 
-config_follower_left = ViperXConfig(
+# Initialize the robot and teleoperator
-    port="/dev/tty.usbserial-FT89FM09",
+robot = Aloha(aloha_robot_config)
-    id="viperx_left",
+teleop = AlohaTeleop(aloha_teleop_config)
-    max_relative_target=20.0,  # increased from default 5.0
+
-    use_degrees=True,
+# Configure the dataset features
+action_features = hw_to_dataset_features(robot.action_features, "action")
+obs_features = hw_to_dataset_features(robot.observation_features, "observation")
+dataset_features = {**action_features, **obs_features}
+
+# Create the dataset
+dataset = LeRobotDataset.create(
+    repo_id=REPO_ID,
+    fps=FPS,
+    features=dataset_features,
+    robot_type=robot.name,
+    use_videos=True,
+    image_writer_threads=4,
 )
 
-config_leader_left = WidowXConfig(
+# Initialize the keyboard listener and rerun visualization
-    port="/dev/tty.usbserial-FT891KPN",
+_, events = init_keyboard_listener()
-    id="widowx_left",
+_init_rerun(session_name="aloha_recording")
-    gripper_motor="xl430-w250",
-    use_degrees=True,
-)
 
-follower_right = ViperX(config_follower_right)
+# Connect the robot and teleoperator
-follower_right.connect()
+robot.connect()
+teleop.connect()
 
-leader_right = WidowX(config_leader_right)
+episode_idx = 0
-leader_right.connect()
+while episode_idx < NUM_EPISODES and not events["stop_recording"]:
+    log_say(f"Recording episode {episode_idx + 1} of {NUM_EPISODES}")
 
-follower_left = ViperX(config_follower_left)
+    record_loop(
-follower_left.connect()
+        robot=robot,
+        events=events,
+        fps=FPS,
+        teleop=teleop,
+        dataset=dataset,
+        control_time_s=EPISODE_TIME_SEC,
+        single_task=TASK_DESCRIPTION,
+        display_data=True,
+    )
 
-leader_left = WidowX(config_leader_left)
+    # Reset the environment if not stopping or re-recording
-leader_left.connect()
+    if not events["stop_recording"] and (episode_idx < NUM_EPISODES - 1 or events["rerecord_episode"]):
+        log_say("Reset the environment")
+        record_loop(
+            robot=robot,
+            events=events,
+            fps=FPS,
+            teleop=teleop,
+            control_time_s=RESET_TIME_SEC,
+            single_task=TASK_DESCRIPTION,
+            display_data=True,
+        )
 
-while True:
+    if events["rerecord_episode"]:
-    act_right = leader_right.get_action()
+        log_say("Re-recording episode")
-    obs_right = follower_right.get_observation()
+        events["rerecord_episode"] = False
+        events["exit_early"] = False
+        dataset.clear_episode_buffer()
+        continue
 
-    act_left = leader_left.get_action()
+    dataset.save_episode()
-    obs_left = follower_left.get_observation()
+    episode_idx += 1
 
-    print("=" * 60)
+# Clean up
-    print("ACTION (Leader Right):")
+log_say("Stop recording")
-    for key, value in act_right.items():
+robot.disconnect()
-        print(f"  {key:20}: {value:8.3f}")
+teleop.disconnect()
-
+dataset.push_to_hub()
-    print("\nOBSERVATION (Follower Right):")
-    for key, value in obs_right.items():
-        print(f"  {key:20}: {value:8.3f}")
-    print("=" * 60)
-    print("ACTION (Leader Left):")
-    for key, value in act_left.items():
-        print(f"  {key:20}: {value:8.3f}")
-
-    print("\nOBSERVATION (Follower Left):")
-    for key, value in obs_left.items():
-        print(f"  {key:20}: {value:8.3f}")
-    print("=" * 60)
-
-    # follower_right.send_action(act_right)
-    # follower_left.send_action(act_left)
-
-    time.sleep(0.02)

src/lerobot/robots/aloha/__init__.py

@@ -0,0 +1,4 @@
+from .aloha import Aloha
+from .config_aloha import AlohaConfig
+
+__all__ = ["Aloha", "AlohaConfig"]

src/lerobot/robots/aloha/aloha.py

@@ -0,0 +1,161 @@
+#!/usr/bin/env python
+
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+import time
+from functools import cached_property
+from typing import Any
+
+from lerobot.cameras.utils import make_cameras_from_configs
+from lerobot.robots.viperx import ViperX
+from lerobot.robots.viperx.config_viperx import ViperXConfig
+
+from ..robot import Robot
+from .config_aloha import AlohaConfig
+
+logger = logging.getLogger(__name__)
+
+
+class Aloha(Robot):
+    """
+    ALOHA Bimanual Robot System using dual ViperX follower arms.
+    Based on the ALOHA (A Low-cost Open-source Hardware System for Bimanual Teleoperation) design.
+    """
+
+    config_class = AlohaConfig
+    name = "aloha"
+
+    def __init__(self, config: AlohaConfig):
+        super().__init__(config)
+        self.config = config
+
+        left_arm_config = ViperXConfig(
+            id=f"{config.id}_left" if config.id else None,
+            calibration_dir=config.calibration_dir,
+            port=config.left_arm_port,
+            max_relative_target=config.left_arm_max_relative_target,
+            use_degrees=config.left_arm_use_degrees,
+            cameras={},
+        )
+
+        right_arm_config = ViperXConfig(
+            id=f"{config.id}_right" if config.id else None,
+            calibration_dir=config.calibration_dir,
+            port=config.right_arm_port,
+            max_relative_target=config.right_arm_max_relative_target,
+            use_degrees=config.right_arm_use_degrees,
+            cameras={},
+        )
+
+        self.left_arm = ViperX(left_arm_config)
+        self.right_arm = ViperX(right_arm_config)
+        self.cameras = make_cameras_from_configs(config.cameras)
+
+    @property
+    def _motors_ft(self) -> dict[str, type]:
+        return {f"left_{motor}.pos": float for motor in self.left_arm.bus.motors} | {
+            f"right_{motor}.pos": float for motor in self.right_arm.bus.motors
+        }
+
+    @property
+    def _cameras_ft(self) -> dict[str, tuple]:
+        return {
+            cam: (self.config.cameras[cam].height, self.config.cameras[cam].width, 3) for cam in self.cameras
+        }
+
+    @cached_property
+    def observation_features(self) -> dict[str, type | tuple]:
+        return {**self._motors_ft, **self._cameras_ft}
+
+    @cached_property
+    def action_features(self) -> dict[str, type]:
+        return self._motors_ft
+
+    @property
+    def is_connected(self) -> bool:
+        return (
+            self.left_arm.bus.is_connected
+            and self.right_arm.bus.is_connected
+            and all(cam.is_connected for cam in self.cameras.values())
+        )
+
+    def connect(self, calibrate: bool = True) -> None:
+        self.left_arm.connect(calibrate)
+        self.right_arm.connect(calibrate)
+
+        for cam in self.cameras.values():
+            cam.connect()
+
+    @property
+    def is_calibrated(self) -> bool:
+        return self.left_arm.is_calibrated and self.right_arm.is_calibrated
+
+    def calibrate(self) -> None:
+        self.left_arm.calibrate()
+        self.right_arm.calibrate()
+
+    def configure(self) -> None:
+        self.left_arm.configure()
+        self.right_arm.configure()
+
+    def setup_motors(self) -> None:
+        self.left_arm.setup_motors()
+        self.right_arm.setup_motors()
+
+    def get_observation(self) -> dict[str, Any]:
+        obs_dict = {}
+
+        # Add "left_" prefix
+        left_obs = self.left_arm.get_observation()
+        obs_dict.update({f"left_{key}": value for key, value in left_obs.items()})
+
+        # Add "right_" prefix
+        right_obs = self.right_arm.get_observation()
+        obs_dict.update({f"right_{key}": value for key, value in right_obs.items()})
+
+        for cam_key, cam in self.cameras.items():
+            start = time.perf_counter()
+            obs_dict[cam_key] = cam.async_read()
+            dt_ms = (time.perf_counter() - start) * 1e3
+            logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms")
+
+        return obs_dict
+
+    def send_action(self, action: dict[str, Any]) -> dict[str, Any]:
+        # Remove "left_" prefix
+        left_action = {
+            key.removeprefix("left_"): value for key, value in action.items() if key.startswith("left_")
+        }
+        # Remove "right_" prefix
+        right_action = {
+            key.removeprefix("right_"): value for key, value in action.items() if key.startswith("right_")
+        }
+
+        send_action_left = self.left_arm.send_action(left_action)
+        send_action_right = self.right_arm.send_action(right_action)
+
+        # Add prefixes back
+        prefixed_send_action_left = {f"left_{key}": value for key, value in send_action_left.items()}
+        prefixed_send_action_right = {f"right_{key}": value for key, value in send_action_right.items()}
+
+        return {**prefixed_send_action_left, **prefixed_send_action_right}
+
+    def disconnect(self):
+        self.left_arm.disconnect()
+        self.right_arm.disconnect()
+
+        for cam in self.cameras.values():
+            cam.disconnect()

src/lerobot/robots/aloha/config_aloha.py

@@ -0,0 +1,39 @@
+#!/usr/bin/env python
+
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass, field
+
+from lerobot.cameras import CameraConfig
+
+from ..config import RobotConfig
+
+
+@RobotConfig.register_subclass("aloha")
+@dataclass
+class AlohaConfig(RobotConfig):
+    left_arm_port: str
+    right_arm_port: str
+
+    # Optional parameters for left arm (ViperX)
+    left_arm_max_relative_target: float | dict[str, float] = 20.0
+    left_arm_use_degrees: bool = True
+
+    # Optional parameters for right arm (ViperX)
+    right_arm_max_relative_target: float | dict[str, float] = 20.0
+    right_arm_use_degrees: bool = True
+
+    # cameras (shared between both arms)
+    cameras: dict[str, CameraConfig] = field(default_factory=dict)

src/lerobot/teleoperators/aloha_teleop/__init__.py

@@ -0,0 +1,4 @@
+from .aloha_teleop import AlohaTeleop
+from .config_aloha_teleop import AlohaTeleopConfig
+
+__all__ = ["AlohaTeleop", "AlohaTeleopConfig"]

src/lerobot/teleoperators/aloha_teleop/aloha_teleop.py

@@ -0,0 +1,125 @@
+#!/usr/bin/env python
+
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+from functools import cached_property
+
+from lerobot.teleoperators.widowx.config_widowx import WidowXConfig
+from lerobot.teleoperators.widowx.widowx import WidowX
+
+from ..teleoperator import Teleoperator
+from .config_aloha_teleop import AlohaTeleopConfig
+
+logger = logging.getLogger(__name__)
+
+
+class AlohaTeleop(Teleoperator):
+    """
+    ALOHA Bimanual Teleoperator System using dual WidowX leader arms.
+    Based on the ALOHA (A Low-cost Open-source Hardware System for Bimanual Teleoperation) design.
+    """
+
+    config_class = AlohaTeleopConfig
+    name = "aloha_teleop"
+
+    def __init__(self, config: AlohaTeleopConfig):
+        super().__init__(config)
+        self.config = config
+
+        left_arm_config = WidowXConfig(
+            id=f"{config.id}_left" if config.id else None,
+            calibration_dir=config.calibration_dir,
+            port=config.left_arm_port,
+            gripper_motor=config.left_arm_gripper_motor,
+            use_degrees=config.left_arm_use_degrees,
+        )
+
+        right_arm_config = WidowXConfig(
+            id=f"{config.id}_right" if config.id else None,
+            calibration_dir=config.calibration_dir,
+            port=config.right_arm_port,
+            gripper_motor=config.right_arm_gripper_motor,
+            use_degrees=config.right_arm_use_degrees,
+        )
+
+        self.left_arm = WidowX(left_arm_config)
+        self.right_arm = WidowX(right_arm_config)
+
+    @cached_property
+    def action_features(self) -> dict[str, type]:
+        return {f"left_{motor}.pos": float for motor in self.left_arm.bus.motors} | {
+            f"right_{motor}.pos": float for motor in self.right_arm.bus.motors
+        }
+
+    @cached_property
+    def feedback_features(self) -> dict[str, type]:
+        return {}
+
+    @property
+    def is_connected(self) -> bool:
+        return self.left_arm.is_connected and self.right_arm.is_connected
+
+    def connect(self, calibrate: bool = True) -> None:
+        self.left_arm.connect(calibrate)
+        self.right_arm.connect(calibrate)
+
+    @property
+    def is_calibrated(self) -> bool:
+        return self.left_arm.is_calibrated and self.right_arm.is_calibrated
+
+    def calibrate(self) -> None:
+        self.left_arm.calibrate()
+        self.right_arm.calibrate()
+
+    def configure(self) -> None:
+        self.left_arm.configure()
+        self.right_arm.configure()
+
+    def setup_motors(self) -> None:
+        self.left_arm.setup_motors()
+        self.right_arm.setup_motors()
+
+    def get_action(self) -> dict[str, float]:
+        action_dict = {}
+
+        # Add "left_" prefix
+        left_action = self.left_arm.get_action()
+        action_dict.update({f"left_{key}": value for key, value in left_action.items()})
+
+        # Add "right_" prefix
+        right_action = self.right_arm.get_action()
+        action_dict.update({f"right_{key}": value for key, value in right_action.items()})
+
+        return action_dict
+
+    def send_feedback(self, feedback: dict[str, float]) -> None:
+        # Remove "left_" prefix
+        left_feedback = {
+            key.removeprefix("left_"): value for key, value in feedback.items() if key.startswith("left_")
+        }
+        # Remove "right_" prefix
+        right_feedback = {
+            key.removeprefix("right_"): value for key, value in feedback.items() if key.startswith("right_")
+        }
+
+        if left_feedback:
+            self.left_arm.send_feedback(left_feedback)
+        if right_feedback:
+            self.right_arm.send_feedback(right_feedback)
+
+    def disconnect(self) -> None:
+        self.left_arm.disconnect()
+        self.right_arm.disconnect()

src/lerobot/teleoperators/aloha_teleop/config_aloha_teleop.py

@@ -0,0 +1,34 @@
+#!/usr/bin/env python
+
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+
+from ..config import TeleoperatorConfig
+
+
+@TeleoperatorConfig.register_subclass("aloha_teleop")
+@dataclass
+class AlohaTeleopConfig(TeleoperatorConfig):
+    left_arm_port: str
+    right_arm_port: str
+
+    # Parameters for left arm (WidowX)
+    left_arm_gripper_motor: str = "xl430-w250"
+    left_arm_use_degrees: bool = True
+
+    # Parameters for right arm (WidowX)
+    right_arm_gripper_motor: str = "xc430-w150"
+    right_arm_use_degrees: bool = True