diff --git a/pyproject.toml b/pyproject.toml index 876f3699b5..16f802ec7f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -106,6 +106,7 @@ dynamixel = ["dynamixel-sdk>=3.7.31,<3.9.0"] # Robots gamepad = ["lerobot[pygame-dep]", "hidapi>=0.14.0,<0.15.0"] hopejr = ["lerobot[feetech]", "lerobot[pygame-dep]"] +trossen_arm = ["trossen_arm==1.8.5"] lekiwi = ["lerobot[feetech]", "pyzmq>=26.2.1,<28.0.0"] reachy2 = ["reachy2_sdk>=1.0.14,<1.1.0"] kinematics = ["lerobot[placo-dep]"] diff --git a/src/lerobot/async_inference/constants.py b/src/lerobot/async_inference/constants.py index 1b1dac0f57..101b364555 100644 --- a/src/lerobot/async_inference/constants.py +++ b/src/lerobot/async_inference/constants.py @@ -26,4 +26,4 @@ SUPPORTED_POLICIES = ["act", "smolvla", "diffusion", "tdmpc", "vqbet", "pi0", "pi05"] # TODO: Add all other robots -SUPPORTED_ROBOTS = ["so100_follower", "so101_follower", "bi_so100_follower"] +SUPPORTED_ROBOTS = ["so100_follower", "so101_follower", "bi_so100_follower", "widow_ai_follower"] diff --git a/src/lerobot/async_inference/robot_client.py b/src/lerobot/async_inference/robot_client.py index f9d70a64ec..b6f8ec4e6d 100644 --- a/src/lerobot/async_inference/robot_client.py +++ b/src/lerobot/async_inference/robot_client.py @@ -56,6 +56,7 @@ make_robot_from_config, so100_follower, so101_follower, + widow_ai_follower, ) from lerobot.transport import ( services_pb2, # type: ignore diff --git a/src/lerobot/motors/trossen/__init__.py b/src/lerobot/motors/trossen/__init__.py new file mode 100644 index 0000000000..3a12f5c279 --- /dev/null +++ b/src/lerobot/motors/trossen/__init__.py @@ -0,0 +1,3 @@ +from .widow_ai_arm_driver import TrossenArmDriver + +__all__ = ["TrossenArmDriver"] \ No newline at end of file diff --git a/src/lerobot/motors/trossen/widow_ai_arm_driver.py b/src/lerobot/motors/trossen/widow_ai_arm_driver.py new file mode 100644 index 0000000000..8e9d2476f2 --- /dev/null +++ b/src/lerobot/motors/trossen/widow_ai_arm_driver.py @@ -0,0 +1,137 @@ +import logging +import traceback +from typing import List, Union + +import numpy as np +import trossen_arm as trossen + + +class TrossenArmDriver: + def __init__( + self, + port: str, + model: str = "V0_LEADER", + velocity_limit_scale: float | None = None, + ): + self.port = port + self.model = model + self.velocity_limit_scale = velocity_limit_scale + self.driver = None + self.is_connected = False + + self.MIN_TIME_TO_MOVE = 0.1 # seconds + + @property + def is_connected(self) -> bool: + return self._is_connected and self.driver is not None + + @is_connected.setter + def is_connected(self, value: bool): + self._is_connected = value + + def _get_model_config(self): + trossen_arm_models = { + "V0_LEADER": [trossen.Model.wxai_v0, trossen.StandardEndEffector.wxai_v0_leader], + "V0_FOLLOWER": [trossen.Model.wxai_v0, trossen.StandardEndEffector.wxai_v0_follower], + } + + try: + return trossen_arm_models[self.model] + except KeyError as e: + raise ValueError(f"Unsupported model: {self.model}") from e + + def _scale_velocity_limits(self, scale: float) -> None: + logging.debug(f"Scaling {self.model} arm velocity limits to {scale * 100}% for safety...") + + joint_limits = self.driver.get_joint_limits() + + for i in range(len(joint_limits)): + original_velocity = joint_limits[i].velocity_max + joint_limits[i].velocity_max *= scale + logging.debug(f" Joint {i}: velocity_max scaled from {original_velocity:.3f} to {joint_limits[i].velocity_max:.3f}") + + self.driver.set_joint_limits(joint_limits) + logging.debug(f"{self.model} arm velocity limits successfully scaled.") + + def connect(self) -> None: + if self.is_connected: + logging.debug(f"TrossenArmDriver({self.port}) is already connected.") + return + + logging.debug(f"Connecting to {self.model} arm at {self.port}...") + + self.driver = trossen.TrossenArmDriver() + + model_name, model_end_effector = self._get_model_config() + + logging.debug("Configuring the drivers...") + + try: + self.driver.configure(model_name, model_end_effector, self.port, False) + except Exception as e: + traceback.print_exc() + logging.error(f"Failed to configure the driver for the {self.model} arm at {self.port}.") + raise e + + self.is_connected = True + + if self.velocity_limit_scale is not None: + self._scale_velocity_limits(self.velocity_limit_scale) + + def disconnect(self) -> None: + if not self.is_connected: + return + self.is_connected = False + logging.debug(f"{self.model} arm disconnected.") + + def read(self, data_name: str) -> np.ndarray: + if not self.is_connected: + raise RuntimeError(f"TrossenArmDriver({self.port}) is not connected. You need to run `connect()`.") + + if data_name == "Present_Position": + values = self.driver.get_all_positions() + elif data_name == "External_Efforts": + values = self.driver.get_all_external_efforts() + else: + raise ValueError(f"Data name: {data_name} is not supported for reading.") + + return np.array(values, dtype=np.float32) + + def write(self, data_name: str, values: Union[float, List[float], np.ndarray]) -> None: + if not self.is_connected: + raise RuntimeError(f"TrossenArmDriver({self.port}) is not connected. You need to run `connect()`.") + + if data_name == "Goal_Position": + values = np.array(values, dtype=np.float32) + self.driver.set_all_positions(values.tolist(), self.MIN_TIME_TO_MOVE, False) + elif data_name == "External_Efforts": + values = np.array(values, dtype=np.float32) + self.driver.set_all_external_efforts(values.tolist(), 0.0, False) + else: + raise ValueError(f"Data name: {data_name} is not supported for writing.") + + def initialize_for_teleoperation(self, is_leader: bool = True) -> None: + logging.debug(f"Initializing {self.model} arm for teleoperation...") + + logging.debug(f"Setting {self.model} arm mode for teleoperation...") + if is_leader: + self.driver.set_all_modes(trossen.Mode.external_effort) + # Send zero efforts immediately to prevent instability + logging.debug(f"Sending zero efforts to {self.model} arm for stability...") + try: + zero_efforts = [0.0] * self.driver.get_num_joints() + self.driver.set_all_external_efforts(zero_efforts, 0.0, False) + logging.debug(f"Zero efforts sent to {self.model} arm - should be stable now") + except Exception as e: + logging.warning(f"Failed to send zero efforts to {self.model} arm: {e}") + else: + self.driver.set_all_modes(trossen.Mode.position) + + logging.debug(f"{self.model} arm initialization complete!") + + def __del__(self): + try: + if getattr(self, "_is_connected", False): + self.disconnect() + except Exception as e: + logging.warning(f"Exception during cleanup: {e}") diff --git a/src/lerobot/robots/utils.py b/src/lerobot/robots/utils.py index 4e80015385..a1e1f52679 100644 --- a/src/lerobot/robots/utils.py +++ b/src/lerobot/robots/utils.py @@ -56,6 +56,10 @@ def make_robot_from_config(config: RobotConfig) -> Robot: from .reachy2 import Reachy2Robot return Reachy2Robot(config) + elif config.type == "widow_ai_follower": + from .widow_ai_follower import WidowAIFollower + + return WidowAIFollower(config) elif config.type == "mock_robot": from tests.mocks.mock_robot import MockRobot diff --git a/src/lerobot/robots/widow_ai_follower/__init__.py b/src/lerobot/robots/widow_ai_follower/__init__.py new file mode 100644 index 0000000000..5955d2ace1 --- /dev/null +++ b/src/lerobot/robots/widow_ai_follower/__init__.py @@ -0,0 +1,2 @@ +from .config_widow_ai_follower import WidowAIFollowerConfig +from .widow_ai_follower import WidowAIFollower diff --git a/src/lerobot/robots/widow_ai_follower/config_widow_ai_follower.py b/src/lerobot/robots/widow_ai_follower/config_widow_ai_follower.py new file mode 100644 index 0000000000..7ce46e7ce6 --- /dev/null +++ b/src/lerobot/robots/widow_ai_follower/config_widow_ai_follower.py @@ -0,0 +1,32 @@ +#!/usr/bin/env python + +from dataclasses import dataclass, field + +from lerobot.cameras import CameraConfig + +from ..config import RobotConfig + + +@RobotConfig.register_subclass("widow_ai_follower") +@dataclass +class WidowAIFollowerConfig(RobotConfig): + # Port to connect to the arm (IP address for Trossen arms) + port: str + + # Trossen arm model to use + model: str = "V0_FOLLOWER" # Options: "V0_LEADER", "V0_FOLLOWER" + + # Velocity limit scale for safety (1.0 = 100% of max velocity, 0.5 = 50% of max velocity) + # Scales maximum joint velocities on connection. Set to None to skip scaling. + velocity_limit_scale: float | None = None + + # `max_relative_target` limits the magnitude of the relative positional target vector for safety purposes. + # Set this to a positive scalar to have the same value for all motors, or a list that is the same length as + # the number of motors in your follower arms. + max_relative_target: float | None = None + + # cameras + cameras: dict[str, CameraConfig] = field(default_factory=dict) + + # Enable effort sensing to include effort measurements in observations + effort_sensing: bool = False diff --git a/src/lerobot/robots/widow_ai_follower/widow_ai_follower.py b/src/lerobot/robots/widow_ai_follower/widow_ai_follower.py new file mode 100644 index 0000000000..7496f82b14 --- /dev/null +++ b/src/lerobot/robots/widow_ai_follower/widow_ai_follower.py @@ -0,0 +1,183 @@ +#!/usr/bin/env python + +import logging +import time +from functools import cached_property +from typing import Any + +from lerobot.cameras.utils import make_cameras_from_configs +from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError +from lerobot.motors.trossen import TrossenArmDriver + +from ..robot import Robot +from ..utils import ensure_safe_goal_position +from .config_widow_ai_follower import WidowAIFollowerConfig + +logger = logging.getLogger(__name__) + + +class WidowAIFollower(Robot): + config_class = WidowAIFollowerConfig + name = "widow_ai_follower" + + def __init__(self, config: WidowAIFollowerConfig): + super().__init__(config) + self.config = config + + self.bus = TrossenArmDriver( + port=self.config.port, + model=self.config.model, + velocity_limit_scale=self.config.velocity_limit_scale, + ) + + self.motor_names = ["shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_1", "wrist_2", "wrist_3", "gripper"] + + self.cameras = make_cameras_from_configs(config.cameras) + + @property + def _motors_ft(self) -> dict[str, type]: + return {f"{motor}.pos": float for motor in self.motor_names} + + @property + def _efforts_ft(self) -> dict[str, type]: + return {f"{motor}.effort": float for motor in self.motor_names} + + @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]: + if self.config.effort_sensing: + return {**self._motors_ft, **self._efforts_ft, **self._cameras_ft} + else: + 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.bus.is_connected and all(cam.is_connected for cam in self.cameras.values()) + + def connect(self, calibrate: bool = True) -> None: + if self.is_connected: + raise DeviceAlreadyConnectedError(f"{self} already connected") + + self.bus.connect() + + for cam in self.cameras.values(): + cam.connect() + + self.configure() + logger.info(f"{self} connected.") + + @property + def is_calibrated(self) -> bool: + # Trossen arms are pre-calibrated + return True + + def calibrate(self) -> None: + logger.info(f"{self} is pre-calibrated, no calibration needed.") + + def configure(self) -> None: + self.bus.initialize_for_teleoperation(is_leader=False) + + def setup_motors(self) -> None: + logger.info(f"{self} motors are pre-configured.") + + def get_observation(self) -> dict[str, Any]: + if not self.is_connected: + raise DeviceNotConnectedError(f"{self} is not connected.") + + start = time.perf_counter() + positions = self.bus.read("Present_Position") + + # Convert numpy array to dict with motor names + obs_dict = {} + for i, motor in enumerate(self.motor_names): + if i < len(positions): + obs_dict[f"{motor}.pos"] = float(positions[i]) + else: + obs_dict[f"{motor}.pos"] = 0.0 + + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read state: {dt_ms:.1f}ms") + + if self.config.effort_sensing: + try: + start = time.perf_counter() + efforts = self.bus.read("External_Efforts") + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read external efforts: {dt_ms:.1f}ms") + + # Convert numpy array to dict with motor names + for i, motor in enumerate(self.motor_names): + if i < len(efforts): + obs_dict[f"{motor}.effort"] = float(efforts[i]) + else: + obs_dict[f"{motor}.effort"] = 0.0 + except Exception as e: + logger.debug(f"{self} failed to read external efforts, using zeros: {e}") + for motor in self.motor_names: + obs_dict[f"{motor}.effort"] = 0.0 + + # Capture images from cameras + 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]: + """Command arm to move to a target joint configuration. + + Raises: + RobotDeviceNotConnectedError: if robot is not connected. + + Returns: + the action sent to the motors, potentially clipped. + """ + if not self.is_connected: + raise DeviceNotConnectedError(f"{self} is not connected.") + + start = time.perf_counter() + + # Extract goal positions from action dict + goal_pos = [] + for motor in self.motor_names: + pos_key = f"{motor}.pos" + if pos_key in action: + goal_pos.append(action[pos_key]) + else: + goal_pos.append(0.0) + + # Cap goal position when too far away from present position. + if self.config.max_relative_target is not None: + present_pos = self.bus.read("Present_Position") + present_dict = {motor: float(present_pos[i]) for i, motor in enumerate(self.motor_names) if i < len(present_pos)} + goal_dict = {motor: goal_pos[i] for i, motor in enumerate(self.motor_names)} + goal_present_pos = {key: (goal_dict[key], present_dict.get(key, 0.0)) for key in goal_dict} + safe_goal_dict = ensure_safe_goal_position(goal_present_pos, self.config.max_relative_target) + goal_pos = [safe_goal_dict.get(motor, 0.0) for motor in self.motor_names] + + self.bus.write("Goal_Position", goal_pos) + + dt_ms = (time.perf_counter() - start) * 1e3 + + return {f"{motor}.pos": val for motor, val in zip(self.motor_names, goal_pos)} + + def disconnect(self): + if not self.is_connected: + raise DeviceNotConnectedError(f"{self} is not connected.") + + self.bus.disconnect() + for cam in self.cameras.values(): + cam.disconnect() + + logger.info(f"{self} disconnected.") diff --git a/src/lerobot/scripts/lerobot_record.py b/src/lerobot/scripts/lerobot_record.py index 6df92d893b..7ff112a9eb 100644 --- a/src/lerobot/scripts/lerobot_record.py +++ b/src/lerobot/scripts/lerobot_record.py @@ -98,6 +98,7 @@ make_robot_from_config, so100_follower, so101_follower, + widow_ai_follower, ) from lerobot.teleoperators import ( # noqa: F401 Teleoperator, @@ -108,6 +109,7 @@ make_teleoperator_from_config, so100_leader, so101_leader, + widow_ai_leader, ) from lerobot.teleoperators.keyboard.teleop_keyboard import KeyboardTeleop from lerobot.utils.constants import ACTION, OBS_STR diff --git a/src/lerobot/scripts/lerobot_setup_motors.py b/src/lerobot/scripts/lerobot_setup_motors.py index c1d256c211..d8c6eff78e 100644 --- a/src/lerobot/scripts/lerobot_setup_motors.py +++ b/src/lerobot/scripts/lerobot_setup_motors.py @@ -35,6 +35,7 @@ make_robot_from_config, so100_follower, so101_follower, + widow_ai_follower, ) from lerobot.teleoperators import ( # noqa: F401 TeleoperatorConfig, @@ -42,6 +43,7 @@ make_teleoperator_from_config, so100_leader, so101_leader, + widow_ai_leader, ) COMPATIBLE_DEVICES = [ @@ -52,6 +54,8 @@ "so101_follower", "so101_leader", "lekiwi", + "widow_ai_follower", + "widow_ai_leader", ] diff --git a/src/lerobot/scripts/lerobot_teleoperate.py b/src/lerobot/scripts/lerobot_teleoperate.py index 0a418f3bca..84e53be470 100644 --- a/src/lerobot/scripts/lerobot_teleoperate.py +++ b/src/lerobot/scripts/lerobot_teleoperate.py @@ -76,6 +76,7 @@ make_robot_from_config, so100_follower, so101_follower, + widow_ai_follower, ) from lerobot.teleoperators import ( # noqa: F401 Teleoperator, @@ -87,6 +88,7 @@ make_teleoperator_from_config, so100_leader, so101_leader, + widow_ai_leader, ) from lerobot.utils.import_utils import register_third_party_devices from lerobot.utils.robot_utils import busy_wait @@ -156,6 +158,11 @@ def teleop_loop( # Send processed action to robot (robot_action_processor.to_output should return dict[str, Any]) _ = robot.send_action(robot_action_to_send) + # Effort feedback + effort_feedback = {key: val for key, val in obs.items() if key.endswith(".effort")} + if effort_feedback: + teleop.send_feedback(effort_feedback) + if display_data: # Process robot observation through pipeline obs_transition = robot_observation_processor(obs) diff --git a/src/lerobot/teleoperators/utils.py b/src/lerobot/teleoperators/utils.py index 2103a1669c..a3410a0d27 100644 --- a/src/lerobot/teleoperators/utils.py +++ b/src/lerobot/teleoperators/utils.py @@ -49,6 +49,10 @@ def make_teleoperator_from_config(config: TeleoperatorConfig) -> Teleoperator: from .so101_leader import SO101Leader return SO101Leader(config) + elif config.type == "widow_ai_leader": + from .widow_ai_leader import WidowAILeader + + return WidowAILeader(config) elif config.type == "mock_teleop": from tests.mocks.mock_teleop import MockTeleop diff --git a/src/lerobot/teleoperators/widow_ai_leader/__init__.py b/src/lerobot/teleoperators/widow_ai_leader/__init__.py new file mode 100644 index 0000000000..0912cb4571 --- /dev/null +++ b/src/lerobot/teleoperators/widow_ai_leader/__init__.py @@ -0,0 +1,2 @@ +from .config_widow_ai_leader import WidowAILeaderConfig +from .widow_ai_leader import WidowAILeader diff --git a/src/lerobot/teleoperators/widow_ai_leader/config_widow_ai_leader.py b/src/lerobot/teleoperators/widow_ai_leader/config_widow_ai_leader.py new file mode 100644 index 0000000000..102adf42e5 --- /dev/null +++ b/src/lerobot/teleoperators/widow_ai_leader/config_widow_ai_leader.py @@ -0,0 +1,22 @@ +#!/usr/bin/env python + +from dataclasses import dataclass + +from ..config import TeleoperatorConfig + + +@TeleoperatorConfig.register_subclass("widow_ai_leader") +@dataclass +class WidowAILeaderConfig(TeleoperatorConfig): + # Port to connect to the arm (IP address for Trossen arms) + port: str + + # Trossen arm model to use + model: str = "V0_LEADER" + + # Velocity limit scale for safety (1.0 = 100% of max velocity, 0.5 = 50% of max velocity) + # Scales maximum joint velocities on connection. Set to None to skip scaling. + velocity_limit_scale: float | None = None + + # Effort feedback gain for haptic feedback + effort_feedback_gain: float = 0.1 diff --git a/src/lerobot/teleoperators/widow_ai_leader/widow_ai_leader.py b/src/lerobot/teleoperators/widow_ai_leader/widow_ai_leader.py new file mode 100644 index 0000000000..c6f81c4839 --- /dev/null +++ b/src/lerobot/teleoperators/widow_ai_leader/widow_ai_leader.py @@ -0,0 +1,100 @@ +#!/usr/bin/env python + +import logging +import time + +from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError +from lerobot.motors.trossen import TrossenArmDriver + +from ..teleoperator import Teleoperator +from .config_widow_ai_leader import WidowAILeaderConfig + +logger = logging.getLogger(__name__) + + +class WidowAILeader(Teleoperator): + config_class = WidowAILeaderConfig + name = "widow_ai_leader" + + def __init__(self, config: WidowAILeaderConfig): + super().__init__(config) + self.config = config + + self.bus = TrossenArmDriver( + port=self.config.port, + model=self.config.model, + velocity_limit_scale=self.config.velocity_limit_scale, + ) + + # Define motor names for compatibility + self.motor_names = ["shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_1", "wrist_2", "wrist_3", "gripper"] + + @property + def action_features(self) -> dict[str, type]: + return {f"{motor}.pos": float for motor in self.motor_names} + + @property + def feedback_features(self) -> dict[str, type]: + return {f"{motor}.effort": float for motor in self.motor_names} + + @property + def is_connected(self) -> bool: + return self.bus.is_connected + + def connect(self, calibrate: bool = True) -> None: + if self.is_connected: + raise DeviceAlreadyConnectedError(f"{self} already connected") + + self.bus.connect() + self.configure() + logger.info(f"{self} connected.") + + @property + def is_calibrated(self) -> bool: + return True + + def calibrate(self) -> None: + logger.info(f"{self} is pre-calibrated, no calibration needed.") + + def configure(self) -> None: + self.bus.initialize_for_teleoperation(is_leader=True) + + def setup_motors(self) -> None: + logger.info(f"{self} motors are pre-configured.") + + def get_action(self) -> dict[str, float]: + start = time.perf_counter() + positions = self.bus.read("Present_Position") + + # Convert numpy array to dict with motor names + action = {} + for i, motor in enumerate(self.motor_names): + if i < len(positions): + action[f"{motor}.pos"] = float(positions[i]) + else: + action[f"{motor}.pos"] = 0.0 + + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read action: {dt_ms:.1f}ms") + return action + + def send_feedback(self, feedback: dict[str, float]) -> None: + """Send effort feedback to the leader arm.""" + effort_feedback = [] + for motor in self.motor_names: + effort_key = f"{motor}.effort" + if effort_key in feedback: + effort_feedback.append(-1 * self.config.effort_feedback_gain * feedback[effort_key]) + else: + effort_feedback.append(0.0) + + if effort_feedback: + self.bus.write("External_Efforts", effort_feedback) + logger.debug(f"{self} sent effort feedback: {effort_feedback}") + + def disconnect(self) -> None: + if not self.is_connected: + raise DeviceNotConnectedError(f"{self} is not connected.") + + self.bus.disconnect() + logger.info(f"{self} disconnected.")