properly comment config, example locomotion and unitree_g1 class

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
+82 -86
View File
@@ -44,26 +44,28 @@ GROOT_DEFAULT_ANGLES = np.array(
0.0, 0.0,
0.0, 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, 0.0,
0.0, 0.0,
0.0, 0.0,
0.0, # right arm (zeroed) 0.0, # right arm
], ],
dtype=np.float32, dtype=np.float32,
) )
JOINTS_TO_ZERO = [12, 14, 20, 21, 27, 28] # waist yaw/pitch, wrist pitch/yaw G1_MODEL = "g1_23"
PROBLEMATIC_JOINTS = [12, 14, 20, 21, 27, 28] 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_ACTION_SCALE = 0.25
LOCOMOTION_CONTROL_DT = 0.02 LOCOMOTION_CONTROL_DT = 0.02
ANG_VEL_SCALE: float = 0.25 ANG_VEL_SCALE: float = 0.25
DOF_POS_SCALE: float = 1.0 DOF_POS_SCALE: float = 1.0
DOF_VEL_SCALE: float = 0.05 DOF_VEL_SCALE: float = 0.05
@@ -97,21 +99,13 @@ class GrootLocomotionController:
""" """
def __init__(self, policy_balance, policy_walk, robot, config): 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_balance = policy_balance
self.policy_walk = policy_walk self.policy_walk = policy_walk
self.robot = robot self.robot = robot
self.config = config 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 # GR00T-specific state
self.groot_qj_all = np.zeros(29, dtype=np.float32) 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_height_cmd = 0.74 # Default base height
self.groot_orientation_cmd = np.array([0.0, 0.0, 0.0], dtype=np.float32) 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): for _ in range(6):
self.groot_obs_history.append(np.zeros(86, dtype=np.float32)) self.groot_obs_history.append(np.zeros(86, dtype=np.float32))
@@ -134,42 +128,39 @@ class GrootLocomotionController:
logger.info("GrootLocomotionController initialized") logger.info("GrootLocomotionController initialized")
def groot_locomotion_run(self): def groot_locomotion_run(self):
# Get current obs
# get current observation
robot_state = self.robot.get_observation() robot_state = self.robot.get_observation()
if robot_state is None: if robot_state is None:
return return
# Update remote controller from lowstate # get command from remote controller
if robot_state.wireless_remote is not None: if robot_state.wireless_remote is not None:
self.robot.remote_controller.set(robot_state.wireless_remote) self.robot.remote_controller.set(robot_state.wireless_remote)
if self.robot.remote_controller.button[0]: # R1 - raise waist
# R1/R2 buttons for height control on real robot (button indices 0 and 4) self.groot_height_cmd += 0.001
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)
self.groot_height_cmd = np.clip(self.groot_height_cmd, 0.50, 1.00) self.groot_height_cmd = np.clip(self.groot_height_cmd, 0.50, 1.00)
if self.robot.remote_controller.button[4]: # R2 - lower height if self.robot.remote_controller.button[4]: # R2 - lower waist
self.groot_height_cmd -= 0.001 # Small decrement per timestep self.groot_height_cmd -= 0.001
self.groot_height_cmd = np.clip(self.groot_height_cmd, 0.50, 1.00) self.groot_height_cmd = np.clip(self.groot_height_cmd, 0.50, 1.00)
else: else:
# Default to zero commands if no remote data
self.robot.remote_controller.lx = 0.0 self.robot.remote_controller.lx = 0.0
self.robot.remote_controller.ly = 0.0 self.robot.remote_controller.ly = 0.0
self.robot.remote_controller.rx = 0.0 self.robot.remote_controller.rx = 0.0
self.robot.remote_controller.ry = 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): for i in range(29):
self.groot_qj_all[i] = robot_state.motor_state[i].q self.groot_qj_all[i] = robot_state.motor_state[i].q
self.groot_dqj_all[i] = robot_state.motor_state[i].dq 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) # adapt observation for g1_23dof
for idx in MISSING_JOINTS:
# Zero out specific joints in observation
for idx in JOINTS_TO_ZERO:
self.groot_qj_all[idx] = 0.0 self.groot_qj_all[idx] = 0.0
self.groot_dqj_all[idx] = 0.0 self.groot_dqj_all[idx] = 0.0
@@ -177,24 +168,25 @@ class GrootLocomotionController:
qj_obs = self.groot_qj_all.copy() qj_obs = self.groot_qj_all.copy()
dqj_obs = self.groot_dqj_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 qj_obs = (qj_obs - GROOT_DEFAULT_ANGLES) * DOF_POS_SCALE
dqj_obs = dqj_obs * DOF_VEL_SCALE dqj_obs = dqj_obs * DOF_VEL_SCALE
ang_vel_scaled = ang_vel * ANG_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.locomotion_cmd * np.array(CMD_SCALE)
self.groot_obs_single[3] = self.groot_height_cmd self.groot_obs_single[3] = self.groot_height_cmd
self.groot_obs_single[4:7] = self.groot_orientation_cmd self.groot_obs_single[4:7] = self.groot_orientation_cmd
self.groot_obs_single[7:10] = ang_vel_scaled self.groot_obs_single[7:10] = ang_vel_scaled
self.groot_obs_single[10:13] = gravity_orientation self.groot_obs_single[10:13] = gravity_orientation
self.groot_obs_single[13:42] = qj_obs # 29D joint positions self.groot_obs_single[13:42] = qj_obs
self.groot_obs_single[42:71] = dqj_obs # 29D joint velocities self.groot_obs_single[42:71] = dqj_obs
self.groot_obs_single[71:86] = self.groot_action # 15D previous actions self.groot_obs_single[71:86] = self.groot_action # 15D previous actions
# Add to history and stack observations (6 frames × 86D = 516D) # Add to history and stack observations (6 frames × 86D = 516D)
@@ -209,29 +201,25 @@ class GrootLocomotionController:
# Run policy inference (ONNX) with 516D stacked observation # Run policy inference (ONNX) with 516D stacked observation
obs_tensor = torch.from_numpy(self.groot_obs_stacked).unsqueeze(0) 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) cmd_magnitude = np.linalg.norm(self.locomotion_cmd)
if cmd_magnitude < 0.05: if cmd_magnitude < 0.05:
# Use balance/standing policy for small commands # balance/standing policy for small commands
selected_policy = self.policy_balance selected_policy = self.policy_balance
else: else:
# Use walking policy for movement commands # walking policy for movement commands
selected_policy = self.policy_walk selected_policy = self.policy_walk
# run policy inference
ort_inputs = {selected_policy.get_inputs()[0].name: obs_tensor.cpu().numpy()} ort_inputs = {selected_policy.get_inputs()[0].name: obs_tensor.cpu().numpy()}
ort_outs = selected_policy.run(None, ort_inputs) ort_outs = selected_policy.run(None, ort_inputs)
self.groot_action = ort_outs[0].squeeze() self.groot_action = ort_outs[0].squeeze()
# Zero out waist actions (yaw=12, roll=13, pitch=14) - only use leg actions (0-11) # transform action back to target joint positions
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)
target_dof_pos_15 = GROOT_DEFAULT_ANGLES[:15] + self.groot_action * LOCOMOTION_ACTION_SCALE target_dof_pos_15 = GROOT_DEFAULT_ANGLES[:15] + self.groot_action * LOCOMOTION_ACTION_SCALE
# Send commands to LEG motors (0-11) # command motors
for i in range(12): for i in range(15):
motor_idx = i motor_idx = i
self.robot.msg.motor_cmd[motor_idx].q = target_dof_pos_15[i] self.robot.msg.motor_cmd[motor_idx].q = target_dof_pos_15[i]
self.robot.msg.motor_cmd[motor_idx].qd = 0 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].kd = self.robot.kd[motor_idx]
self.robot.msg.motor_cmd[motor_idx].tau = 0 self.robot.msg.motor_cmd[motor_idx].tau = 0
# Send WAIST commands - but SKIP waist yaw (12) and waist pitch (14) # adapt action for g1_23dof
# Only send waist roll (13) for joint_idx in MISSING_JOINTS:
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:
self.robot.msg.motor_cmd[joint_idx].q = 0.0 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].qd = 0
self.robot.msg.motor_cmd[joint_idx].kp = self.robot.kp[joint_idx] 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].kd = self.robot.kd[joint_idx]
self.robot.msg.motor_cmd[joint_idx].tau = 0 self.robot.msg.motor_cmd[joint_idx].tau = 0
#send action to robot
self.robot.send_action(self.robot.msg) self.robot.send_action(self.robot.msg)
def _locomotion_thread_loop(self): def _locomotion_thread_loop(self):
@@ -276,7 +255,6 @@ class GrootLocomotionController:
logger.info("Locomotion thread stopped") logger.info("Locomotion thread stopped")
def start_locomotion_thread(self): def start_locomotion_thread(self):
"""Start the background locomotion control thread."""
if self.locomotion_running: if self.locomotion_running:
logger.warning("Locomotion thread already running") logger.warning("Locomotion thread already running")
return return
@@ -285,10 +263,10 @@ class GrootLocomotionController:
self.locomotion_running = True self.locomotion_running = True
self.locomotion_thread = threading.Thread(target=self._locomotion_thread_loop, daemon=True) self.locomotion_thread = threading.Thread(target=self._locomotion_thread_loop, daemon=True)
self.locomotion_thread.start() self.locomotion_thread.start()
logger.info("Locomotion control thread started!") logger.info("Locomotion control thread started!")
def stop_locomotion_thread(self): def stop_locomotion_thread(self):
"""Stop the background locomotion control thread."""
if not self.locomotion_running: if not self.locomotion_running:
return return
@@ -298,32 +276,48 @@ class GrootLocomotionController:
self.locomotion_thread.join(timeout=2.0) self.locomotion_thread.join(timeout=2.0)
logger.info("Locomotion control thread stopped") logger.info("Locomotion control thread stopped")
def init_groot_locomotion(self): def reset_robot(self):
"""Initialize GR00T-style locomotion for ONNX policies (29 DOF, 15D actions).""" """Move robot legs to default standing position over 2 seconds (arms are not moved)."""
logger.info("Starting GR00T locomotion initialization...") total_time = 3.0
num_step = int(total_time / self.robot.control_dt)
# Reset legs to default position # Only control legs, not arms (first 12 joints)
self.robot.reset_legs() default_pos = GROOT_DEFAULT_ANGLES # First 12 values are leg angles
dof_size = len(default_pos)
# Wait 3 seconds # Get current lowstate
time.sleep(3.0) robot_state = self.robot.get_observation()
# Start locomotion policy thread # Record the current leg positions
logger.info("Starting GR00T locomotion policy control...") init_dof_pos = np.zeros(dof_size, dtype=np.float32)
self.start_locomotion_thread() 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__": if __name__ == "__main__":
# 1. Load policies externally (separate from robot initialization)
#load policies
policy_balance, policy_walk = load_groot_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() config = UnitreeG1Config()
# 3. Initialize robot
robot = UnitreeG1(config) robot = UnitreeG1(config)
# 4. Create GR00T locomotion controller with loaded policies #initialize gr00t locomotion controller
groot_controller = GrootLocomotionController( groot_controller = GrootLocomotionController(
policy_balance=policy_balance, policy_balance=policy_balance,
policy_walk=policy_walk, policy_walk=policy_walk,
@@ -331,14 +325,16 @@ if __name__ == "__main__":
config=config, config=config,
) )
# 5. Initialize and start locomotion #reset legs and start locomotion thread
groot_controller.init_groot_locomotion() groot_controller.reset_robot()
groot_controller.start_locomotion_thread()
# Robot is now ready with locomotion control! #log status
print("Robot initialized with GR00T locomotion policies") logger.info("Robot initialized with GR00T locomotion policies")
print("Locomotion controller running in background thread") logger.info("Locomotion controller running in background thread")
print("Press Ctrl+C to stop") logger.info("Press Ctrl+C to stop")
#keep robot alive
try: try:
while True: while True:
time.sleep(1.0) time.sleep(1.0)
@@ -26,7 +26,6 @@ from ..config import RobotConfig
@dataclass @dataclass
class UnitreeG1Config(RobotConfig): class UnitreeG1Config(RobotConfig):
# id: str = "unitree_g1" # id: str = "unitree_g1"
simulation_mode: bool = False
kp: list = field(default_factory=lambda: [ kp: list = field(default_factory=lambda: [
150, 150, 150, 300, 40, 40, # Left leg pitch, roll, yaw, knee, ankle pitch, ankle roll 150, 150, 150, 300, 40, 40, # Left leg pitch, roll, yaw, knee, ankle pitch, ankle roll
@@ -50,12 +49,6 @@ class UnitreeG1Config(RobotConfig):
3, 3, 3, 3, 3, 3, # Other 3, 3, 3, 3, 3, 3, # Other
]) ])
arm_velocity_limit = 100.0 control_dt = 1.0 / 250.0 # 250Hz
control_dt = 1.0 / 250.0
leg_joint2motor_idx: list = field(default_factory=lambda: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11])
default_leg_angles: list = field(
default_factory=lambda: [-0.1, 0.0, 0.0, 0.3, -0.2, 0.0, -0.1, 0.0, 0.0, 0.3, -0.2, 0.0]
)
+43 -189
View File
@@ -12,10 +12,10 @@ from unitree_sdk2py.idl.default import unitree_hg_msg_dds__LowCmd_
from unitree_sdk2py.idl.unitree_hg.msg.dds_ import ( from unitree_sdk2py.idl.unitree_hg.msg.dds_ import (
LowCmd_ as hg_LowCmd, LowCmd_ as hg_LowCmd,
LowState_ as hg_LowState, LowState_ as hg_LowState,
) # idl for g1, h1_2 )
from unitree_sdk2py.utils.crc import CRC from unitree_sdk2py.utils.crc import CRC
from lerobot.robots.unitree_g1.g1_utils import G1_29_JointArmIndex, G1_29_JointIndex from lerobot.robots.unitree_g1.g1_utils import G1_29_JointIndex
from lerobot.robots.unitree_g1.unitree_sdk2_socket import ( from lerobot.robots.unitree_g1.unitree_sdk2_socket import (
ChannelFactoryInitialize, ChannelFactoryInitialize,
ChannelPublisher, ChannelPublisher,
@@ -36,12 +36,13 @@ H1_2_Num_Motors = 35
H1_Num_Motors = 20 H1_Num_Motors = 20
class MotorState: class MotorState:
def __init__(self): def __init__(self):
self.q = None self.q = None # position
self.dq = None self.dq = None # velocity
self.tau_est = None # Estimated torque self.tau_est = None # estimated torque
self.temperature = None # Motor temperature self.temperature = None # motor temperature
class IMUState: class IMUState:
@@ -52,7 +53,7 @@ class IMUState:
self.rpy = None # [roll, pitch, yaw] (rad) self.rpy = None # [roll, pitch, yaw] (rad)
self.temperature = None # IMU temperature self.temperature = None # IMU temperature
#g1 observation class
class G1_29_LowState: class G1_29_LowState:
def __init__(self): def __init__(self):
self.motor_state = [MotorState() for _ in range(G1_29_Num_Motors)] self.motor_state = [MotorState() for _ in range(G1_29_Num_Motors)]
@@ -78,7 +79,9 @@ class UnitreeG1(Robot):
config_class = UnitreeG1Config config_class = UnitreeG1Config
name = "unitree_g1" name = "unitree_g1"
#unitree remote controller
class RemoteController: class RemoteController:
def __init__(self): def __init__(self):
self.lx = 0 self.lx = 0
self.ly = 0 self.ly = 0
@@ -102,60 +105,48 @@ class UnitreeG1(Robot):
logger.info("Initialize UnitreeG1...") logger.info("Initialize UnitreeG1...")
self.config = config self.config = config
self.q_target = np.zeros(14)
self.tauff_target = np.zeros(14)
self.simulation_mode = config.simulation_mode
# Unified kp/kd arrays for all 35 motors
self.kp = np.array(config.kp, dtype=np.float32)
self.kd = np.array(config.kd, dtype=np.float32)
self.arm_velocity_limit = config.arm_velocity_limit
self.control_dt = config.control_dt self.control_dt = config.control_dt
# Initialize DDS # connect robot
ChannelFactoryInitialize(0) self.connect()
# Always use debug mode (direct motor control) # initialize direct motor control interface
self.lowcmd_publisher = ChannelPublisher(kTopicLowCommand_Debug, hg_LowCmd) self.lowcmd_publisher = ChannelPublisher(kTopicLowCommand_Debug, hg_LowCmd)
self.lowcmd_publisher.Init() self.lowcmd_publisher.Init()
self.lowstate_subscriber = ChannelSubscriber(kTopicLowState, hg_LowState) self.lowstate_subscriber = ChannelSubscriber(kTopicLowState, hg_LowState)
self.lowstate_subscriber.Init() self.lowstate_subscriber.Init()
self.lowstate_buffer = DataBuffer() self.lowstate_buffer = DataBuffer()
# initialize subscribe thread # initialize subscribe thread to read robot state
self.subscribe_thread = threading.Thread(target=self._subscribe_motor_state) self.subscribe_thread = threading.Thread(target=self._subscribe_motor_state)
self.subscribe_thread.daemon = True self.subscribe_thread.daemon = True
self.subscribe_thread.start() self.subscribe_thread.start()
while not self.lowstate_buffer.GetData():
time.sleep(0.1)
logger.warning("[UnitreeG1] Waiting to subscribe dds...")
logger.warning("[UnitreeG1] Subscribe dds ok.")
# initialize hg's lowcmd msg # initialize hg's lowcmd msg
self.crc = CRC() self.crc = CRC()
self.msg = unitree_hg_msg_dds__LowCmd_() self.msg = unitree_hg_msg_dds__LowCmd_()
self.msg.mode_pr = 0 self.msg.mode_pr = 0
self.msg.mode_machine = self.get_mode_machine() self.msg.mode_machine = self.lowstate_subscriber.Read().mode_machine
# initialize all motors with unified kp/kd from config
lowstate = self.lowstate_buffer.GetData()
self.kp = np.array(config.kp, dtype=np.float32)
self.kd = np.array(config.kd, dtype=np.float32)
# Initialize all motors with unified kp/kd from config
for id in G1_29_JointIndex: for id in G1_29_JointIndex:
self.msg.motor_cmd[id].mode = 1 self.msg.motor_cmd[id].mode = 1
self.msg.motor_cmd[id].kp = self.kp[id.value] self.msg.motor_cmd[id].kp = self.kp[id.value]
self.msg.motor_cmd[id].kd = self.kd[id.value] self.msg.motor_cmd[id].kd = self.kd[id.value]
self.msg.motor_cmd[id].q = self.get_current_motor_q()[id.value] self.msg.motor_cmd[id].q = lowstate.motor_state[id.value].q
# Both update different parts of self.msg, both call Write() # Initialize remote controller
self.publish_thread = None
self.ctrl_lock = threading.Lock()
self.publish_thread = threading.Thread(target=self._ctrl_motor_state)
self.publish_thread.daemon = True
self.publish_thread.start()
logger.warning("Arm control publish thread started")
self.remote_controller = self.RemoteController() self.remote_controller = self.RemoteController()
def _subscribe_motor_state(self): def _subscribe_motor_state(self): #polls robot state @ 250Hz
while True: while True:
start_time = time.time() start_time = time.time()
msg = self.lowstate_subscriber.Read() msg = self.lowstate_subscriber.Read()
@@ -186,143 +177,43 @@ class UnitreeG1(Robot):
sleep_time = max(0, (self.control_dt - all_t_elapsed)) # maintina constant control dt sleep_time = max(0, (self.control_dt - all_t_elapsed)) # maintina constant control dt
time.sleep(sleep_time) time.sleep(sleep_time)
def ctrl_dual_arm_go_home(self):
"""Move both the left and right arms of the robot to their home position by setting the target joint angles (q) and torques (tau) to zero."""
logger.info("[G1_29_ArmController] ctrl_dual_arm_go_home start...")
max_attempts = 100
current_attempts = 0
with self.ctrl_lock:
self.q_target = np.zeros(14)
# self.q_target[G1_29_JointIndex.kLeftElbow] = 0.5
# self.tauff_target = np.zeros(14)
tolerance = 0.05 # Tolerance threshold for joint angles to determine "close to zero", can be adjusted based on your motor's precision requirements
while current_attempts < max_attempts:
current_q = self.get_current_dual_arm_q()
if np.all(np.abs(current_q) < tolerance):
if self.motion_mode:
for weight in np.linspace(1, 0, num=101):
self.msg.motor_cmd[G1_29_JointIndex.kNotUsedJoint0].q = weight
time.sleep(0.02)
logger.info("[G1_29_ArmController] both arms have reached the home position.")
break
current_attempts += 1
time.sleep(0.05)
def clip_arm_q_target(self, target_q, velocity_limit):
current_q = self.get_current_dual_arm_q()
delta = target_q - current_q
motion_scale = np.max(np.abs(delta)) / (velocity_limit * self.control_dt)
cliped_arm_q_target = current_q + delta / max(motion_scale, 1.0)
return cliped_arm_q_target
def _ctrl_motor_state(self):
"""Arm control thread - publishes commands for arms only."""
while True:
start_time = time.time()
with self.ctrl_lock:
arm_q_target = self.q_target
arm_tauff_target = self.tauff_target
if self.simulation_mode:
cliped_arm_q_target = arm_q_target
else:
cliped_arm_q_target = self.clip_arm_q_target(
arm_q_target, velocity_limit=self.arm_velocity_limit
)
for idx, id in enumerate(G1_29_JointArmIndex):
self.msg.motor_cmd[id].q = cliped_arm_q_target[idx]
self.msg.motor_cmd[id].dq = 0
self.msg.motor_cmd[id].tau = arm_tauff_target[idx]
# Zero out specific joints when in simulation mode
if self.simulation_mode:
# Waist joints
self.msg.motor_cmd[G1_29_JointIndex.kWaistYaw].q = 0.0
self.msg.motor_cmd[G1_29_JointIndex.kWaistPitch].q = 0.0
# Wrist joints
self.msg.motor_cmd[G1_29_JointIndex.kLeftWristPitch].q = 0.0
self.msg.motor_cmd[G1_29_JointIndex.kLeftWristyaw].q = 0.0
self.msg.motor_cmd[G1_29_JointIndex.kRightWristPitch].q = 0.0
self.msg.motor_cmd[G1_29_JointIndex.kRightWristYaw].q = 0.0
self.msg.crc = self.crc.Crc(self.msg)
self.lowcmd_publisher.Write(self.msg)
current_time = time.time()
all_t_elapsed = current_time - start_time
sleep_time = max(0, (self.control_dt - all_t_elapsed))
time.sleep(sleep_time)
# logger.debug(f"arm_velocity_limit:{self.arm_velocity_limit}")
# logger.debug(f"sleep_time:{sleep_time}")
def ctrl_dual_arm(self, q_target, tauff_target):
"""Set control target values q & tau of the left and right arm motors."""
with self.ctrl_lock:
self.q_target = q_target
self.tauff_target = tauff_target
def get_mode_machine(self):
"""Return current dds mode machine."""
return self.lowstate_subscriber.Read().mode_machine
def get_current_motor_q(self):
"""Return current state q of all body motors."""
return np.array([self.lowstate_buffer.GetData().motor_state[id].q for id in G1_29_JointIndex])
def get_current_dual_arm_q(self):
"""Return current state q of the left and right arm motors."""
return np.array([self.lowstate_buffer.GetData().motor_state[id].q for id in G1_29_JointArmIndex])
def get_current_dual_arm_dq(self):
"""Return current state dq of the left and right arm motors."""
return np.array([self.lowstate_buffer.GetData().motor_state[id].dq for id in G1_29_JointArmIndex])
@cached_property @cached_property
def action_features(self) -> dict[str, type]: def action_features(self) -> dict[str, type]:
return {f"{G1_29_JointArmIndex(motor).name}.pos": float for motor in G1_29_JointArmIndex} return {f"{G1_29_JointIndex(motor).name}.pos": float for motor in G1_29_JointIndex}
def calibrate(self) -> None: def calibrate(self) -> None:#robot is already calibrated
self.calibration = json.load(open("src/lerobot/robots/unitree_g1/arm_calibration.json")) pass
self.calibrated = True
def configure(self) -> None: def configure(self) -> None:
pass pass
def connect(self, calibrate: bool = True) -> None: def connect(self, calibrate: bool = True) -> None: #connect to DDS
# Connect cameras ChannelFactoryInitialize(0)
for cam in self.cameras.values():
cam.connect() while not self.lowstate_buffer.GetData():
logger.info(f"{self} connected with {len(self.cameras)} camera(s).") time.sleep(0.1)
logger.warning("[UnitreeG1] Waiting to subscribe dds...")
logger.warning("[UnitreeG1] Subscribe dds ok.")
def disconnect(self): def disconnect(self):
# Disconnect cameras pass
for cam in self.cameras.values():
cam.disconnect()
# Close MuJoCo environment if in simulation mode
if self.simulation_mode and hasattr(self, "mujoco_env"):
logger.info("Closing MuJoCo environment...")
print(self.mujoco_env)
self.mujoco_env["hub_env"][0].envs[0].kill_sim()
logger.info(f"{self} disconnected.")
def get_observation(self) -> dict[str, Any]: def get_observation(self) -> dict[str, Any]:
return self.lowstate_buffer.GetData() return self.lowstate_buffer.GetData()
@property @property
def is_calibrated(self) -> bool: def is_calibrated(self) -> bool:
return self.calibrated return True
@property @property
def is_connected(self) -> bool: def is_connected(self) -> bool:
return all(cam.is_connected for cam in self.cameras.values()) if self.lowstate_buffer.GetData() is None:
return False
return True
@property @property
def _motors_ft(self) -> dict[str, type]: def _motors_ft(self) -> dict[str, type]:
return {f"{G1_29_JointArmIndex(motor).name}.pos": float for motor in G1_29_JointArmIndex} return {f"{G1_29_JointIndex(motor).name}.pos": float for motor in G1_29_JointIndex}
@property @property
def _cameras_ft(self) -> dict[str, tuple]: def _cameras_ft(self) -> dict[str, tuple]:
@@ -338,44 +229,7 @@ class UnitreeG1(Robot):
self.msg.crc = self.crc.Crc(action) self.msg.crc = self.crc.Crc(action)
self.lowcmd_publisher.Write(action) self.lowcmd_publisher.Write(action)
def reset_legs(self): def get_gravity_orientation(self, quaternion):#get gravity orientation from quaternion
"""Move robot legs to default standing position over 2 seconds (arms are not moved)."""
total_time = 2.0
num_step = int(total_time / self.control_dt)
# Only control legs, not arms
dof_idx = self.config.leg_joint2motor_idx
default_pos = np.array(self.config.default_leg_angles, dtype=np.float32)
dof_size = len(dof_idx)
# Get current lowstate
lowstate = self.lowstate_buffer.GetData()
if lowstate is None:
logger.error("Cannot get lowstate")
return
# 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] = lowstate.motor_state[dof_idx[i]].q
# Move legs to default pos
for i in range(num_step):
alpha = i / num_step
for j in range(dof_size):
motor_idx = dof_idx[j]
target_pos = default_pos[j]
self.msg.motor_cmd[motor_idx].q = init_dof_pos[j] * (1 - alpha) + target_pos * alpha
self.msg.motor_cmd[motor_idx].qd = 0
self.msg.motor_cmd[motor_idx].kp = self.kp[motor_idx]
self.msg.motor_cmd[motor_idx].kd = self.kd[motor_idx]
self.msg.motor_cmd[motor_idx].tau = 0
self.msg.crc = self.crc.Crc(self.msg)
self.lowcmd_publisher.Write(self.msg)
time.sleep(self.control_dt)
logger.info("Reached default position (legs only)")
def get_gravity_orientation(self, quaternion):
"""Get gravity orientation from quaternion.""" """Get gravity orientation from quaternion."""
qw = quaternion[0] qw = quaternion[0]
qx = quaternion[1] qx = quaternion[1]
@@ -388,7 +242,7 @@ class UnitreeG1(Robot):
gravity_orientation[2] = 1 - 2 * (qw * qw + qz * qz) gravity_orientation[2] = 1 - 2 * (qw * qw + qz * qz)
return gravity_orientation return gravity_orientation
def transform_imu_data(self, waist_yaw, waist_yaw_omega, imu_quat, imu_omega): def transform_imu_data(self, waist_yaw, waist_yaw_omega, imu_quat, imu_omega):#transform imu data from torso to pelvis frame
"""Transform IMU data from torso to pelvis frame.""" """Transform IMU data from torso to pelvis frame."""
RzWaist = R.from_euler("z", waist_yaw).as_matrix() RzWaist = R.from_euler("z", waist_yaw).as_matrix()
R_torso = R.from_quat([imu_quat[1], imu_quat[2], imu_quat[3], imu_quat[0]]).as_matrix() R_torso = R.from_quat([imu_quat[1], imu_quat[2], imu_quat[3], imu_quat[0]]).as_matrix()