style nit

This commit is contained in:
Michel Aractingi
2025-08-12 18:04:28 +02:00
parent f76a108b08
commit 73fab32c26
4 changed files with 71 additions and 75 deletions
+1
View File
@@ -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
@@ -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()
self.leader_device.reset()
+13 -7
View File
@@ -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 = (
@@ -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
}
self.keyboard_events = {"intervention": False, "success": False, "failure": False, "rerecord": False}