mirror of
https://github.com/huggingface/lerobot.git
synced 2026-05-23 12:40:08 +00:00
add full bimanual gravity comp
This commit is contained in:
croissant
@@ -1,73 +1,21 @@
|
||||
"""
|
||||
OpenArms Teleoperation with Gravity Compensation
|
||||
|
||||
Leader RIGHT arm: Gravity compensation (weightless, easy to move)
|
||||
Follower RIGHT arm: Mirrors leader movements
|
||||
Both LEFT arms: Free to move (disabled)
|
||||
Leader arms (both LEFT and RIGHT): Gravity compensation (weightless, easy to move)
|
||||
Follower arms (both LEFT and RIGHT): Mirror leader movements
|
||||
|
||||
The urdf file tested with this script is found in:
|
||||
https://github.com/michel-aractingi/openarm_description/blob/main/openarm_bimanual_pybullet.urdf
|
||||
Uses the URDF file from the lerobot repository.
|
||||
"""
|
||||
|
||||
import time
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import pinocchio as pin
|
||||
|
||||
from lerobot.robots.openarms.config_openarms_follower import OpenArmsFollowerConfig
|
||||
from lerobot.robots.openarms.openarms_follower import OpenArmsFollower
|
||||
from lerobot.teleoperators.openarms.config_openarms_leader import OpenArmsLeaderConfig
|
||||
from lerobot.teleoperators.openarms.openarms_leader import OpenArmsLeader
|
||||
|
||||
# Path to the URDF file
|
||||
URDF_PATH = "/home/croissant/Documents/openarm_description/openarm_bimanual_pybullet.urdf"
|
||||
|
||||
def compute_gravity_torques(leader, positions_rad):
|
||||
"""
|
||||
Compute gravity torques for all joints in the robot.
|
||||
|
||||
Args:
|
||||
leader: OpenArmsLeader instance with pin_robot set
|
||||
positions_rad: Dictionary mapping motor names (with arm prefix) to positions in radians
|
||||
|
||||
Returns:
|
||||
Dictionary mapping motor names to gravity torques in N·m
|
||||
"""
|
||||
if not hasattr(leader, "pin_robot") or leader.pin_robot is None:
|
||||
raise RuntimeError("URDF model not loaded on leader")
|
||||
|
||||
# Build position vector in the order of motors (right arm, then left arm)
|
||||
q = np.zeros(leader.pin_robot.model.nq)
|
||||
idx = 0
|
||||
|
||||
# Right arm motors
|
||||
for motor_name in leader.bus_right.motors:
|
||||
full_name = f"right_{motor_name}"
|
||||
q[idx] = positions_rad.get(full_name, 0.0)
|
||||
idx += 1
|
||||
|
||||
# Left arm motors
|
||||
for motor_name in leader.bus_left.motors:
|
||||
full_name = f"left_{motor_name}"
|
||||
q[idx] = positions_rad.get(full_name, 0.0)
|
||||
idx += 1
|
||||
|
||||
# Compute generalized gravity vector
|
||||
g = pin.computeGeneralizedGravity(leader.pin_robot.model, leader.pin_robot.data, q)
|
||||
|
||||
# Map back to motor names
|
||||
result = {}
|
||||
idx = 0
|
||||
for motor_name in leader.bus_right.motors:
|
||||
result[f"right_{motor_name}"] = float(g[idx])
|
||||
idx += 1
|
||||
for motor_name in leader.bus_left.motors:
|
||||
result[f"left_{motor_name}"] = float(g[idx])
|
||||
idx += 1
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def main():
|
||||
"""Main teleoperation loop with gravity compensation"""
|
||||
@@ -78,8 +26,8 @@ def main():
|
||||
|
||||
# Configuration
|
||||
follower_config = OpenArmsFollowerConfig(
|
||||
port_right="can0",
|
||||
port_left="can1",
|
||||
port_left="can0",
|
||||
port_right="can1",
|
||||
can_interface="socketcan",
|
||||
id="openarms_follower",
|
||||
disable_torque_on_disconnect=True,
|
||||
@@ -87,8 +35,8 @@ def main():
|
||||
)
|
||||
|
||||
leader_config = OpenArmsLeaderConfig(
|
||||
port_right="can2",
|
||||
port_left="can3",
|
||||
port_left="can2",
|
||||
port_right="can3",
|
||||
can_interface="socketcan",
|
||||
id="openarms_leader",
|
||||
manual_control=False, # Enable torque control for gravity compensation
|
||||
@@ -102,18 +50,15 @@ def main():
|
||||
follower.connect(calibrate=True)
|
||||
leader.connect(calibrate=True)
|
||||
|
||||
# Load URDF for gravity compensation
|
||||
if not os.path.exists(URDF_PATH):
|
||||
raise FileNotFoundError(f"URDF file not found at {URDF_PATH}")
|
||||
pin_robot = pin.RobotWrapper.BuildFromURDF(URDF_PATH, os.path.dirname(URDF_PATH))
|
||||
pin_robot.data = pin_robot.model.createData()
|
||||
leader.pin_robot = pin_robot
|
||||
# URDF is automatically loaded in the leader constructor
|
||||
if leader.pin_robot is None:
|
||||
raise RuntimeError("URDF model not loaded on leader. Gravity compensation not available.")
|
||||
|
||||
print("\nLeader RIGHT: G-comp | Follower RIGHT: Teleop")
|
||||
print("\nLeader BOTH arms: G-comp | Follower BOTH arms: Teleop")
|
||||
print("Press ENTER to start...")
|
||||
input()
|
||||
|
||||
# Enable motors
|
||||
# Enable motors on both leader arms for gravity compensation
|
||||
leader.bus_right.enable_torque()
|
||||
leader.bus_left.enable_torque()
|
||||
time.sleep(0.1)
|
||||
@@ -124,17 +69,12 @@ def main():
|
||||
loop_times = []
|
||||
last_print_time = time.perf_counter()
|
||||
|
||||
# Right arm joints only
|
||||
right_joints = [
|
||||
"right_joint_1",
|
||||
"right_joint_2",
|
||||
"right_joint_3",
|
||||
"right_joint_4",
|
||||
"right_joint_5",
|
||||
"right_joint_6",
|
||||
"right_joint_7",
|
||||
"right_gripper",
|
||||
]
|
||||
# All joints (both arms)
|
||||
all_joints = []
|
||||
for motor in leader.bus_right.motors:
|
||||
all_joints.append(f"right_{motor}")
|
||||
for motor in leader.bus_left.motors:
|
||||
all_joints.append(f"left_{motor}")
|
||||
|
||||
try:
|
||||
while True:
|
||||
@@ -143,6 +83,7 @@ def main():
|
||||
# Get leader state
|
||||
leader_action = leader.get_action()
|
||||
|
||||
# Extract positions in degrees
|
||||
leader_positions_deg = {}
|
||||
for motor in leader.bus_right.motors:
|
||||
key = f"right_{motor}.pos"
|
||||
@@ -154,12 +95,26 @@ def main():
|
||||
if key in leader_action:
|
||||
leader_positions_deg[f"left_{motor}"] = leader_action[key]
|
||||
|
||||
# Calculate gravity torques for leader
|
||||
# Calculate gravity torques for leader using built-in method
|
||||
leader_positions_rad = {k: np.deg2rad(v) for k, v in leader_positions_deg.items()}
|
||||
leader_torques_nm = compute_gravity_torques(leader, leader_positions_rad)
|
||||
leader_torques_nm = leader._gravity_from_q(leader_positions_rad)
|
||||
|
||||
# Apply gravity compensation to leader right arm
|
||||
# Apply gravity compensation to leader RIGHT arm (all joints except gripper)
|
||||
for motor in leader.bus_right.motors:
|
||||
if motor == "gripper":
|
||||
# Skip gripper - keep it free
|
||||
full_name = f"right_{motor}"
|
||||
position = leader_positions_deg.get(full_name, 0.0)
|
||||
leader.bus_right._mit_control(
|
||||
motor=motor,
|
||||
kp=0.0,
|
||||
kd=0.0,
|
||||
position_degrees=position,
|
||||
velocity_deg_per_sec=0.0,
|
||||
torque=0.0,
|
||||
)
|
||||
continue
|
||||
|
||||
full_name = f"right_{motor}"
|
||||
position = leader_positions_deg.get(full_name, 0.0)
|
||||
torque = leader_torques_nm.get(full_name, 0.0)
|
||||
@@ -173,10 +128,25 @@ def main():
|
||||
torque=torque,
|
||||
)
|
||||
|
||||
# Keep leader left arm free
|
||||
# Apply gravity compensation to leader LEFT arm (all joints except gripper)
|
||||
for motor in leader.bus_left.motors:
|
||||
if motor == "gripper":
|
||||
# Skip gripper - keep it free
|
||||
full_name = f"left_{motor}"
|
||||
position = leader_positions_deg.get(full_name, 0.0)
|
||||
leader.bus_left._mit_control(
|
||||
motor=motor,
|
||||
kp=0.0,
|
||||
kd=0.0,
|
||||
position_degrees=position,
|
||||
velocity_deg_per_sec=0.0,
|
||||
torque=0.0,
|
||||
)
|
||||
continue
|
||||
|
||||
full_name = f"left_{motor}"
|
||||
position = leader_positions_deg.get(full_name, 0.0)
|
||||
torque = leader_torques_nm.get(full_name, 0.0)
|
||||
|
||||
leader.bus_left._mit_control(
|
||||
motor=motor,
|
||||
@@ -184,12 +154,12 @@ def main():
|
||||
kd=0.0,
|
||||
position_degrees=position,
|
||||
velocity_deg_per_sec=0.0,
|
||||
torque=0.0,
|
||||
torque=torque,
|
||||
)
|
||||
|
||||
# Send leader positions to follower right arm
|
||||
# Send leader positions to follower (both arms)
|
||||
follower_action = {}
|
||||
for joint in right_joints:
|
||||
for joint in all_joints:
|
||||
pos_key = f"{joint}.pos"
|
||||
if pos_key in leader_action:
|
||||
follower_action[pos_key] = leader_action[pos_key]
|
||||
@@ -206,10 +176,8 @@ def main():
|
||||
if loop_times:
|
||||
avg_time = sum(loop_times) / len(loop_times)
|
||||
current_hz = 1.0 / avg_time if avg_time > 0 else 0
|
||||
sample_pos = leader_positions_deg.get("right_joint_2", 0.0)
|
||||
sample_torque = leader_torques_nm.get("right_joint_2", 0.0)
|
||||
|
||||
print(f"[{current_hz:.1f} Hz] J2: {sample_pos:5.1f}° | Torque: {sample_torque:5.2f} N·m")
|
||||
print(f"{current_hz:.1f} Hz ({avg_time*1000:.1f} ms)")
|
||||
|
||||
loop_times = []
|
||||
last_print_time = loop_end
|
||||
|
||||
Reference in New Issue
Block a user