mirror of
https://github.com/huggingface/lerobot.git
synced 2026-05-21 03:30:10 +00:00
Martino Russi
448 lines
17 KiB
Python
448 lines
17 KiB
Python
#!/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.
|
|
"""
|
|
Example: Unitree RL 12-DOF Legs-Only Locomotion (TorchScript)
|
|
|
|
This example demonstrates loading a 12-DOF legs-only locomotion policy
|
|
(TorchScript .pt format) and running it on the Unitree G1 robot.
|
|
|
|
Key characteristics:
|
|
- Single TorchScript policy (.pt)
|
|
- 47D observations, 12D actions (legs only)
|
|
- Phase-based gait timing
|
|
- Arms and waist held at fixed positions
|
|
"""
|
|
|
|
import argparse
|
|
import logging
|
|
import threading
|
|
import time
|
|
|
|
import numpy as np
|
|
import torch
|
|
from huggingface_hub import hf_hub_download
|
|
from scipy.spatial.transform import Rotation as R
|
|
|
|
from lerobot.robots.unitree_g1.config_unitree_g1 import UnitreeG1Config
|
|
from lerobot.robots.unitree_g1.unitree_g1 import UnitreeG1
|
|
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# 12-DOF leg joint configuration
|
|
# Joint order: [L_hip_pitch, L_hip_roll, L_hip_yaw, L_knee, L_ankle_pitch, L_ankle_roll,
|
|
# R_hip_pitch, R_hip_roll, R_hip_yaw, R_knee, R_ankle_pitch, R_ankle_roll]
|
|
LEG_JOINT_INDICES = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
|
|
|
|
# Default leg angles for standing
|
|
DEFAULT_LEG_ANGLES = np.array([
|
|
-0.1, 0.0, 0.0, 0.3, -0.2, 0.0, # left leg
|
|
-0.1, 0.0, 0.0, 0.3, -0.2, 0.0, # right leg
|
|
], dtype=np.float32)
|
|
|
|
# KP/KD for leg joints
|
|
LEG_KPS = np.array([150, 150, 150, 300, 40, 40, 150, 150, 150, 300, 40, 40], dtype=np.float32)
|
|
LEG_KDS = np.array([6, 6, 6, 4, 2, 2, 6, 6, 6, 4, 2, 2], dtype=np.float32)
|
|
|
|
# Waist configuration (held at zero)
|
|
WAIST_JOINT_INDICES = [12, 13, 14] # yaw, roll, pitch
|
|
WAIST_KPS = np.array([250, 250, 250], dtype=np.float32)
|
|
WAIST_KDS = np.array([5, 5, 5], dtype=np.float32)
|
|
|
|
# Arm configuration (indices 15-28, held at initial position)
|
|
ARM_JOINT_INDICES = list(range(15, 29))
|
|
ARM_KPS = np.array([80, 80, 80, 80, 40, 40, 40, # left arm (shoulder + wrist)
|
|
80, 80, 80, 80, 40, 40, 40], dtype=np.float32) # right arm
|
|
ARM_KDS = np.array([3, 3, 3, 3, 1.5, 1.5, 1.5,
|
|
3, 3, 3, 3, 1.5, 1.5, 1.5], dtype=np.float32)
|
|
|
|
# Control parameters
|
|
LOCOMOTION_CONTROL_DT = 0.02 # 50Hz control rate
|
|
LOCOMOTION_ACTION_SCALE = 0.25
|
|
ANG_VEL_SCALE = 0.25
|
|
DOF_POS_SCALE = 1.0
|
|
DOF_VEL_SCALE = 0.05
|
|
CMD_SCALE = np.array([2.0, 2.0, 0.25], dtype=np.float32)
|
|
MAX_CMD = np.array([0.8, 0.5, 1.57], dtype=np.float32) # max vx, vy, yaw_rate
|
|
|
|
# Gait parameters
|
|
GAIT_PERIOD = 0.8 # seconds
|
|
|
|
DEFAULT_REPO_ID = "nepyope/unitree_rl_locomotion"
|
|
|
|
|
|
def load_torchscript_policy(
|
|
repo_id: str = DEFAULT_REPO_ID,
|
|
filename: str = "motion.pt",
|
|
) -> torch.jit.ScriptModule:
|
|
"""Load TorchScript locomotion policy from Hugging Face Hub.
|
|
|
|
Args:
|
|
repo_id: Hugging Face Hub repository ID containing the policy.
|
|
filename: Policy filename (default: motion.pt).
|
|
"""
|
|
logger.info(f"Loading TorchScript policy from Hugging Face Hub ({repo_id}/{filename})...")
|
|
|
|
policy_path = hf_hub_download(
|
|
repo_id=repo_id,
|
|
filename=filename,
|
|
)
|
|
|
|
policy = torch.jit.load(policy_path)
|
|
policy.eval()
|
|
|
|
logger.info("TorchScript policy loaded successfully")
|
|
|
|
return policy
|
|
|
|
|
|
class UnitreeRLLocomotionController:
|
|
"""
|
|
Handles 12-DOF legs-only locomotion control for the Unitree G1 robot.
|
|
|
|
This controller manages:
|
|
- Single TorchScript policy
|
|
- 47D observations (single frame)
|
|
- 12D action output (legs only)
|
|
- Arms and waist held at fixed positions
|
|
- Phase-based gait timing
|
|
"""
|
|
|
|
def __init__(self, policy, robot, config):
|
|
self.policy = policy
|
|
self.robot = robot
|
|
self.config = config
|
|
|
|
# Velocity commands (vx, vy, yaw_rate)
|
|
self.locomotion_cmd = np.array([0.0, 0.0, 0.0], dtype=np.float32)
|
|
|
|
# State variables (12 DOF legs)
|
|
self.qj = np.zeros(12, dtype=np.float32)
|
|
self.dqj = np.zeros(12, dtype=np.float32)
|
|
self.locomotion_action = np.zeros(12, dtype=np.float32)
|
|
self.locomotion_obs = np.zeros(47, dtype=np.float32)
|
|
|
|
# Initial arm positions (captured on reset)
|
|
self.initial_arm_positions = np.zeros(14, dtype=np.float32)
|
|
|
|
# Counter for phase calculation
|
|
self.counter = 0
|
|
|
|
# Thread management
|
|
self.locomotion_running = False
|
|
self.locomotion_thread = None
|
|
|
|
logger.info("UnitreeRLLocomotionController initialized")
|
|
logger.info(" Observation dim: 47, Action dim: 12 (legs only)")
|
|
|
|
def locomotion_run(self):
|
|
"""12-DOF legs-only locomotion policy loop."""
|
|
self.counter += 1
|
|
|
|
if self.counter == 1:
|
|
print("\n" + "=" * 60)
|
|
print("🚀 RUNNING UNITREE RL 12-DOF LOCOMOTION POLICY")
|
|
print(" 47D observations → 12D actions (legs only)")
|
|
print(" Arms and waist held at fixed positions")
|
|
print("=" * 60 + "\n")
|
|
|
|
# Get current observation
|
|
robot_state = self.robot.get_observation()
|
|
if robot_state is None:
|
|
return
|
|
|
|
# Get command from remote controller
|
|
if robot_state.wireless_remote is not None:
|
|
self.robot.remote_controller.set(robot_state.wireless_remote)
|
|
else:
|
|
self.robot.remote_controller.lx = 0.0
|
|
self.robot.remote_controller.ly = 0.0
|
|
self.robot.remote_controller.rx = 0.0
|
|
self.robot.remote_controller.ry = 0.0
|
|
|
|
self.locomotion_cmd[0] = self.robot.remote_controller.ly # forward/backward
|
|
self.locomotion_cmd[1] = self.robot.remote_controller.lx * -1 # left/right (inverted)
|
|
self.locomotion_cmd[2] = self.robot.remote_controller.rx * -1 # yaw (inverted)
|
|
|
|
# Get leg joint positions and velocities (12 DOF)
|
|
for i, motor_idx in enumerate(LEG_JOINT_INDICES):
|
|
self.qj[i] = robot_state.motor_state[motor_idx].q
|
|
self.dqj[i] = robot_state.motor_state[motor_idx].dq
|
|
|
|
# Get IMU data
|
|
quat = robot_state.imu_state.quaternion
|
|
ang_vel = np.array(robot_state.imu_state.gyroscope, dtype=np.float32)
|
|
|
|
# Scale observations
|
|
gravity_orientation = self.robot.get_gravity_orientation(quat)
|
|
qj_obs = (self.qj - DEFAULT_LEG_ANGLES) * DOF_POS_SCALE
|
|
dqj_obs = self.dqj * DOF_VEL_SCALE
|
|
ang_vel_scaled = ang_vel * ANG_VEL_SCALE
|
|
|
|
# Calculate phase
|
|
count = self.counter * LOCOMOTION_CONTROL_DT
|
|
phase = (count % GAIT_PERIOD) / GAIT_PERIOD
|
|
sin_phase = np.sin(2 * np.pi * phase)
|
|
cos_phase = np.cos(2 * np.pi * phase)
|
|
|
|
# Build 47D observation vector
|
|
# [0:3] - angular velocity (scaled)
|
|
# [3:6] - gravity orientation
|
|
# [6:9] - velocity command (scaled)
|
|
# [9:21] - joint positions (12D, relative to default)
|
|
# [21:33] - joint velocities (12D, scaled)
|
|
# [33:45] - previous actions (12D)
|
|
# [45] - sin_phase
|
|
# [46] - cos_phase
|
|
self.locomotion_obs[0:3] = ang_vel_scaled
|
|
self.locomotion_obs[3:6] = gravity_orientation
|
|
self.locomotion_obs[6:9] = self.locomotion_cmd * CMD_SCALE * MAX_CMD
|
|
self.locomotion_obs[9:21] = qj_obs
|
|
self.locomotion_obs[21:33] = dqj_obs
|
|
self.locomotion_obs[33:45] = self.locomotion_action
|
|
self.locomotion_obs[45] = sin_phase
|
|
self.locomotion_obs[46] = cos_phase
|
|
|
|
# Run policy inference (TorchScript)
|
|
obs_tensor = torch.from_numpy(self.locomotion_obs).unsqueeze(0).float()
|
|
with torch.no_grad():
|
|
action_tensor = self.policy(obs_tensor)
|
|
self.locomotion_action = action_tensor.squeeze().numpy()
|
|
|
|
# Transform action to target joint positions
|
|
target_leg_pos = DEFAULT_LEG_ANGLES + self.locomotion_action * LOCOMOTION_ACTION_SCALE
|
|
|
|
# Debug logging (first 3 iterations)
|
|
if self.counter <= 3:
|
|
print(f"\n[Unitree RL Debug #{self.counter}]")
|
|
print(f" Phase: {phase:.3f} (sin={sin_phase:.3f}, cos={cos_phase:.3f})")
|
|
print(f" Cmd (vx, vy, yaw): ({self.locomotion_cmd[0]:.2f}, {self.locomotion_cmd[1]:.2f}, {self.locomotion_cmd[2]:.2f})")
|
|
print(f" Action range: [{self.locomotion_action.min():.3f}, {self.locomotion_action.max():.3f}]")
|
|
|
|
# Send commands to LEG motors (0-11)
|
|
for i, motor_idx in enumerate(LEG_JOINT_INDICES):
|
|
self.robot.msg.motor_cmd[motor_idx].q = target_leg_pos[i]
|
|
self.robot.msg.motor_cmd[motor_idx].qd = 0
|
|
self.robot.msg.motor_cmd[motor_idx].kp = LEG_KPS[i]
|
|
self.robot.msg.motor_cmd[motor_idx].kd = LEG_KDS[i]
|
|
self.robot.msg.motor_cmd[motor_idx].tau = 0
|
|
|
|
# Hold WAIST motors at zero (12, 13, 14)
|
|
for i, motor_idx in enumerate(WAIST_JOINT_INDICES):
|
|
self.robot.msg.motor_cmd[motor_idx].q = 0.0
|
|
self.robot.msg.motor_cmd[motor_idx].qd = 0
|
|
self.robot.msg.motor_cmd[motor_idx].kp = WAIST_KPS[i]
|
|
self.robot.msg.motor_cmd[motor_idx].kd = WAIST_KDS[i]
|
|
self.robot.msg.motor_cmd[motor_idx].tau = 0
|
|
|
|
# Hold ARM motors at initial position (15-28)
|
|
for i, motor_idx in enumerate(ARM_JOINT_INDICES):
|
|
self.robot.msg.motor_cmd[motor_idx].q = self.initial_arm_positions[i]
|
|
self.robot.msg.motor_cmd[motor_idx].qd = 0
|
|
self.robot.msg.motor_cmd[motor_idx].kp = ARM_KPS[i]
|
|
self.robot.msg.motor_cmd[motor_idx].kd = ARM_KDS[i]
|
|
self.robot.msg.motor_cmd[motor_idx].tau = 0
|
|
|
|
# Send command
|
|
self.robot.send_action(self.robot.msg)
|
|
|
|
def _locomotion_thread_loop(self):
|
|
"""Background thread that runs the locomotion policy at specified rate."""
|
|
logger.info("Locomotion thread started")
|
|
while self.locomotion_running:
|
|
start_time = time.time()
|
|
try:
|
|
self.locomotion_run()
|
|
except Exception as e:
|
|
logger.error(f"Error in locomotion loop: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
# Sleep to maintain control rate
|
|
elapsed = time.time() - start_time
|
|
sleep_time = max(0, LOCOMOTION_CONTROL_DT - elapsed)
|
|
time.sleep(sleep_time)
|
|
logger.info("Locomotion thread stopped")
|
|
|
|
def start_locomotion_thread(self):
|
|
if self.locomotion_running:
|
|
logger.warning("Locomotion thread already running")
|
|
return
|
|
|
|
logger.info("Starting locomotion control thread...")
|
|
self.locomotion_running = True
|
|
self.locomotion_thread = threading.Thread(target=self._locomotion_thread_loop, daemon=True)
|
|
self.locomotion_thread.start()
|
|
|
|
logger.info("Locomotion control thread started!")
|
|
|
|
def stop_locomotion_thread(self):
|
|
if not self.locomotion_running:
|
|
return
|
|
|
|
logger.info("Stopping locomotion control thread...")
|
|
self.locomotion_running = False
|
|
if self.locomotion_thread:
|
|
self.locomotion_thread.join(timeout=2.0)
|
|
logger.info("Locomotion control thread stopped")
|
|
|
|
def reset_robot(self):
|
|
"""Move legs to default standing position over 2 seconds (arms are captured and held)."""
|
|
logger.info("Moving legs to default position...")
|
|
|
|
total_time = 2.0
|
|
num_step = int(total_time / self.robot.control_dt)
|
|
|
|
# Get current state
|
|
robot_state = self.robot.get_observation()
|
|
|
|
# Capture initial arm positions (to hold during locomotion)
|
|
for i, motor_idx in enumerate(ARM_JOINT_INDICES):
|
|
self.initial_arm_positions[i] = robot_state.motor_state[motor_idx].q
|
|
logger.info(f"Captured initial arm positions: {self.initial_arm_positions[:4]}...")
|
|
|
|
# Record current leg positions
|
|
init_leg_pos = np.zeros(12, dtype=np.float32)
|
|
for i, motor_idx in enumerate(LEG_JOINT_INDICES):
|
|
init_leg_pos[i] = robot_state.motor_state[motor_idx].q
|
|
|
|
# Interpolate legs to default position
|
|
for step in range(num_step):
|
|
alpha = step / num_step
|
|
|
|
# Interpolate leg positions
|
|
for i, motor_idx in enumerate(LEG_JOINT_INDICES):
|
|
target_pos = DEFAULT_LEG_ANGLES[i]
|
|
self.robot.msg.motor_cmd[motor_idx].q = (
|
|
init_leg_pos[i] * (1 - alpha) + target_pos * alpha
|
|
)
|
|
self.robot.msg.motor_cmd[motor_idx].qd = 0
|
|
self.robot.msg.motor_cmd[motor_idx].kp = LEG_KPS[i]
|
|
self.robot.msg.motor_cmd[motor_idx].kd = LEG_KDS[i]
|
|
self.robot.msg.motor_cmd[motor_idx].tau = 0
|
|
|
|
# Hold waist at zero
|
|
for i, motor_idx in enumerate(WAIST_JOINT_INDICES):
|
|
self.robot.msg.motor_cmd[motor_idx].q = 0.0
|
|
self.robot.msg.motor_cmd[motor_idx].qd = 0
|
|
self.robot.msg.motor_cmd[motor_idx].kp = WAIST_KPS[i]
|
|
self.robot.msg.motor_cmd[motor_idx].kd = WAIST_KDS[i]
|
|
self.robot.msg.motor_cmd[motor_idx].tau = 0
|
|
|
|
# Hold arms at initial position
|
|
for i, motor_idx in enumerate(ARM_JOINT_INDICES):
|
|
self.robot.msg.motor_cmd[motor_idx].q = self.initial_arm_positions[i]
|
|
self.robot.msg.motor_cmd[motor_idx].qd = 0
|
|
self.robot.msg.motor_cmd[motor_idx].kp = ARM_KPS[i]
|
|
self.robot.msg.motor_cmd[motor_idx].kd = ARM_KDS[i]
|
|
self.robot.msg.motor_cmd[motor_idx].tau = 0
|
|
|
|
self.robot.msg.crc = self.robot.crc.Crc(self.robot.msg)
|
|
self.robot.lowcmd_publisher.Write(self.robot.msg)
|
|
time.sleep(self.robot.control_dt)
|
|
|
|
logger.info("Reached default leg position")
|
|
|
|
# Hold position for 2 seconds
|
|
logger.info("Holding default position for 2 seconds...")
|
|
hold_time = 2.0
|
|
num_hold_steps = int(hold_time / self.robot.control_dt)
|
|
|
|
for _ in range(num_hold_steps):
|
|
# Hold legs at default
|
|
for i, motor_idx in enumerate(LEG_JOINT_INDICES):
|
|
self.robot.msg.motor_cmd[motor_idx].q = DEFAULT_LEG_ANGLES[i]
|
|
self.robot.msg.motor_cmd[motor_idx].qd = 0
|
|
self.robot.msg.motor_cmd[motor_idx].kp = LEG_KPS[i]
|
|
self.robot.msg.motor_cmd[motor_idx].kd = LEG_KDS[i]
|
|
self.robot.msg.motor_cmd[motor_idx].tau = 0
|
|
|
|
# Hold waist at zero
|
|
for i, motor_idx in enumerate(WAIST_JOINT_INDICES):
|
|
self.robot.msg.motor_cmd[motor_idx].q = 0.0
|
|
self.robot.msg.motor_cmd[motor_idx].qd = 0
|
|
self.robot.msg.motor_cmd[motor_idx].kp = WAIST_KPS[i]
|
|
self.robot.msg.motor_cmd[motor_idx].kd = WAIST_KDS[i]
|
|
self.robot.msg.motor_cmd[motor_idx].tau = 0
|
|
|
|
# Hold arms at initial position
|
|
for i, motor_idx in enumerate(ARM_JOINT_INDICES):
|
|
self.robot.msg.motor_cmd[motor_idx].q = self.initial_arm_positions[i]
|
|
self.robot.msg.motor_cmd[motor_idx].qd = 0
|
|
self.robot.msg.motor_cmd[motor_idx].kp = ARM_KPS[i]
|
|
self.robot.msg.motor_cmd[motor_idx].kd = ARM_KDS[i]
|
|
self.robot.msg.motor_cmd[motor_idx].tau = 0
|
|
|
|
self.robot.msg.crc = self.robot.crc.Crc(self.robot.msg)
|
|
self.robot.lowcmd_publisher.Write(self.robot.msg)
|
|
time.sleep(self.robot.control_dt)
|
|
|
|
logger.info("Ready to start locomotion!")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Unitree RL 12-DOF Locomotion Controller for Unitree G1")
|
|
parser.add_argument(
|
|
"--repo-id",
|
|
type=str,
|
|
default=DEFAULT_REPO_ID,
|
|
help=f"Hugging Face Hub repo ID for policy (default: {DEFAULT_REPO_ID})",
|
|
)
|
|
parser.add_argument(
|
|
"--filename",
|
|
type=str,
|
|
default="motion.pt",
|
|
help="Policy filename (default: motion.pt)",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
# Load policy
|
|
policy = load_torchscript_policy(repo_id=args.repo_id, filename=args.filename)
|
|
|
|
# Initialize robot
|
|
config = UnitreeG1Config()
|
|
robot = UnitreeG1(config)
|
|
|
|
# Initialize locomotion controller
|
|
locomotion_controller = UnitreeRLLocomotionController(
|
|
policy=policy,
|
|
robot=robot,
|
|
config=config,
|
|
)
|
|
|
|
# Reset robot and start locomotion thread
|
|
try:
|
|
locomotion_controller.reset_robot()
|
|
locomotion_controller.start_locomotion_thread()
|
|
|
|
# Log status
|
|
logger.info("Robot initialized with Unitree RL locomotion policy")
|
|
logger.info("Locomotion controller running in background thread")
|
|
logger.info("Use remote controller to command velocity:")
|
|
logger.info(" Left stick Y: forward/backward")
|
|
logger.info(" Left stick X: left/right")
|
|
logger.info(" Right stick X: rotate")
|
|
logger.info("Press Ctrl+C to stop")
|
|
|
|
# Keep robot alive
|
|
while True:
|
|
time.sleep(1.0)
|
|
except KeyboardInterrupt:
|
|
print("\nStopping locomotion...")
|
|
locomotion_controller.stop_locomotion_thread()
|
|
print("Done!")
|
|
|