admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-05-17 09:39:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

properly comment config, example locomotion and unitree_g1 class

Browse Source
This commit is contained in:
Martino Russi
2025-11-26 21:27:45 +01:00
parent 55ee13aec1
commit 3ec332fabc
3 changed files with 127 additions and 284 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/unitree_g1/gr00t_locomotion.py
+83 -87
View File
@@ -44,26 +44,28 @@ GROOT_DEFAULT_ANGLES = np.array(
0.0,
0.0,
0.0,
0.0, # left arm (zeroed)
0.0, # left arm
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0, # right arm (zeroed)
0.0, # right arm
],
dtype=np.float32,
)
JOINTS_TO_ZERO = [12, 14, 20, 21, 27, 28] # waist yaw/pitch, wrist pitch/yaw
PROBLEMATIC_JOINTS = [12, 14, 20, 21, 27, 28]
G1_MODEL = "g1_23"
if G1_MODEL == "g1_23":
MISSING_JOINTS = [12, 14, 20, 21, 27, 28] # waist yaw/pitch, wrist pitch/yaw
elif G1_MODEL == "g1_29":
MISSING_JOINTS = [] # waist yaw/pitch, wrist pitch/yaw
LOCOMOTION_ACTION_SCALE = 0.25
LOCOMOTION_CONTROL_DT = 0.02
ANG_VEL_SCALE: float = 0.25
DOF_POS_SCALE: float = 1.0
DOF_VEL_SCALE: float = 0.05
@@ -97,21 +99,13 @@ class GrootLocomotionController:
"""
def __init__(self, policy_balance, policy_walk, robot, config):
"""
Initialize the GR00T locomotion controller.
Args:
policy_balance: ONNX InferenceSession for balance/standing policy
policy_walk: ONNX InferenceSession for walking policy
robot: Reference to the UnitreeG1 robot instance
config: UnitreeG1Config object with locomotion parameters
"""
self.policy_balance = policy_balance
self.policy_walk = policy_walk
self.robot = robot
self.config = config
self.locomotion_cmd = np.array([0.0, 0.0, 0.0], dtype=np.float32) # vx, vy, yaw_rate
self.locomotion_cmd = np.array([0.0, 0.0, 0.0], dtype=np.float32) # vx, vy, theta_dot
# GR00T-specific state
self.groot_qj_all = np.zeros(29, dtype=np.float32)
@@ -123,7 +117,7 @@ class GrootLocomotionController:
self.groot_height_cmd = 0.74 # Default base height
self.groot_orientation_cmd = np.array([0.0, 0.0, 0.0], dtype=np.float32)
# Initialize history with zeros
# input to gr00t is 6 frames (6*86D=516)
for _ in range(6):
self.groot_obs_history.append(np.zeros(86, dtype=np.float32))
@@ -134,42 +128,39 @@ class GrootLocomotionController:
logger.info("GrootLocomotionController initialized")
def groot_locomotion_run(self):
# Get current obs
# get current observation
robot_state = self.robot.get_observation()
if robot_state is None:
return
# Update remote controller from lowstate
# get command from remote controller
if robot_state.wireless_remote is not None:
self.robot.remote_controller.set(robot_state.wireless_remote)
# R1/R2 buttons for height control on real robot (button indices 0 and 4)
if self.robot.remote_controller.button[0]: # R1 - raise height
self.groot_height_cmd += 0.001 # Small increment per timestep (~0.05m per second at 50Hz)
if self.robot.remote_controller.button[0]: # R1 - raise waist
self.groot_height_cmd += 0.001
self.groot_height_cmd = np.clip(self.groot_height_cmd, 0.50, 1.00)
if self.robot.remote_controller.button[4]: # R2 - lower height
self.groot_height_cmd -= 0.001 # Small decrement per timestep
if self.robot.remote_controller.button[4]: # R2 - lower waist
self.groot_height_cmd -= 0.001
self.groot_height_cmd = np.clip(self.groot_height_cmd, 0.50, 1.00)
else:
# Default to zero commands if no remote data
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
# Get ALL 29 joint positions and velocities
self.locomotion_cmd[0] = self.robot.remote_controller.ly # forward/backward
self.locomotion_cmd[1] = self.robot.remote_controller.lx * -1 # left/right
self.locomotion_cmd[2] = self.robot.remote_controller.rx * -1 # rotation rate
for i in range(29):
self.groot_qj_all[i] = robot_state.motor_state[i].q
self.groot_dqj_all[i] = robot_state.motor_state[i].dq
# Get IMU data
quat = robot_state.imu_state.quaternion
ang_vel = np.array(robot_state.imu_state.gyroscope, dtype=np.float32)
gravity_orientation = self.robot.get_gravity_orientation(quat)
# Zero out specific joints in observation
for idx in JOINTS_TO_ZERO:
# adapt observation for g1_23dof
for idx in MISSING_JOINTS:
self.groot_qj_all[idx] = 0.0
self.groot_dqj_all[idx] = 0.0
@@ -177,24 +168,25 @@ class GrootLocomotionController:
qj_obs = self.groot_qj_all.copy()
dqj_obs = self.groot_dqj_all.copy()
# express imu data in gravity frame of reference
quat = robot_state.imu_state.quaternion
ang_vel = np.array(robot_state.imu_state.gyroscope, dtype=np.float32)
gravity_orientation = self.robot.get_gravity_orientation(quat)
#scale joint positions and velocities before policy inference
qj_obs = (qj_obs - GROOT_DEFAULT_ANGLES) * DOF_POS_SCALE
dqj_obs = dqj_obs * DOF_VEL_SCALE
ang_vel_scaled = ang_vel * ANG_VEL_SCALE
# Get velocity commands (keyboard or remote)
if not self.robot.simulation_mode:
self.locomotion_cmd[0] = self.robot.remote_controller.ly
self.locomotion_cmd[1] = self.robot.remote_controller.lx * -1
self.locomotion_cmd[2] = self.robot.remote_controller.rx * -1
# Build 86D single frame observation (GR00T format)
# build single frame observation
self.groot_obs_single[:3] = self.locomotion_cmd * np.array(CMD_SCALE)
self.groot_obs_single[3] = self.groot_height_cmd
self.groot_obs_single[4:7] = self.groot_orientation_cmd
self.groot_obs_single[7:10] = ang_vel_scaled
self.groot_obs_single[10:13] = gravity_orientation
self.groot_obs_single[13:42] = qj_obs # 29D joint positions
self.groot_obs_single[42:71] = dqj_obs # 29D joint velocities
self.groot_obs_single[13:42] = qj_obs
self.groot_obs_single[42:71] = dqj_obs
self.groot_obs_single[71:86] = self.groot_action # 15D previous actions
# Add to history and stack observations (6 frames × 86D = 516D)
@@ -209,29 +201,25 @@ class GrootLocomotionController:
# Run policy inference (ONNX) with 516D stacked observation
obs_tensor = torch.from_numpy(self.groot_obs_stacked).unsqueeze(0)
# Select appropriate policy based on command magnitude (dual-policy system)
cmd_magnitude = np.linalg.norm(self.locomotion_cmd)
if cmd_magnitude < 0.05:
# Use balance/standing policy for small commands
# balance/standing policy for small commands
selected_policy = self.policy_balance
else:
# Use walking policy for movement commands
# walking policy for movement commands
selected_policy = self.policy_walk
# run policy inference
ort_inputs = {selected_policy.get_inputs()[0].name: obs_tensor.cpu().numpy()}
ort_outs = selected_policy.run(None, ort_inputs)
self.groot_action = ort_outs[0].squeeze()
# Zero out waist actions (yaw=12, roll=13, pitch=14) - only use leg actions (0-11)
self.groot_action[12] = 0.0 # Waist yaw
self.groot_action[13] = 0.0 # Waist roll
self.groot_action[14] = 0.0 # Waist pitch
# Transform action to target joint positions (15D: legs + waist)
# transform action back to target joint positions
target_dof_pos_15 = GROOT_DEFAULT_ANGLES[:15] + self.groot_action * LOCOMOTION_ACTION_SCALE
# Send commands to LEG motors (0-11)
for i in range(12):
# command motors
for i in range(15):
motor_idx = i
self.robot.msg.motor_cmd[motor_idx].q = target_dof_pos_15[i]
self.robot.msg.motor_cmd[motor_idx].qd = 0
@@ -239,24 +227,15 @@ class GrootLocomotionController:
self.robot.msg.motor_cmd[motor_idx].kd = self.robot.kd[motor_idx]
self.robot.msg.motor_cmd[motor_idx].tau = 0
# Send WAIST commands - but SKIP waist yaw (12) and waist pitch (14)
# Only send waist roll (13)
waist_roll_idx = 13
waist_roll_action_idx = 13
self.robot.msg.motor_cmd[waist_roll_idx].q = target_dof_pos_15[waist_roll_action_idx]
self.robot.msg.motor_cmd[waist_roll_idx].qd = 0
self.robot.msg.motor_cmd[waist_roll_idx].kp = self.robot.kp[waist_roll_idx]
self.robot.msg.motor_cmd[waist_roll_idx].kd = self.robot.kd[waist_roll_idx]
self.robot.msg.motor_cmd[waist_roll_idx].tau = 0
# Zero out the problematic joints (waist yaw, waist pitch, wrist pitch/yaw)
for joint_idx in PROBLEMATIC_JOINTS:
# adapt action for g1_23dof
for joint_idx in MISSING_JOINTS:
self.robot.msg.motor_cmd[joint_idx].q = 0.0
self.robot.msg.motor_cmd[joint_idx].qd = 0
self.robot.msg.motor_cmd[joint_idx].kp = self.robot.kp[joint_idx]
self.robot.msg.motor_cmd[joint_idx].kd = self.robot.kd[joint_idx]
self.robot.msg.motor_cmd[joint_idx].tau = 0
#send action to robot
self.robot.send_action(self.robot.msg)
def _locomotion_thread_loop(self):
@@ -276,7 +255,6 @@ class GrootLocomotionController:
logger.info("Locomotion thread stopped")
def start_locomotion_thread(self):
"""Start the background locomotion control thread."""
if self.locomotion_running:
logger.warning("Locomotion thread already running")
return
@@ -285,10 +263,10 @@ class GrootLocomotionController:
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):
"""Stop the background locomotion control thread."""
if not self.locomotion_running:
return
@@ -298,32 +276,48 @@ class GrootLocomotionController:
self.locomotion_thread.join(timeout=2.0)
logger.info("Locomotion control thread stopped")
def init_groot_locomotion(self):
"""Initialize GR00T-style locomotion for ONNX policies (29 DOF, 15D actions)."""
logger.info("Starting GR00T locomotion initialization...")
def reset_robot(self):
"""Move robot legs to default standing position over 2 seconds (arms are not moved)."""
total_time = 3.0
num_step = int(total_time / self.robot.control_dt)
# Reset legs to default position
self.robot.reset_legs()
# Only control legs, not arms (first 12 joints)
default_pos = GROOT_DEFAULT_ANGLES # First 12 values are leg angles
dof_size = len(default_pos)
# Wait 3 seconds
time.sleep(3.0)
# Get current lowstate
robot_state = self.robot.get_observation()
# Start locomotion policy thread
logger.info("Starting GR00T locomotion policy control...")
self.start_locomotion_thread()
# Record the current leg positions
init_dof_pos = np.zeros(dof_size, dtype=np.float32)
for i in range(dof_size):
init_dof_pos[i] = robot_state.motor_state[i].q
# Move legs to default pos
for i in range(num_step):
alpha = i / num_step
for motor_idx in range(dof_size):
target_pos = default_pos[motor_idx]
self.robot.msg.motor_cmd[motor_idx].q = init_dof_pos[motor_idx] * (1 - alpha) + target_pos * alpha
self.robot.msg.motor_cmd[motor_idx].qd = 0
self.robot.msg.motor_cmd[motor_idx].kp = self.robot.kp[motor_idx]
self.robot.msg.motor_cmd[motor_idx].kd = self.robot.kd[motor_idx]
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 position (legs only)")
if __name__ == "__main__":
# 1. Load policies externally (separate from robot initialization)
#load policies
policy_balance, policy_walk = load_groot_policies()
# 2. Create config (no locomotion_control=True since we're using external controller)
#initialize robot
config = UnitreeG1Config()
# 3. Initialize robot
robot = UnitreeG1(config)
# 4. Create GR00T locomotion controller with loaded policies
#initialize gr00t locomotion controller
groot_controller = GrootLocomotionController(
policy_balance=policy_balance,
policy_walk=policy_walk,
@@ -331,18 +325,20 @@ if __name__ == "__main__":
config=config,
)
# 5. Initialize and start locomotion
groot_controller.init_groot_locomotion()
#reset legs and start locomotion thread
groot_controller.reset_robot()
groot_controller.start_locomotion_thread()
# Robot is now ready with locomotion control!
print("Robot initialized with GR00T locomotion policies")
print("Locomotion controller running in background thread")
print("Press Ctrl+C to stop")
#log status
logger.info("Robot initialized with GR00T locomotion policies")
logger.info("Locomotion controller running in background thread")
logger.info("Press Ctrl+C to stop")
#keep robot alive
try:
while True:
time.sleep(1.0)
except KeyboardInterrupt:
print("\nStopping locomotion...")
groot_controller.stop_locomotion_thread()
print("Done!")
print("Done!")