From 7f7b45cfbbf42eeabb154932081083e877761c62 Mon Sep 17 00:00:00 2001 From: Pepijn Date: Fri, 4 Jul 2025 17:33:25 +0200 Subject: [PATCH] Add rerun vis --- src/lerobot/bi_teleoperate.py | 68 ++++++++++++++++++++++++++++++----- 1 file changed, 60 insertions(+), 8 deletions(-) diff --git a/src/lerobot/bi_teleoperate.py b/src/lerobot/bi_teleoperate.py index a405b067a..766d822ef 100644 --- a/src/lerobot/bi_teleoperate.py +++ b/src/lerobot/bi_teleoperate.py @@ -8,9 +8,11 @@ import numpy as np from lerobot.robots.so101_follower_torque.config_so101_follower_t import SO101FollowerTConfig from lerobot.robots.so101_follower_torque.so101_follower_t import SO101FollowerT from lerobot.utils.robot_utils import busy_wait +from lerobot.utils.visualization_utils import _init_rerun, log_rerun_data FRQ = 300 PRINT_HZ = 10 +RERUN_HZ = 100 ESC_CLR_EOL = "\x1b[K" CURSOR_UP = "\x1b[F" @@ -27,12 +29,12 @@ def to_cmd(tau_vec, joints): follower_cfg = SO101FollowerTConfig( port="/dev/tty.usbmodem58760431551", - id="my_awesome_follower_arm", + id="follower_arm_t", ) leader_cfg = SO101FollowerTConfig( port="/dev/tty.usbmodem58760428721", - id="my_awesome_leader_arm", + id="leader_arm_t", ) follower = SO101FollowerT(follower_cfg) @@ -43,11 +45,18 @@ leader.connect() exec_f = ThreadPoolExecutor(max_workers=1) exec_l = ThreadPoolExecutor(max_workers=1) +# Initialize Rerun for visualization +_init_rerun("bilateral_teleoperation") + print("Starting 4-channel bilateral teleoperation") first_print = True loop_count = 0 tic_prev = time.perf_counter() +# Initialize previous positions for velocity calculation +pos_f_prev = None +pos_l_prev = None + while True: tic = time.perf_counter() @@ -64,17 +73,32 @@ while True: tau_cmd_f, tau_cmd_l = [], [] debug_info_f, debug_info_l = {}, {} + # Collect torques with deterministic components removed for model + tau_interaction_l = {} + tau_interaction_f = {} + # Collect data for all motors pos_f = {j: obs_f[f"{j}.pos"] for j in follower.bus.motors} - vel_f = {j: obs_f[f"{j}.vel"] for j in follower.bus.motors} acc_f = {j: obs_f[f"{j}.acc"] for j in follower.bus.motors} tau_meas_f = {j: obs_f[f"{j}.tau_meas"] for j in follower.bus.motors} pos_l = {j: obs_l[f"{j}.pos"] for j in leader.bus.motors} - vel_l = {j: obs_l[f"{j}.vel"] for j in leader.bus.motors} acc_l = {j: obs_l[f"{j}.acc"] for j in leader.bus.motors} tau_meas_l = {j: obs_l[f"{j}.tau_meas"] for j in leader.bus.motors} + # Calculate velocities from position differences + if pos_f_prev is not None and pos_l_prev is not None: + vel_f = {j: (pos_f[j] - pos_f_prev[j]) / dt for j in follower.bus.motors} + vel_l = {j: (pos_l[j] - pos_l_prev[j]) / dt for j in leader.bus.motors} + else: + # First iteration: initialize with zeros + vel_f = dict.fromkeys(follower.bus.motors, 0.0) + vel_l = dict.fromkeys(leader.bus.motors, 0.0) + + # Update previous positions for next iteration + pos_f_prev = pos_f.copy() + pos_l_prev = pos_l.copy() + # Compute reaction torques using model-based approach tau_reaction_f = follower._compute_model_based_disturbance(pos_f, vel_f, acc_f, tau_meas_f) tau_reaction_l = leader._compute_model_based_disturbance(pos_l, vel_l, acc_l, tau_meas_l) @@ -87,8 +111,8 @@ while True: inertia_f = follower._inertia_from_q_dq(pos_f, vel_f, acc_f) inertia_l = leader._inertia_from_q_dq(pos_l, vel_l, acc_l) - Kp = 5.0 # Position gain - Kd = 0.1 # Velocity gain + Kp = 8.0 # Position gain + Kd = 0.5 # Velocity gain Kf = 0.05 # Force reflection gain # Friction compensation parameters (feedforward) @@ -151,6 +175,12 @@ while True: tau_cmd_f.append(tau_ref_f) tau_cmd_l.append(tau_ref_l) + # Store interaction torques (cmd minus deterministic components) for model + tau_interaction_l[j] = tau_ref_l - grav_l[j] - inertia_l[j] - tau_friction_l + tau_interaction_f[j] = -( + tau_meas_f[j] - grav_f[j] - inertia_f[j] - tau_friction_f + ) # Invert to match leader + # Store debug info debug_info_f[j] = { "τ_measured": tau_meas_f[j], @@ -182,8 +212,30 @@ while True: } # Send torques to both arms - follower.send_action(to_cmd(np.array(tau_cmd_f), list(follower.bus.motors.keys()))) - leader.send_action(to_cmd(np.array(tau_cmd_l), list(leader.bus.motors.keys()))) + follower_action = to_cmd(np.array(tau_cmd_f), list(follower.bus.motors.keys())) + leader_action = to_cmd(np.array(tau_cmd_l), list(leader.bus.motors.keys())) + + follower.send_action(follower_action) + leader.send_action(leader_action) + + # Create structured observation and action for rerun + # Observation: follower side only (θ_f, ω_f, τ_ext) + observation = { + "follower_joint_angles": pos_f, # θ_f: current angles + "follower_angular_velocities": vel_f, # ω_f: current velocities + "follower_external_torques": tau_interaction_f, # τ_ext: measured minus deterministic components + } + + # Action: leader targets (θ_leader[τ], ω_leader[τ], τ_leader[τ]) + action = { + "leader_target_angles": pos_l, # θ_leader[τ]: absolute target angles + "leader_target_velocities": vel_l, # ω_leader[τ]: absolute target velocities + "leader_interaction_torques": tau_interaction_l, # τ_leader[τ]: cmd minus deterministic components + } + + # Log data for visualization (100 Hz) + if loop_count % (FRQ // RERUN_HZ) == 0: + log_rerun_data(observation, action) # Console diagnostics (10 Hz) loop_count += 1