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: 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 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 - 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 # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
from dataclasses import dataclass
import numpy as np import numpy as np
import torch import torch
from dataclasses import dataclass
from lerobot.model.kinematics import RobotKinematics from lerobot.model.kinematics import RobotKinematics
from lerobot.processor.pipeline import EnvTransition, ProcessorStepRegistry, TransitionKey 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.robots import Robot
from lerobot.teleoperators import Teleoperator
from lerobot.teleoperators.utils import TeleopEvents
@ProcessorStepRegistry.register("leader_follower_processor") @ProcessorStepRegistry.register("leader_follower_processor")
@@ -53,10 +53,7 @@ class LeaderFollowerProcessor:
raw_joint_pos = transition.get(TransitionKey.COMPLEMENTARY_DATA, {}).get("raw_joint_positions") raw_joint_pos = transition.get(TransitionKey.COMPLEMENTARY_DATA, {}).get("raw_joint_positions")
if raw_joint_pos is not None: if raw_joint_pos is not None:
# Send follower position to leader (for follow mode) # Send follower position to leader (for follow mode)
follower_action = { follower_action = {f"{motor}.pos": float(raw_joint_pos[motor]) for motor in self.motor_names}
f"{motor}.pos": float(raw_joint_pos[motor])
for motor in self.motor_names
}
self.leader_device.send_action(follower_action) self.leader_device.send_action(follower_action)
# Only compute EE action if intervention is active # Only compute EE action if intervention is active
@@ -80,9 +77,7 @@ class LeaderFollowerProcessor:
# Compute normalized EE delta # Compute normalized EE delta
if self.end_effector_step_sizes is not None: if self.end_effector_step_sizes is not None:
ee_delta = np.clip( ee_delta = np.clip(
leader_ee - follower_ee, leader_ee - follower_ee, -self.end_effector_step_sizes, self.end_effector_step_sizes
-self.end_effector_step_sizes,
self.end_effector_step_sizes
) )
ee_delta_normalized = ee_delta / self.end_effector_step_sizes ee_delta_normalized = ee_delta / self.end_effector_step_sizes
else: else:
@@ -91,9 +86,7 @@ class LeaderFollowerProcessor:
# Handle gripper # Handle gripper
if self.use_gripper and len(leader_pos) > 3: if self.use_gripper and len(leader_pos) > 3:
if self.prev_leader_gripper is None: if self.prev_leader_gripper is None:
self.prev_leader_gripper = np.clip( self.prev_leader_gripper = np.clip(leader_pos[-1], 0, self.max_gripper_pos)
leader_pos[-1], 0, self.max_gripper_pos
)
leader_gripper = leader_pos[-1] leader_gripper = leader_pos[-1]
gripper_delta = leader_gripper - self.prev_leader_gripper gripper_delta = leader_gripper - self.prev_leader_gripper
+11 -5
View File
@@ -467,11 +467,13 @@ def make_processors(
action_pipeline_steps = [ action_pipeline_steps = [
AddTeleopActionAsComplimentaryData(teleop_device=teleop_device), AddTeleopActionAsComplimentaryData(teleop_device=teleop_device),
AddTeleopEventsAsInfo(teleop_device=teleop_device), AddTeleopEventsAsInfo(teleop_device=teleop_device),
AddRobotObservationAsComplimentaryData(robot=env.robot) AddRobotObservationAsComplimentaryData(robot=env.robot),
] ]
# Check for leader control mode # Check for leader control mode
if control_mode == "leader": 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( action_pipeline_steps.append(
LeaderFollowerProcessor( LeaderFollowerProcessor(
@@ -479,11 +481,15 @@ def make_processors(
motor_names=motor_names, motor_names=motor_names,
robot=env.robot, robot=env.robot,
kinematics=kinematics_solver, 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, 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.) # Standard teleop mode (gamepad, keyboard, etc.)
action_pipeline_steps.append( action_pipeline_steps.append(
InterventionActionProcessor( InterventionActionProcessor(
@@ -85,6 +85,7 @@ class SO101LeaderFollower(SO101Leader):
def _start_keyboard_listener(self): def _start_keyboard_listener(self):
"""Start keyboard listener thread for intervention control.""" """Start keyboard listener thread for intervention control."""
def on_press(key): def on_press(key):
try: try:
if key == keyboard.Key.space: if key == keyboard.Key.space:
@@ -94,10 +95,10 @@ class SO101LeaderFollower(SO101Leader):
logger.info(f"Toggled to {state}") logger.info(f"Toggled to {state}")
elif key == keyboard.Key.esc: elif key == keyboard.Key.esc:
self.keyboard_events["failure"] = True self.keyboard_events["failure"] = True
elif hasattr(key, 'char'): elif hasattr(key, "char"):
if key.char == 's': if key.char == "s":
self.keyboard_events["success"] = True self.keyboard_events["success"] = True
elif key.char == 'r': elif key.char == "r":
self.keyboard_events["rerecord"] = True self.keyboard_events["rerecord"] = True
except Exception as e: except Exception as e:
logger.error(f"Error handling key press: {e}") logger.error(f"Error handling key press: {e}")
@@ -204,9 +205,4 @@ class SO101LeaderFollower(SO101Leader):
self.is_intervening = False self.is_intervening = False
self.leader_torque_enabled = True self.leader_torque_enabled = True
self.leader_tracking_error_queue.clear() self.leader_tracking_error_queue.clear()
self.keyboard_events = { self.keyboard_events = {"intervention": False, "success": False, "failure": False, "rerecord": False}
"intervention": False,
"success": False,
"failure": False,
"rerecord": False
}