mirror of
https://github.com/huggingface/lerobot.git
synced 2026-05-16 17:20:05 +00:00
Add code for interpolation, pid tuning, raw action scaling and vel FF
This commit is contained in:
Pepijn
@@ -0,0 +1,587 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""
|
||||
OpenArms Policy Evaluation with Interpolation
|
||||
|
||||
Evaluates a trained policy with smooth action interpolation:
|
||||
- Decoupled camera capture (CAMERA_FPS) from robot control (ROBOT_FPS)
|
||||
- Speed multiplier to execute actions faster than training
|
||||
- Velocity feedforward for smoother tracking
|
||||
- Adjustable PID gains
|
||||
|
||||
Example usage:
|
||||
python examples/openarms/evaluate_interpolation.py
|
||||
"""
|
||||
|
||||
import time
|
||||
from collections import deque
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig
|
||||
from lerobot.configs.policies import PreTrainedConfig
|
||||
from lerobot.datasets.lerobot_dataset import LeRobotDataset
|
||||
from lerobot.datasets.pipeline_features import aggregate_pipeline_dataset_features, create_initial_features
|
||||
from lerobot.datasets.utils import combine_feature_dicts
|
||||
from lerobot.policies.factory import make_policy, make_pre_post_processors
|
||||
from lerobot.processor import make_default_processors
|
||||
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
|
||||
from lerobot.utils.control_utils import init_keyboard_listener, predict_action
|
||||
from lerobot.utils.utils import log_say, get_safe_torch_device
|
||||
from lerobot.utils.visualization_utils import init_rerun
|
||||
|
||||
|
||||
# ======================== MODEL & TASK CONFIG ========================
|
||||
HF_MODEL_ID = "lerobot-data-collection/three-folds-pi0" # TODO: Replace with your trained model
|
||||
HF_EVAL_DATASET_ID = "lerobot-data-collection/three-folds-pi0_eval_interp" # TODO: Replace
|
||||
TASK_DESCRIPTION = "three-folds-dataset" # TODO: Replace with your task
|
||||
|
||||
# ======================== TIMING CONFIG ========================
|
||||
CAMERA_FPS = 30 # Camera hardware limit (fixed)
|
||||
POLICY_FPS = 30 # What the policy was trained with
|
||||
SPEED_MULTIPLIER = 1.2 # Execute actions faster (1.0 = normal, 1.2 = 20% faster)
|
||||
ROBOT_FPS = 50 # Robot command rate (higher = smoother interpolation)
|
||||
|
||||
# Derived values
|
||||
EFFECTIVE_POLICY_FPS = int(POLICY_FPS * SPEED_MULTIPLIER) # How fast we consume actions (36Hz at 1.2x)
|
||||
|
||||
NUM_EPISODES = 1
|
||||
EPISODE_TIME_SEC = 300
|
||||
RESET_TIME_SEC = 60
|
||||
|
||||
# ======================== PID TUNING ========================
|
||||
# Set to None to use robot config defaults
|
||||
CUSTOM_KP_SCALE = 0.7 # Scale factor for position gain (0.5-1.0, lower = smoother)
|
||||
CUSTOM_KD_SCALE = 1.3 # Scale factor for damping gain (1.0-2.0, higher = less overshoot)
|
||||
USE_VELOCITY_FEEDFORWARD = True # Enable velocity feedforward for smoother tracking
|
||||
|
||||
# ======================== ROBOT CONFIG ========================
|
||||
FOLLOWER_LEFT_PORT = "can0"
|
||||
FOLLOWER_RIGHT_PORT = "can1"
|
||||
|
||||
USE_LEADER_FOR_RESETS = True
|
||||
LEADER_LEFT_PORT = "can2"
|
||||
LEADER_RIGHT_PORT = "can3"
|
||||
|
||||
# Camera config uses CAMERA_FPS (hardware limit)
|
||||
CAMERA_CONFIG = {
|
||||
"left_wrist": OpenCVCameraConfig(index_or_path="/dev/video5", width=640, height=480, fps=CAMERA_FPS),
|
||||
"right_wrist": OpenCVCameraConfig(index_or_path="/dev/video1", width=640, height=480, fps=CAMERA_FPS),
|
||||
"base": OpenCVCameraConfig(index_or_path="/dev/video3", width=640, height=480, fps=CAMERA_FPS),
|
||||
}
|
||||
|
||||
|
||||
class ActionInterpolator:
|
||||
"""Interpolate between policy actions for smoother robot control with velocity estimation."""
|
||||
|
||||
def __init__(self, effective_policy_fps: int, robot_fps: int):
|
||||
self.effective_policy_fps = effective_policy_fps
|
||||
self.robot_fps = robot_fps
|
||||
self.substeps_per_policy_step = robot_fps / effective_policy_fps
|
||||
self.prev_action: dict | None = None
|
||||
self.curr_action: dict | None = None
|
||||
self.substep = 0
|
||||
self.last_interpolated: dict | None = None
|
||||
|
||||
def update(self, new_action: dict) -> None:
|
||||
self.prev_action = self.curr_action
|
||||
self.curr_action = new_action
|
||||
self.substep = 0
|
||||
|
||||
def get_interpolated_action(self) -> tuple[dict | None, dict | None]:
|
||||
"""Returns (interpolated_position, estimated_velocity_deg_per_sec)"""
|
||||
if self.curr_action is None:
|
||||
return None, None
|
||||
if self.prev_action is None:
|
||||
self.last_interpolated = self.curr_action.copy()
|
||||
return self.curr_action, {k: 0.0 for k in self.curr_action}
|
||||
|
||||
t = min(self.substep / self.substeps_per_policy_step, 1.0)
|
||||
self.substep += 1
|
||||
|
||||
interpolated = {}
|
||||
velocity = {}
|
||||
dt = 1.0 / self.robot_fps
|
||||
|
||||
for key in self.curr_action:
|
||||
prev = self.prev_action.get(key, self.curr_action[key])
|
||||
curr = self.curr_action[key]
|
||||
interpolated[key] = prev * (1 - t) + curr * t
|
||||
|
||||
if self.last_interpolated is not None and key in self.last_interpolated:
|
||||
velocity[key] = (interpolated[key] - self.last_interpolated[key]) / dt
|
||||
else:
|
||||
velocity[key] = (curr - prev) * self.effective_policy_fps
|
||||
|
||||
self.last_interpolated = interpolated.copy()
|
||||
return interpolated, velocity
|
||||
|
||||
def reset(self):
|
||||
self.prev_action = None
|
||||
self.curr_action = None
|
||||
self.substep = 0
|
||||
self.last_interpolated = None
|
||||
|
||||
|
||||
class HzTracker:
|
||||
"""Track and display actual loop frequency."""
|
||||
|
||||
def __init__(self, name: str = "Robot", window_size: int = 100, print_interval: float = 2.0):
|
||||
self.name = name
|
||||
self.timestamps = deque(maxlen=window_size)
|
||||
self.last_print_time = 0
|
||||
self.print_interval = print_interval
|
||||
|
||||
def tick(self) -> float | None:
|
||||
now = time.perf_counter()
|
||||
self.timestamps.append(now)
|
||||
|
||||
if len(self.timestamps) < 2:
|
||||
return None
|
||||
|
||||
hz = (len(self.timestamps) - 1) / (self.timestamps[-1] - self.timestamps[0])
|
||||
|
||||
if now - self.last_print_time >= self.print_interval:
|
||||
print(f"{self.name} Hz: {hz:.1f}")
|
||||
self.last_print_time = now
|
||||
|
||||
return hz
|
||||
|
||||
def get_avg_hz(self) -> float | None:
|
||||
if len(self.timestamps) < 2:
|
||||
return None
|
||||
return (len(self.timestamps) - 1) / (self.timestamps[-1] - self.timestamps[0])
|
||||
|
||||
def reset(self):
|
||||
self.timestamps.clear()
|
||||
self.last_print_time = 0
|
||||
|
||||
|
||||
def interpolated_eval_loop(
|
||||
robot,
|
||||
policy,
|
||||
preprocessor,
|
||||
postprocessor,
|
||||
robot_observation_processor,
|
||||
robot_action_processor,
|
||||
dataset,
|
||||
events,
|
||||
interpolator: ActionInterpolator,
|
||||
robot_hz_tracker: HzTracker,
|
||||
camera_fps: int,
|
||||
effective_policy_fps: int,
|
||||
robot_fps: int,
|
||||
control_time_s: float,
|
||||
task: str,
|
||||
kp_scale: float | None = None,
|
||||
kd_scale: float | None = None,
|
||||
use_velocity_ff: bool = False,
|
||||
):
|
||||
"""
|
||||
Run evaluation with decoupled camera and robot control:
|
||||
- Camera captures at camera_fps (hardware limit)
|
||||
- Policy inference runs when new camera frame is available
|
||||
- Actions are consumed at effective_policy_fps (sped up by SPEED_MULTIPLIER)
|
||||
- Robot receives interpolated commands at robot_fps (smoothest)
|
||||
"""
|
||||
from lerobot.scripts.lerobot_record import build_dataset_frame, make_robot_action
|
||||
from lerobot.utils.visualization_utils import log_rerun_data
|
||||
|
||||
camera_dt = 1.0 / camera_fps
|
||||
policy_dt = 1.0 / effective_policy_fps
|
||||
robot_dt = 1.0 / robot_fps
|
||||
|
||||
interpolator.reset()
|
||||
robot_hz_tracker.reset()
|
||||
policy.reset()
|
||||
|
||||
# Build custom gains if scaling is enabled
|
||||
custom_kp = None
|
||||
custom_kd = None
|
||||
if kp_scale is not None or kd_scale is not None:
|
||||
custom_kp = {}
|
||||
custom_kd = {}
|
||||
for arm in ["right", "left"]:
|
||||
bus = robot.bus_right if arm == "right" else robot.bus_left
|
||||
for i, motor_name in enumerate(bus.motors):
|
||||
full_name = f"{arm}_{motor_name}"
|
||||
default_kp = robot.config.position_kp[i] if isinstance(robot.config.position_kp, list) else robot.config.position_kp
|
||||
default_kd = robot.config.position_kd[i] if isinstance(robot.config.position_kd, list) else robot.config.position_kd
|
||||
custom_kp[full_name] = default_kp * (kp_scale or 1.0)
|
||||
custom_kd[full_name] = default_kd * (kd_scale or 1.0)
|
||||
print(f"Custom gains: kp_scale={kp_scale}, kd_scale={kd_scale}")
|
||||
|
||||
if use_velocity_ff:
|
||||
print("Velocity feedforward: enabled")
|
||||
|
||||
last_camera_time = -camera_dt
|
||||
last_policy_action_time = -policy_dt
|
||||
cached_observation = None
|
||||
cached_robot_action = None
|
||||
|
||||
start_time = time.perf_counter()
|
||||
|
||||
print(f"\nStarting interpolated eval loop:")
|
||||
print(f" Camera: {camera_fps}Hz | Policy actions consumed: {effective_policy_fps}Hz | Robot: {robot_fps}Hz")
|
||||
|
||||
while time.perf_counter() - start_time < control_time_s:
|
||||
if events["exit_early"] or events["stop_recording"]:
|
||||
break
|
||||
|
||||
loop_start = time.perf_counter()
|
||||
elapsed = loop_start - start_time
|
||||
|
||||
# === CAMERA CAPTURE (at camera_fps, decoupled from robot) ===
|
||||
if elapsed - last_camera_time >= camera_dt:
|
||||
obs = robot.get_observation()
|
||||
obs_processed = robot_observation_processor(obs)
|
||||
observation_frame = build_dataset_frame(dataset.features, obs_processed, prefix="observation")
|
||||
|
||||
# Run policy inference with fresh observation
|
||||
action_values = predict_action(
|
||||
observation=observation_frame,
|
||||
policy=policy,
|
||||
device=get_safe_torch_device(policy.config.device),
|
||||
preprocessor=preprocessor,
|
||||
postprocessor=postprocessor,
|
||||
use_amp=policy.config.use_amp,
|
||||
task=task,
|
||||
robot_type=robot.robot_type,
|
||||
)
|
||||
|
||||
act_processed = make_robot_action(action_values, dataset.features)
|
||||
cached_robot_action = robot_action_processor((act_processed, obs))
|
||||
cached_observation = (obs_processed, observation_frame, act_processed)
|
||||
|
||||
last_camera_time = elapsed
|
||||
|
||||
# === ACTION UPDATE (at effective_policy_fps, faster than camera if speed > 1) ===
|
||||
if elapsed - last_policy_action_time >= policy_dt and cached_robot_action is not None:
|
||||
interpolator.update(cached_robot_action)
|
||||
last_policy_action_time = elapsed
|
||||
|
||||
# Save to dataset at effective policy rate
|
||||
if dataset is not None and cached_observation is not None:
|
||||
obs_processed, observation_frame, act_processed = cached_observation
|
||||
action_frame = build_dataset_frame(dataset.features, act_processed, prefix="action")
|
||||
frame = {**observation_frame, **action_frame, "task": task}
|
||||
dataset.add_frame(frame)
|
||||
log_rerun_data(observation=obs_processed, action=act_processed)
|
||||
|
||||
# === ROBOT COMMAND (at robot_fps, highest rate for smoothness) ===
|
||||
smooth_action, velocity = interpolator.get_interpolated_action()
|
||||
if smooth_action is not None:
|
||||
vel_ff = velocity if use_velocity_ff else None
|
||||
robot.send_action(smooth_action, custom_kp=custom_kp, custom_kd=custom_kd, velocity_feedforward=vel_ff)
|
||||
|
||||
robot_hz_tracker.tick()
|
||||
|
||||
# Maintain robot control rate
|
||||
sleep_time = robot_dt - (time.perf_counter() - loop_start)
|
||||
if sleep_time > 0:
|
||||
time.sleep(sleep_time)
|
||||
|
||||
# Print final stats
|
||||
robot_hz = robot_hz_tracker.get_avg_hz()
|
||||
if robot_hz:
|
||||
print(f"\nFinal average robot Hz: {robot_hz:.1f}")
|
||||
|
||||
|
||||
def main():
|
||||
"""Main evaluation function."""
|
||||
print("=" * 60)
|
||||
print("OpenArms Policy Evaluation with Interpolation")
|
||||
print("=" * 60)
|
||||
print(f"\nModel: {HF_MODEL_ID}")
|
||||
print(f"Dataset: {HF_EVAL_DATASET_ID}")
|
||||
print(f"Task: {TASK_DESCRIPTION}")
|
||||
print(f"\n--- Timing ---")
|
||||
print(f"Camera FPS: {CAMERA_FPS} (hardware limit)")
|
||||
print(f"Policy trained at: {POLICY_FPS}Hz")
|
||||
print(f"Speed multiplier: {SPEED_MULTIPLIER}x")
|
||||
print(f"Effective policy FPS: {EFFECTIVE_POLICY_FPS}Hz (actions consumed)")
|
||||
print(f"Robot FPS: {ROBOT_FPS}Hz (interpolated commands)")
|
||||
print(f"\n--- PID Tuning ---")
|
||||
print(f"KP scale: {CUSTOM_KP_SCALE}")
|
||||
print(f"KD scale: {CUSTOM_KD_SCALE}")
|
||||
print(f"Velocity feedforward: {USE_VELOCITY_FEEDFORWARD}")
|
||||
print(f"\n--- Episodes ---")
|
||||
print(f"Episodes: {NUM_EPISODES}")
|
||||
print(f"Duration: {EPISODE_TIME_SEC}s per episode")
|
||||
print(f"Reset time: {RESET_TIME_SEC}s")
|
||||
print(f"Leader for resets: {USE_LEADER_FOR_RESETS}")
|
||||
print("=" * 60)
|
||||
|
||||
follower_config = OpenArmsFollowerConfig(
|
||||
port_left=FOLLOWER_LEFT_PORT,
|
||||
port_right=FOLLOWER_RIGHT_PORT,
|
||||
can_interface="socketcan",
|
||||
id="openarms_follower",
|
||||
disable_torque_on_disconnect=True,
|
||||
max_relative_target=10.0,
|
||||
cameras=CAMERA_CONFIG,
|
||||
)
|
||||
|
||||
follower = OpenArmsFollower(follower_config)
|
||||
follower.connect(calibrate=False)
|
||||
|
||||
if not follower.is_connected:
|
||||
raise RuntimeError("Follower robot failed to connect!")
|
||||
|
||||
leader = None
|
||||
if USE_LEADER_FOR_RESETS:
|
||||
leader_config = OpenArmsLeaderConfig(
|
||||
port_left=LEADER_LEFT_PORT,
|
||||
port_right=LEADER_RIGHT_PORT,
|
||||
can_interface="socketcan",
|
||||
id="openarms_leader",
|
||||
manual_control=False,
|
||||
)
|
||||
|
||||
leader = OpenArmsLeader(leader_config)
|
||||
leader.connect(calibrate=False)
|
||||
|
||||
if not leader.is_connected:
|
||||
raise RuntimeError("Leader robot failed to connect!")
|
||||
|
||||
if leader.pin_robot is not None:
|
||||
leader.bus_right.enable_torque()
|
||||
leader.bus_left.enable_torque()
|
||||
time.sleep(0.1)
|
||||
print(f"Leader connected with gravity compensation")
|
||||
else:
|
||||
print(f"Leader connected (no gravity compensation)")
|
||||
|
||||
teleop_action_processor, robot_action_processor, robot_observation_processor = make_default_processors()
|
||||
|
||||
action_features_hw = {}
|
||||
for key, value in follower.action_features.items():
|
||||
if key.endswith(".pos"):
|
||||
action_features_hw[key] = value
|
||||
|
||||
dataset_features = combine_feature_dicts(
|
||||
aggregate_pipeline_dataset_features(
|
||||
pipeline=teleop_action_processor,
|
||||
initial_features=create_initial_features(action=action_features_hw),
|
||||
use_videos=True,
|
||||
),
|
||||
aggregate_pipeline_dataset_features(
|
||||
pipeline=robot_observation_processor,
|
||||
initial_features=create_initial_features(observation=follower.observation_features),
|
||||
use_videos=True,
|
||||
),
|
||||
)
|
||||
|
||||
dataset_path = Path.home() / ".cache" / "huggingface" / "lerobot" / HF_EVAL_DATASET_ID
|
||||
if dataset_path.exists():
|
||||
print(f"\nDataset exists at: {dataset_path}")
|
||||
choice = input("Continue and append? (y/n): ").strip().lower()
|
||||
if choice != 'y':
|
||||
print("Aborting.")
|
||||
follower.disconnect()
|
||||
if leader:
|
||||
leader.disconnect()
|
||||
return
|
||||
|
||||
# Dataset uses effective policy FPS (sped up rate)
|
||||
dataset = LeRobotDataset.create(
|
||||
repo_id=HF_EVAL_DATASET_ID,
|
||||
fps=EFFECTIVE_POLICY_FPS,
|
||||
features=dataset_features,
|
||||
robot_type=follower.name,
|
||||
use_videos=True,
|
||||
image_writer_processes=0,
|
||||
image_writer_threads=12,
|
||||
)
|
||||
|
||||
policy_config = PreTrainedConfig.from_pretrained(HF_MODEL_ID)
|
||||
policy_config.pretrained_path = HF_MODEL_ID
|
||||
policy = make_policy(policy_config, ds_meta=dataset.meta)
|
||||
|
||||
preprocessor, postprocessor = make_pre_post_processors(
|
||||
policy_cfg=policy.config,
|
||||
pretrained_path=HF_MODEL_ID,
|
||||
dataset_stats=dataset.meta.stats,
|
||||
preprocessor_overrides={
|
||||
"device_processor": {"device": str(policy.config.device)}
|
||||
},
|
||||
)
|
||||
|
||||
print(f"\nRunning evaluation...")
|
||||
listener, events = init_keyboard_listener()
|
||||
init_rerun(session_name="openarms_evaluation_interp")
|
||||
|
||||
interpolator = ActionInterpolator(effective_policy_fps=EFFECTIVE_POLICY_FPS, robot_fps=ROBOT_FPS)
|
||||
robot_hz_tracker = HzTracker(name="Robot", window_size=100, print_interval=2.0)
|
||||
|
||||
episode_idx = 0
|
||||
|
||||
try:
|
||||
while episode_idx < NUM_EPISODES and not events["stop_recording"]:
|
||||
log_say(f"Evaluating episode {episode_idx + 1} of {NUM_EPISODES}")
|
||||
print(f"\n--- Episode {episode_idx + 1}/{NUM_EPISODES} ---")
|
||||
|
||||
interpolated_eval_loop(
|
||||
robot=follower,
|
||||
policy=policy,
|
||||
preprocessor=preprocessor,
|
||||
postprocessor=postprocessor,
|
||||
robot_observation_processor=robot_observation_processor,
|
||||
robot_action_processor=robot_action_processor,
|
||||
dataset=dataset,
|
||||
events=events,
|
||||
interpolator=interpolator,
|
||||
robot_hz_tracker=robot_hz_tracker,
|
||||
camera_fps=CAMERA_FPS,
|
||||
effective_policy_fps=EFFECTIVE_POLICY_FPS,
|
||||
robot_fps=ROBOT_FPS,
|
||||
control_time_s=EPISODE_TIME_SEC,
|
||||
task=TASK_DESCRIPTION,
|
||||
kp_scale=CUSTOM_KP_SCALE,
|
||||
kd_scale=CUSTOM_KD_SCALE,
|
||||
use_velocity_ff=USE_VELOCITY_FEEDFORWARD,
|
||||
)
|
||||
|
||||
if events["rerecord_episode"]:
|
||||
log_say("Re-recording episode")
|
||||
events["rerecord_episode"] = False
|
||||
events["exit_early"] = False
|
||||
dataset.clear_episode_buffer()
|
||||
continue
|
||||
|
||||
if dataset.episode_buffer is not None and dataset.episode_buffer.get("size", 0) > 0:
|
||||
print(f"Saving episode ({dataset.episode_buffer['size']} frames)...")
|
||||
dataset.save_episode()
|
||||
episode_idx += 1
|
||||
|
||||
if not events["stop_recording"] and episode_idx < NUM_EPISODES:
|
||||
if USE_LEADER_FOR_RESETS and leader:
|
||||
log_say("Reset the environment using leader arms")
|
||||
print(f"\nManual reset ({RESET_TIME_SEC}s)...")
|
||||
|
||||
dt = 1 / CAMERA_FPS
|
||||
reset_start_time = time.perf_counter()
|
||||
|
||||
while time.perf_counter() - reset_start_time < RESET_TIME_SEC:
|
||||
if events["exit_early"] or events["stop_recording"]:
|
||||
break
|
||||
|
||||
loop_start = time.perf_counter()
|
||||
leader_action = leader.get_action()
|
||||
|
||||
leader_positions_deg = {}
|
||||
leader_velocities_deg_per_sec = {}
|
||||
|
||||
for motor in leader.bus_right.motors:
|
||||
pos_key = f"right_{motor}.pos"
|
||||
vel_key = f"right_{motor}.vel"
|
||||
if pos_key in leader_action:
|
||||
leader_positions_deg[f"right_{motor}"] = leader_action[pos_key]
|
||||
if vel_key in leader_action:
|
||||
leader_velocities_deg_per_sec[f"right_{motor}"] = leader_action[vel_key]
|
||||
|
||||
for motor in leader.bus_left.motors:
|
||||
pos_key = f"left_{motor}.pos"
|
||||
vel_key = f"left_{motor}.vel"
|
||||
if pos_key in leader_action:
|
||||
leader_positions_deg[f"left_{motor}"] = leader_action[pos_key]
|
||||
if vel_key in leader_action:
|
||||
leader_velocities_deg_per_sec[f"left_{motor}"] = leader_action[vel_key]
|
||||
|
||||
leader_positions_rad = {k: np.deg2rad(v) for k, v in leader_positions_deg.items()}
|
||||
leader_gravity_torques_nm = leader._gravity_from_q(leader_positions_rad)
|
||||
|
||||
leader_velocities_rad_per_sec = {k: np.deg2rad(v) for k, v in leader_velocities_deg_per_sec.items()}
|
||||
leader_friction_torques_nm = leader._friction_from_velocity(
|
||||
leader_velocities_rad_per_sec, friction_scale=1.0
|
||||
)
|
||||
|
||||
leader_total_torques_nm = {}
|
||||
for motor_name in leader_gravity_torques_nm:
|
||||
gravity = leader_gravity_torques_nm.get(motor_name, 0.0)
|
||||
friction = leader_friction_torques_nm.get(motor_name, 0.0)
|
||||
leader_total_torques_nm[motor_name] = gravity + friction
|
||||
|
||||
for motor in leader.bus_right.motors:
|
||||
full_name = f"right_{motor}"
|
||||
position = leader_positions_deg.get(full_name, 0.0)
|
||||
torque = leader_total_torques_nm.get(full_name, 0.0)
|
||||
kd = leader.get_damping_kd(motor)
|
||||
leader.bus_right._mit_control(
|
||||
motor=motor, kp=0.0, kd=kd,
|
||||
position_degrees=position, velocity_deg_per_sec=0.0, torque=torque,
|
||||
)
|
||||
|
||||
for motor in leader.bus_left.motors:
|
||||
full_name = f"left_{motor}"
|
||||
position = leader_positions_deg.get(full_name, 0.0)
|
||||
torque = leader_total_torques_nm.get(full_name, 0.0)
|
||||
kd = leader.get_damping_kd(motor)
|
||||
leader.bus_left._mit_control(
|
||||
motor=motor, kp=0.0, kd=kd,
|
||||
position_degrees=position, velocity_deg_per_sec=0.0, torque=torque,
|
||||
)
|
||||
|
||||
follower_action = {}
|
||||
for joint in leader_positions_deg.keys():
|
||||
pos_key = f"{joint}.pos"
|
||||
if pos_key in leader_action:
|
||||
follower_action[pos_key] = leader_action[pos_key]
|
||||
|
||||
if follower_action:
|
||||
follower.send_action(follower_action)
|
||||
|
||||
loop_duration = time.perf_counter() - loop_start
|
||||
sleep_time = dt - loop_duration
|
||||
if sleep_time > 0:
|
||||
time.sleep(sleep_time)
|
||||
|
||||
print("Reset complete")
|
||||
else:
|
||||
log_say("Waiting for manual reset")
|
||||
input("Press ENTER when ready...")
|
||||
|
||||
print(f"\nEvaluation complete! {episode_idx} episodes recorded")
|
||||
log_say("Evaluation complete", blocking=True)
|
||||
|
||||
except KeyboardInterrupt:
|
||||
print("\n\nInterrupted by user")
|
||||
|
||||
finally:
|
||||
if leader:
|
||||
leader.bus_right.disable_torque()
|
||||
leader.bus_left.disable_torque()
|
||||
time.sleep(0.1)
|
||||
leader.disconnect()
|
||||
|
||||
follower.disconnect()
|
||||
|
||||
if listener is not None:
|
||||
listener.stop()
|
||||
|
||||
dataset.finalize()
|
||||
print("\nUploading to Hugging Face Hub...")
|
||||
dataset.push_to_hub(private=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
Reference in New Issue
Block a user