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")
@@ -53,10 +53,7 @@ class LeaderFollowerProcessor:
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
@@ -80,9 +77,7 @@ class LeaderFollowerProcessor:
# 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:
@@ -91,9 +86,7 @@ class LeaderFollowerProcessor:
# 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
+11 -5
View File
@@ -467,11 +467,13 @@ def make_processors(
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(
@@ -85,6 +85,7 @@ class SO101LeaderFollower(SO101Leader):
def _start_keyboard_listener(self):
"""Start keyboard listener thread for intervention control."""
def on_press(key):
try:
if key == keyboard.Key.space:
@@ -94,10 +95,10 @@ 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}")
@@ -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}