diff --git a/docs/source/hilserl.mdx b/docs/source/hilserl.mdx index 42add2677..cea2ec728 100644 --- a/docs/source/hilserl.mdx +++ b/docs/source/hilserl.mdx @@ -500,6 +500,7 @@ To setup the SO101 leader, you need to set the `control_mode` to `"leader"` and ``` The `leader_follower_mode` enables the leader arm to automatically track the follower's position when you're not intervening. This creates a seamless teleoperation experience where: + - When not intervening: the leader arm follows the follower arm's position - When intervening (press `space`): you control the leader arm, and the follower tracks it in end-effector space diff --git a/src/lerobot/processor/leader_follower_processor.py b/src/lerobot/processor/leader_follower_processor.py index 489994db7..dc5a28706 100644 --- a/src/lerobot/processor/leader_follower_processor.py +++ b/src/lerobot/processor/leader_follower_processor.py @@ -14,16 +14,16 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dataclasses import dataclass + import numpy as np import torch -from dataclasses import dataclass - from lerobot.model.kinematics import RobotKinematics from lerobot.processor.pipeline import EnvTransition, ProcessorStepRegistry, TransitionKey -from lerobot.teleoperators.utils import TeleopEvents -from lerobot.teleoperators import Teleoperator from lerobot.robots import Robot +from lerobot.teleoperators import Teleoperator +from lerobot.teleoperators.utils import TeleopEvents @ProcessorStepRegistry.register("leader_follower_processor") @@ -31,13 +31,13 @@ from lerobot.robots import Robot class LeaderFollowerProcessor: """ Processor for leader-follower teleoperation mode. - + This processor: 1. Sends follower positions to leader arm when not intervening 2. Computes EE delta actions from leader when intervening 3. Handles teleop events from the leader device """ - + leader_device: Teleoperator motor_names: list[str] robot: Robot @@ -46,19 +46,16 @@ class LeaderFollowerProcessor: use_gripper: bool = True prev_leader_gripper: float | None = None max_gripper_pos: float = 100.0 - + def __call__(self, transition: EnvTransition) -> EnvTransition: """Process transition with leader-follower logic.""" # Get current follower position from complementary data raw_joint_pos = transition.get(TransitionKey.COMPLEMENTARY_DATA, {}).get("raw_joint_positions") if raw_joint_pos is not None: # Send follower position to leader (for follow mode) - follower_action = { - f"{motor}.pos": float(raw_joint_pos[motor]) - for motor in self.motor_names - } + follower_action = {f"{motor}.pos": float(raw_joint_pos[motor]) for motor in self.motor_names} self.leader_device.send_action(follower_action) - + # Only compute EE action if intervention is active # (AddTeleopEventsAsInfo already added IS_INTERVENTION to info) info = transition.get(TransitionKey.INFO, {}) @@ -67,38 +64,34 @@ class LeaderFollowerProcessor: # (AddTeleopActionAsComplimentaryData already got the action) complementary = transition.get(TransitionKey.COMPLEMENTARY_DATA, {}) teleop_action = complementary.get("teleop_action", {}) - + if isinstance(teleop_action, dict) and raw_joint_pos is not None: # Extract leader positions from teleop action dict leader_pos = np.array([teleop_action.get(f"{motor}.pos", 0) for motor in self.motor_names]) follower_pos = np.array([raw_joint_pos[motor] for motor in self.motor_names]) - + # Compute EE positions leader_ee = self.kinematics.forward_kinematics(leader_pos)[:3, 3] follower_ee = self.kinematics.forward_kinematics(follower_pos)[:3, 3] - + # Compute normalized EE delta if self.end_effector_step_sizes is not None: ee_delta = np.clip( - leader_ee - follower_ee, - -self.end_effector_step_sizes, - self.end_effector_step_sizes + leader_ee - follower_ee, -self.end_effector_step_sizes, self.end_effector_step_sizes ) ee_delta_normalized = ee_delta / self.end_effector_step_sizes else: ee_delta_normalized = leader_ee - follower_ee - + # Handle gripper if self.use_gripper and len(leader_pos) > 3: if self.prev_leader_gripper is None: - self.prev_leader_gripper = np.clip( - leader_pos[-1], 0, self.max_gripper_pos - ) - + self.prev_leader_gripper = np.clip(leader_pos[-1], 0, self.max_gripper_pos) + leader_gripper = leader_pos[-1] gripper_delta = leader_gripper - self.prev_leader_gripper normalized_delta = gripper_delta / self.max_gripper_pos - + # Quantize gripper action if normalized_delta >= 0.3: gripper_action = 2 @@ -106,22 +99,22 @@ class LeaderFollowerProcessor: gripper_action = 0 else: gripper_action = 1 - + self.prev_leader_gripper = leader_gripper - + # Create intervention action intervention_action = np.append(ee_delta_normalized, gripper_action) else: intervention_action = ee_delta_normalized - + # Override teleop_action with computed EE action complementary["teleop_action"] = torch.from_numpy(intervention_action).float() transition[TransitionKey.COMPLEMENTARY_DATA] = complementary # type: ignore[misc] - + return transition - + def reset(self) -> None: """Reset leader-follower state.""" self.prev_leader_gripper = None if hasattr(self.leader_device, "reset"): - self.leader_device.reset() \ No newline at end of file + self.leader_device.reset() diff --git a/src/lerobot/scripts/rl/gym_manipulator.py b/src/lerobot/scripts/rl/gym_manipulator.py index 4691c3599..3de0a822a 100644 --- a/src/lerobot/scripts/rl/gym_manipulator.py +++ b/src/lerobot/scripts/rl/gym_manipulator.py @@ -463,15 +463,17 @@ def make_processors( # Get control mode control_mode = cfg.processor.control_mode if cfg.processor is not None else "gamepad" - + action_pipeline_steps = [ AddTeleopActionAsComplimentaryData(teleop_device=teleop_device), AddTeleopEventsAsInfo(teleop_device=teleop_device), - AddRobotObservationAsComplimentaryData(robot=env.robot) + AddRobotObservationAsComplimentaryData(robot=env.robot), ] # Check for leader control mode if control_mode == "leader": - assert isinstance(teleop_device, SO101LeaderFollower), "Leader control mode requires SO101LeaderFollower teleop device" + assert isinstance(teleop_device, SO101LeaderFollower), ( + "Leader control mode requires SO101LeaderFollower teleop device" + ) action_pipeline_steps.append( LeaderFollowerProcessor( @@ -479,11 +481,15 @@ def make_processors( motor_names=motor_names, robot=env.robot, kinematics=kinematics_solver, - end_effector_step_sizes=np.array(list(cfg.processor.inverse_kinematics.end_effector_step_sizes.values())), + end_effector_step_sizes=np.array( + list(cfg.processor.inverse_kinematics.end_effector_step_sizes.values()) + ), use_gripper=cfg.processor.gripper.use_gripper if cfg.processor.gripper is not None else False, - max_gripper_pos=cfg.processor.max_gripper_pos if cfg.processor.max_gripper_pos is not None else 100.0, - ) + max_gripper_pos=cfg.processor.max_gripper_pos + if cfg.processor.max_gripper_pos is not None + else 100.0, ) + ) # Standard teleop mode (gamepad, keyboard, etc.) action_pipeline_steps.append( InterventionActionProcessor( @@ -596,7 +602,7 @@ def control_loop( dt = 1.0 / cfg.env.fps print(f"Starting control loop at {cfg.env.fps} FPS") - + # Reset environment and processors obs, info = env.reset() complementary_data = ( diff --git a/src/lerobot/teleoperators/so101_leader/so101_leader_follower.py b/src/lerobot/teleoperators/so101_leader/so101_leader_follower.py index 2356b59c6..f2ada02db 100644 --- a/src/lerobot/teleoperators/so101_leader/so101_leader_follower.py +++ b/src/lerobot/teleoperators/so101_leader/so101_leader_follower.py @@ -31,39 +31,39 @@ logger = logging.getLogger(__name__) class SO101LeaderFollower(SO101Leader): """ Extended SO101 Leader that can both lead (human control) and follow (mimic follower). - + This class adds leader-follower functionality where: - In follow mode: The leader arm mimics the follower's position (torque enabled) - In lead mode: Human controls the leader (torque disabled) and provides actions """ - + def __init__(self, config): super().__init__(config) - + # Leader-follower state self.is_intervening = False self.leader_torque_enabled = True - + # Tracking error for automatic intervention detection self.leader_tracking_error_queue = deque(maxlen=4) - + # Keyboard event handling self.keyboard_events = { "intervention": False, - "success": False, + "success": False, "failure": False, "rerecord": False, } self.keyboard_thread = None self.stop_event = Event() - + # Store last follower position for action computation self.last_follower_pos = None - + def connect(self, calibrate: bool = True) -> None: """Connect and configure for leader-follower mode.""" super().connect(calibrate) - + # Configure for leader-follower mode with lower gains # Lower gains allow manual intervention without injury risk self.bus.sync_write("Torque_Enable", 1) @@ -71,7 +71,7 @@ class SO101LeaderFollower(SO101Leader): self.bus.write("P_Coefficient", motor, 16) self.bus.write("I_Coefficient", motor, 0) self.bus.write("D_Coefficient", motor, 16) - + # Start keyboard listener self._start_keyboard_listener() @@ -82,9 +82,10 @@ class SO101LeaderFollower(SO101Leader): print(" - Press 's' to mark episode as success") print(" - Press ESC to end episode as failure") print(" - Press 'r' to re-record episode") - + def _start_keyboard_listener(self): """Start keyboard listener thread for intervention control.""" + def on_press(key): try: if key == keyboard.Key.space: @@ -94,68 +95,68 @@ class SO101LeaderFollower(SO101Leader): logger.info(f"Toggled to {state}") elif key == keyboard.Key.esc: self.keyboard_events["failure"] = True - elif hasattr(key, 'char'): - if key.char == 's': + elif hasattr(key, "char"): + if key.char == "s": self.keyboard_events["success"] = True - elif key.char == 'r': + elif key.char == "r": self.keyboard_events["rerecord"] = True except Exception as e: logger.error(f"Error handling key press: {e}") - + def listen(): with keyboard.Listener(on_press=on_press) as listener: while not self.stop_event.is_set(): time.sleep(0.1) listener.stop() - + self.keyboard_thread = Thread(target=listen, daemon=True) self.keyboard_thread.start() - + def send_action(self, action: dict[str, float]) -> None: """ Send position commands to leader arm (follow mode). - + Args: action: Dictionary of motor positions to command """ # Store follower position for later use self.last_follower_pos = np.array([action.get(f"{motor}.pos", 0) for motor in self.bus.motors]) - + if not self.is_intervening: # Follow mode: enable torque and track follower if not self.leader_torque_enabled: self.bus.sync_write("Torque_Enable", 1) self.leader_torque_enabled = True - + # Send follower positions to leader goal_pos = {motor: action[f"{motor}.pos"] for motor in self.bus.motors} self.bus.sync_write("Goal_Position", goal_pos) - + # Track error for automatic intervention detection current_pos = self.bus.sync_read("Present_Position") current_array = np.array([current_pos[motor] for motor in self.bus.motors]) error = np.linalg.norm(self.last_follower_pos[:-1] - current_array[:-1]) self.leader_tracking_error_queue.append(error) - + def get_action(self) -> dict[str, float]: """ Get action from leader arm. - + In follow mode: Returns neutral/current positions In lead mode: Returns actual leader positions for follower to track """ start = time.perf_counter() - + if self.is_intervening: # Lead mode: disable torque if needed and return leader positions if self.leader_torque_enabled: self.bus.sync_write("Torque_Enable", 0) self.leader_torque_enabled = False - + # Get current leader position action = self.bus.sync_read("Present_Position") action = {f"{motor}.pos": val for motor, val in action.items()} - + # Track error if self.last_follower_pos is not None: current_array = np.array([action[f"{motor}.pos"] for motor in self.bus.motors]) @@ -165,15 +166,15 @@ class SO101LeaderFollower(SO101Leader): # Follow mode: return current/neutral positions action = self.bus.sync_read("Present_Position") action = {f"{motor}.pos": val for motor, val in action.items()} - + dt_ms = (time.perf_counter() - start) * 1e3 logger.debug(f"{self} read action: {dt_ms:.1f}ms") return action - + def get_teleop_events(self) -> dict[TeleopEvents, bool]: """Get current keyboard events.""" events = {} - + # Map keyboard events to TeleopEvents if self.keyboard_events["success"]: events[TeleopEvents.SUCCESS] = True @@ -186,12 +187,12 @@ class SO101LeaderFollower(SO101Leader): events[TeleopEvents.RERECORD_EPISODE] = True events[TeleopEvents.TERMINATE_EPISODE] = True self.keyboard_events["rerecord"] = False - + # Always report intervention state events[TeleopEvents.IS_INTERVENTION] = self.is_intervening - + return events - + def disconnect(self) -> None: """Disconnect and cleanup.""" self.stop_event.set() @@ -204,9 +205,4 @@ class SO101LeaderFollower(SO101Leader): self.is_intervening = False self.leader_torque_enabled = True self.leader_tracking_error_queue.clear() - self.keyboard_events = { - "intervention": False, - "success": False, - "failure": False, - "rerecord": False - } \ No newline at end of file + self.keyboard_events = {"intervention": False, "success": False, "failure": False, "rerecord": False}