From d7481f653ee25bf77d5f12ad30f4b280610ad28a Mon Sep 17 00:00:00 2001 From: Martino Russi Date: Wed, 26 Nov 2025 16:26:14 +0100 Subject: [PATCH] precommit --- .../robots/unitree_g1/config_unitree_g1.py | 46 +- src/lerobot/robots/unitree_g1/g1_utils.py | 67 ++ src/lerobot/robots/unitree_g1/robot_server.py | 11 +- src/lerobot/robots/unitree_g1/unitree_g1.py | 608 +++++++++--------- .../robots/unitree_g1/unitree_sdk2_socket.py | 15 +- 5 files changed, 416 insertions(+), 331 deletions(-) create mode 100644 src/lerobot/robots/unitree_g1/g1_utils.py diff --git a/src/lerobot/robots/unitree_g1/config_unitree_g1.py b/src/lerobot/robots/unitree_g1/config_unitree_g1.py index 2e986b8ea..6e711ed53 100644 --- a/src/lerobot/robots/unitree_g1/config_unitree_g1.py +++ b/src/lerobot/robots/unitree_g1/config_unitree_g1.py @@ -24,7 +24,7 @@ from ..config import RobotConfig @RobotConfig.register_subclass("unitree_g1") @dataclass class UnitreeG1Config(RobotConfig): - # id: str = "unitree_g1" + # id: str = "unitree_g1" simulation_mode: bool = False kp_high = 40.0 kd_high = 3.0 @@ -52,37 +52,49 @@ class UnitreeG1Config(RobotConfig): # This robot class ONLY uses sockets to communicate with a bridge on the Orin # Run 'python dds_to_socket.py' on the Orin first, then set this to the Orin's IP # Example: socket_host="192.168.123.164" (Orin's wlan0 IP) - socket_host: str | None = None# = "172.18.129.215" + socket_host: str | None = None # = "172.18.129.215" socket_port: int | None = None # Locomotion control locomotion_control: bool = True - #policy_path: str = "src/lerobot/robots/unitree_g1/assets/g1/locomotion/motion.pt" - policy_path: str = "src/lerobot/robots/unitree_g1/assets/g1/locomotion/GR00T-WholeBodyControl-Walk.onnx" - + # policy_path: str = "src/lerobot/robots/unitree_g1/assets/g1/locomotion/motion.pt" + policy_path: str | None = None + # Locomotion parameters (from g1.yaml) locomotion_control_dt: float = 0.02 - - leg_joint2motor_idx: list = field(default_factory=lambda: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]) - locomotion_kps: list = field(default_factory=lambda: [150, 150, 150, 300, 40, 40, 150, 150, 150, 300, 40, 40]) - locomotion_kds: list = field(default_factory=lambda: [2, 2, 2, 4, 2, 2, 2, 2, 2, 4, 2, 2]) - 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]) - arm_waist_joint2motor_idx: list = field(default_factory=lambda: [12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28]) - locomotion_arm_waist_kps: list = field(default_factory=lambda: [250, 250, 250, 100, 100, 50, 50, 20, 20, 20, 100, 100, 50, 50, 20, 20, 20]) - locomotion_arm_waist_kds: list = field(default_factory=lambda: [5, 5, 5, 2, 2, 2, 2, 1, 1, 1, 2, 2, 2, 2, 1, 1, 1]) - locomotion_arm_waist_target: list = field(default_factory=lambda: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) - + leg_joint2motor_idx: list = field(default_factory=lambda: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]) + locomotion_kps: list = field( + default_factory=lambda: [150, 150, 150, 300, 40, 40, 150, 150, 150, 300, 40, 40] + ) + locomotion_kds: list = field(default_factory=lambda: [2, 2, 2, 4, 2, 2, 2, 2, 2, 4, 2, 2]) + 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] + ) + + arm_waist_joint2motor_idx: list = field( + default_factory=lambda: [12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28] + ) + locomotion_arm_waist_kps: list = field( + default_factory=lambda: [250, 250, 250, 100, 100, 50, 50, 20, 20, 20, 100, 100, 50, 50, 20, 20, 20] + ) + locomotion_arm_waist_kds: list = field( + default_factory=lambda: [5, 5, 5, 2, 2, 2, 2, 1, 1, 1, 2, 2, 2, 2, 1, 1, 1] + ) + locomotion_arm_waist_target: list = field( + default_factory=lambda: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] + ) + ang_vel_scale: float = 0.25 dof_pos_scale: float = 1.0 dof_vel_scale: float = 0.05 locomotion_action_scale: float = 0.25 cmd_scale: list = field(default_factory=lambda: [2.0, 2.0, 0.25]) - + # GR00T-specific scaling (different from regular locomotion!) groot_ang_vel_scale: float = 0.25 # GR00T uses 0.5, not 0.25 groot_cmd_scale: list = field(default_factory=lambda: [2.0, 2.0, 0.25]) # yaw is 0.5 for GR00T num_locomotion_actions: int = 12 num_locomotion_obs: int = 47 max_cmd: list = field(default_factory=lambda: [0.8, 0.5, 1.57]) - locomotion_imu_type: str = "pelvis" # "torso" or "pelvis" \ No newline at end of file + locomotion_imu_type: str = "pelvis" # "torso" or "pelvis" diff --git a/src/lerobot/robots/unitree_g1/g1_utils.py b/src/lerobot/robots/unitree_g1/g1_utils.py new file mode 100644 index 000000000..0e6fa89fb --- /dev/null +++ b/src/lerobot/robots/unitree_g1/g1_utils.py @@ -0,0 +1,67 @@ +from enum import IntEnum + +class G1_29_JointArmIndex(IntEnum): + # Left arm + kLeftShoulderPitch = 15 + kLeftShoulderRoll = 16 + kLeftShoulderYaw = 17 + kLeftElbow = 18 + kLeftWristRoll = 19 + kLeftWristPitch = 20 + kLeftWristyaw = 21 + + # Right arm + kRightShoulderPitch = 22 + kRightShoulderRoll = 23 + kRightShoulderYaw = 24 + kRightElbow = 25 + kRightWristRoll = 26 + kRightWristPitch = 27 + kRightWristYaw = 28 + +class G1_29_JointIndex(IntEnum): + # Left leg + kLeftHipPitch = 0 + kLeftHipRoll = 1 + kLeftHipYaw = 2 + kLeftKnee = 3 + kLeftAnklePitch = 4 + kLeftAnkleRoll = 5 + + # Right leg + kRightHipPitch = 6 + kRightHipRoll = 7 + kRightHipYaw = 8 + kRightKnee = 9 + kRightAnklePitch = 10 + kRightAnkleRoll = 11 + + kWaistYaw = 12 #we're c + kWaistRoll = 13 + kWaistPitch = 14 + + # Left arm + kLeftShoulderPitch = 15 + kLeftShoulderRoll = 16 + kLeftShoulderYaw = 17 + kLeftElbow = 18 + kLeftWristRoll = 19 + kLeftWristPitch = 20 + kLeftWristyaw = 21 + + # Right arm + kRightShoulderPitch = 22 + kRightShoulderRoll = 23 + kRightShoulderYaw = 24 + kRightElbow = 25 + kRightWristRoll = 26 + kRightWristPitch = 27 + kRightWristYaw = 28 + + # not used + kNotUsedJoint0 = 29 + kNotUsedJoint1 = 30 + kNotUsedJoint2 = 31 + kNotUsedJoint3 = 32 + kNotUsedJoint4 = 33 + kNotUsedJoint5 = 34 \ No newline at end of file diff --git a/src/lerobot/robots/unitree_g1/robot_server.py b/src/lerobot/robots/unitree_g1/robot_server.py index 42a93636c..0a3e70feb 100644 --- a/src/lerobot/robots/unitree_g1/robot_server.py +++ b/src/lerobot/robots/unitree_g1/robot_server.py @@ -1,20 +1,19 @@ #!/usr/bin/env python3 -import time import pickle import threading +import time import zmq - -from unitree_sdk2py.core.channel import ChannelPublisher, ChannelSubscriber, ChannelFactoryInitialize +from unitree_sdk2py.comm.motion_switcher.motion_switcher_client import MotionSwitcherClient +from unitree_sdk2py.core.channel import ChannelFactoryInitialize, ChannelPublisher, ChannelSubscriber from unitree_sdk2py.idl.unitree_hg.msg.dds_ import LowCmd_ as hg_LowCmd, LowState_ as hg_LowState from unitree_sdk2py.utils.crc import CRC -from unitree_sdk2py.comm.motion_switcher.motion_switcher_client import MotionSwitcherClient kTopicLowCommand_Debug = "rt/lowcmd" kTopicLowState = "rt/lowstate" -LOWCMD_PORT = 6000 # laptop -> robot -LOWSTATE_PORT = 6001 # robot -> laptop +LOWCMD_PORT = 6000 # laptop -> robot +LOWSTATE_PORT = 6001 # robot -> laptop def state_forward_loop(lowstate_sub, lowstate_sock, state_period: float): diff --git a/src/lerobot/robots/unitree_g1/unitree_g1.py b/src/lerobot/robots/unitree_g1/unitree_g1.py index 8bf9073ee..cb19acbed 100644 --- a/src/lerobot/robots/unitree_g1/unitree_g1.py +++ b/src/lerobot/robots/unitree_g1/unitree_g1.py @@ -1,46 +1,33 @@ -import logging -import time -import struct -from functools import cached_property -from typing import Any -from pathlib import Path - -from lerobot.cameras.utils import make_cameras_from_configs - import json -from ..robot import Robot -from .config_unitree_g1 import UnitreeG1Config - -import numpy as np +import logging +import select +import struct +import sys +import termios import threading import time -from enum import IntEnum -import sys -import select -import termios import tty from collections import deque +from functools import cached_property +from pathlib import Path +from typing import Any -from typing import Union import numpy as np -import time -import torch import onnxruntime as ort - -from unitree_sdk2py.idl.unitree_hg.msg.dds_ import LowCmd_ as hg_LowCmd, LowState_ as hg_LowState # idl for g1, h1_2 -from unitree_sdk2py.idl.default import unitree_hg_msg_dds__LowCmd_ -from unitree_sdk2py.utils.crc import CRC -from unitree_sdk2py.comm.motion_switcher.motion_switcher_client import ( - MotionSwitcherClient, -) - -from lerobot.envs.factory import make_env -from scipy.spatial.transform import Rotation as R - -import struct - - import torch +from scipy.spatial.transform import Rotation as R +from unitree_sdk2py.idl.default import unitree_hg_msg_dds__LowCmd_ +from unitree_sdk2py.idl.unitree_hg.msg.dds_ import ( + LowCmd_ as hg_LowCmd, + LowState_ as hg_LowState, +) # idl for g1, h1_2 +from unitree_sdk2py.utils.crc import CRC + +from lerobot.cameras.utils import make_cameras_from_configs +from lerobot.robots.unitree_g1.g1_utils import G1_29_JointArmIndex, G1_29_JointIndex + +from ..robot import Robot +from .config_unitree_g1 import UnitreeG1Config logger = logging.getLogger(__name__) @@ -76,6 +63,7 @@ class G1_29_LowState: self.imu_state = IMUState() self.wireless_remote = None # Raw wireless remote data + class DataBuffer: def __init__(self): self.data = None @@ -89,14 +77,14 @@ class DataBuffer: with self.lock: self.data = data -class UnitreeG1(Robot): +class UnitreeG1(Robot): config_class = UnitreeG1Config name = "unitree_g1" def __init__(self, config: UnitreeG1Config): super().__init__(config) - + logger.info("Initialize UnitreeG1...") self.config = config @@ -128,9 +116,12 @@ class UnitreeG1(Robot): self.calibrated = False + from lerobot.robots.unitree_g1.unitree_sdk2_socket import ( + ChannelFactoryInitialize, + ChannelPublisher, + ChannelSubscriber, + ) # dds - from lerobot.robots.unitree_g1.unitree_sdk2_socket import ChannelPublisher, ChannelSubscriber, ChannelFactoryInitialize # dds - ChannelFactoryInitialize(0) # Always use debug mode (direct motor control) @@ -163,7 +154,7 @@ class UnitreeG1(Robot): logger.info(f"Current two arms motor state q:\n{self.get_current_dual_arm_q()}\n") logger.info("Lock all joints except two arms...\n") - arm_indices = set(member.value for member in G1_29_JointArmIndex) + arm_indices = {member.value for member in G1_29_JointArmIndex} for id in G1_29_JointIndex: self.msg.motor_cmd[id].mode = 1 if id.value in arm_indices: @@ -181,19 +172,19 @@ class UnitreeG1(Robot): self.msg.motor_cmd[id].kp = self.kp_high self.msg.motor_cmd[id].kd = self.kd_high self.msg.motor_cmd[id].q = self.all_motor_q[id] - #print current motor q, kp, kd + # print current motor q, kp, kd - logger.warning("Lock OK!\n") #motors are not locked x + logger.warning("Lock OK!\n") # motors are not locked x # for i in range(10000): # print(self.get_current_motor_q()) - # time.sleep(0.05) - + # time.sleep(0.05) + # Initialize control flags BEFORE starting threads self.keyboard_thread = None self.keyboard_running = False self.locomotion_thread = None self.locomotion_running = False - + # Initialize publish thread for arm control # Note: This thread runs alongside locomotion thread # - Arm thread: controls arms (indices 15-28) @@ -213,24 +204,24 @@ class UnitreeG1(Robot): if config.locomotion_control: if config.policy_path is None: raise ValueError("locomotion_control is True but policy_path is not set") - + logger.warning(f"Loading locomotion policy from {config.policy_path}") - + # Check file extension and load accordingly - if config.policy_path.endswith('.pt'): + if config.policy_path.endswith(".pt"): logger.warning("Detected TorchScript (.pt) policy") self.policy = torch.jit.load(config.policy_path) - self.policy_type = 'torchscript' + self.policy_type = "torchscript" logger.info("TorchScript policy loaded successfully") - elif config.policy_path.endswith('.onnx'): + elif config.policy_path.endswith(".onnx"): logger.warning("Detected ONNX (.onnx) policy") - + # For GR00T-style policies, load both Balance and Walk policies # Balance policy for standing (low velocity commands) # Walk policy for locomotion (high velocity commands) - balance_policy_path = config.policy_path.replace('Walk.onnx', 'Balance.onnx') + balance_policy_path = config.policy_path.replace("Walk.onnx", "Balance.onnx") walk_policy_path = config.policy_path - + if Path(balance_policy_path).exists() and Path(walk_policy_path).exists(): logger.info("Loading dual-policy system (Balance + Walk)") self.policy_balance = ort.InferenceSession(balance_policy_path) @@ -238,8 +229,12 @@ class UnitreeG1(Robot): self.policy = None # Not used when dual policies are loaded logger.info(f"Balance policy loaded from: {balance_policy_path}") logger.info(f"Walk policy loaded from: {walk_policy_path}") - logger.info(f"ONNX input: {self.policy_balance.get_inputs()[0].name}, shape: {self.policy_balance.get_inputs()[0].shape}") - logger.info(f"ONNX output: {self.policy_balance.get_outputs()[0].name}, shape: {self.policy_balance.get_outputs()[0].shape}") + logger.info( + f"ONNX input: {self.policy_balance.get_inputs()[0].name}, shape: {self.policy_balance.get_inputs()[0].shape}" + ) + logger.info( + f"ONNX output: {self.policy_balance.get_outputs()[0].name}, shape: {self.policy_balance.get_outputs()[0].shape}" + ) else: # Fallback to single policy logger.info("Loading single ONNX policy") @@ -247,13 +242,19 @@ class UnitreeG1(Robot): self.policy_balance = None self.policy_walk = None logger.info("ONNX policy loaded successfully") - logger.info(f"ONNX input: {self.policy.get_inputs()[0].name}, shape: {self.policy.get_inputs()[0].shape}") - logger.info(f"ONNX output: {self.policy.get_outputs()[0].name}, shape: {self.policy.get_outputs()[0].shape}") - - self.policy_type = 'onnx' + logger.info( + f"ONNX input: {self.policy.get_inputs()[0].name}, shape: {self.policy.get_inputs()[0].shape}" + ) + logger.info( + f"ONNX output: {self.policy.get_outputs()[0].name}, shape: {self.policy.get_outputs()[0].shape}" + ) + + self.policy_type = "onnx" else: - raise ValueError(f"Unsupported policy format: {config.policy_path}. Only .pt (TorchScript) and .onnx (ONNX) are supported.") - + raise ValueError( + f"Unsupported policy format: {config.policy_path}. Only .pt (TorchScript) and .onnx (ONNX) are supported." + ) + # Initialize locomotion variables self.remote_controller = self.RemoteController() self.locomotion_counter = 0 @@ -262,9 +263,9 @@ class UnitreeG1(Robot): self.locomotion_action = np.zeros(config.num_locomotion_actions, dtype=np.float32) self.locomotion_obs = np.zeros(config.num_locomotion_obs, dtype=np.float32) self.locomotion_cmd = np.array([0.0, 0.0, 0.0], dtype=np.float32) - + # GR00T-specific variables (for ONNX policies with 29 joints) - if self.policy_type == 'onnx': + if self.policy_type == "onnx": self.groot_qj_all = np.zeros(29, dtype=np.float32) # All 29 joints self.groot_dqj_all = np.zeros(29, dtype=np.float32) self.groot_action = np.zeros(15, dtype=np.float32) # 15D action (legs + waist) @@ -273,18 +274,18 @@ class UnitreeG1(Robot): self.groot_obs_stacked = np.zeros(516, dtype=np.float32) # 86D × 6 = 516D stacked observation self.groot_height_cmd = 0.74 # Default base height self.groot_orientation_cmd = np.array([0.0, 0.0, 0.0], dtype=np.float32) # roll, pitch, yaw - + # Initialize history with zeros for _ in range(6): self.groot_obs_history.append(np.zeros(86, dtype=np.float32)) - + # Start keyboard controls if in simulation mode if self.simulation_mode: logger.info("Starting keyboard controls for simulation...") self.start_keyboard_controls() - + # Use different init based on policy type - if self.policy_type == 'onnx': + if self.policy_type == "onnx": self.init_groot_locomotion() else: self.init_locomotion() @@ -292,7 +293,6 @@ class UnitreeG1(Robot): # Even without locomotion, provide keyboard feedback in sim logger.info("Simulation mode active (locomotion disabled)") - logger.info("Initialize G1 OK!\n") def _subscribe_motor_state(self): @@ -301,31 +301,52 @@ class UnitreeG1(Robot): msg = self.lowstate_subscriber.Read() if msg is not None: lowstate = G1_29_LowState() - + # Capture motor states for id in range(G1_29_Num_Motors): lowstate.motor_state[id].q = msg.motor_state[id].q lowstate.motor_state[id].dq = msg.motor_state[id].dq lowstate.motor_state[id].tau_est = msg.motor_state[id].tau_est lowstate.motor_state[id].temperature = msg.motor_state[id].temperature - + # Capture IMU state lowstate.imu_state.quaternion = list(msg.imu_state.quaternion) lowstate.imu_state.gyroscope = list(msg.imu_state.gyroscope) lowstate.imu_state.accelerometer = list(msg.imu_state.accelerometer) lowstate.imu_state.rpy = list(msg.imu_state.rpy) lowstate.imu_state.temperature = msg.imu_state.temperature - + # Capture wireless remote data lowstate.wireless_remote = msg.wireless_remote - + self.lowstate_buffer.SetData(lowstate) current_time = time.time() all_t_elapsed = current_time - start_time - 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) - + + 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() @@ -338,7 +359,7 @@ class UnitreeG1(Robot): """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 @@ -346,7 +367,9 @@ class UnitreeG1(Robot): 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) + 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] @@ -396,7 +419,6 @@ class UnitreeG1(Robot): """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]) - def _Is_weak_motor(self, motor_index): weak_motors = [ G1_29_JointIndex.kLeftAnklePitch.value, @@ -430,7 +452,7 @@ class UnitreeG1(Robot): return {f"{G1_29_JointArmIndex(motor).name}.pos": float for motor in G1_29_JointArmIndex} def calibrate(self) -> None: - self.calibration = json.load(open('src/lerobot/robots/unitree_g1/arm_calibration.json')) + self.calibration = json.load(open("src/lerobot/robots/unitree_g1/arm_calibration.json")) self.calibrated = True def configure(self) -> None: @@ -446,19 +468,19 @@ class UnitreeG1(Robot): # Disconnect cameras for cam in self.cameras.values(): cam.disconnect() - + # Close MuJoCo environment if in simulation mode - if self.simulation_mode and hasattr(self, 'mujoco_env'): + 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_full_robot_state(self) -> dict[str, Any]: """ Get full robot state including IMU and extended motor data. - + Returns: dict with keys: - 'imu': dict containing IMU data (quaternion, gyroscope, accelerometer, rpy, temperature) @@ -467,44 +489,51 @@ class UnitreeG1(Robot): lowstate = self.lowstate_buffer.GetData() if lowstate is None: raise RuntimeError("No robot state available. Is the robot connected?") - + # Extract IMU data imu_data = { - 'quaternion': lowstate.imu_state.quaternion, # [w, x, y, z] - 'gyroscope': lowstate.imu_state.gyroscope, # [x, y, z] rad/s - 'accelerometer': lowstate.imu_state.accelerometer, # [x, y, z] m/s² - 'rpy': lowstate.imu_state.rpy, # [roll, pitch, yaw] rad - 'temperature': lowstate.imu_state.temperature, # °C + "quaternion": lowstate.imu_state.quaternion, # [w, x, y, z] + "gyroscope": lowstate.imu_state.gyroscope, # [x, y, z] rad/s + "accelerometer": lowstate.imu_state.accelerometer, # [x, y, z] m/s² + "rpy": lowstate.imu_state.rpy, # [roll, pitch, yaw] rad + "temperature": lowstate.imu_state.temperature, # °C } - + # Extract motor data motors_data = [] for i in range(G1_29_Num_Motors): motor = lowstate.motor_state[i] - motors_data.append({ - 'id': i, - 'q': motor.q, # position (rad) - 'dq': motor.dq, # velocity (rad/s) - 'tau_est': motor.tau_est, # estimated torque (Nm) - 'temperature': motor.temperature[0] if isinstance(motor.temperature, (list, tuple)) else motor.temperature, # °C - }) - + motors_data.append( + { + "id": i, + "q": motor.q, # position (rad) + "dq": motor.dq, # velocity (rad/s) + "tau_est": motor.tau_est, # estimated torque (Nm) + "temperature": motor.temperature[0] + if isinstance(motor.temperature, (list, tuple)) + else motor.temperature, # °C + } + ) + return { - 'imu': imu_data, - 'motors': motors_data, + "imu": imu_data, + "motors": motors_data, } def get_observation(self) -> dict[str, Any]: obs_array = self.get_current_dual_arm_q() - obs_dict = {f"{G1_29_JointArmIndex(motor).name}.pos": val for motor, val in zip(G1_29_JointArmIndex, obs_array, strict=True)} - + obs_dict = { + f"{G1_29_JointArmIndex(motor).name}.pos": val + for motor, val in zip(G1_29_JointArmIndex, obs_array, strict=True) + } + # Capture images from cameras for cam_key, cam in self.cameras.items(): start = time.perf_counter() obs_dict[cam_key] = cam.async_read() dt_ms = (time.perf_counter() - start) * 1e3 logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") - + return obs_dict @property @@ -514,11 +543,11 @@ class UnitreeG1(Robot): @property def is_connected(self) -> bool: return all(cam.is_connected for cam in self.cameras.values()) - + @property def _motors_ft(self) -> dict[str, type]: return {f"{G1_29_JointArmIndex(motor).name}.pos": float for motor in G1_29_JointArmIndex} - + @property def _cameras_ft(self) -> dict[str, tuple]: return { @@ -530,22 +559,24 @@ class UnitreeG1(Robot): return {**self._motors_ft, **self._cameras_ft} def send_action(self, action: dict[str, Any]) -> dict[str, Any]: - #need a any to any teleoperator solution. i wanna teleoperate a horse with a shoe. action - #to action mapping, when you do teleoperate. the keys that are left empty are just set to 0 - #also what would be fun is finding all sorts of robots and adding them to lerobot, see if people do the same. - #then teleop them wiuth the glove hehe - #then we get ALL THE DATA + # need a any to any teleoperator solution. i wanna teleoperate a horse with a shoe. action + # to action mapping, when you do teleoperate. the keys that are left empty are just set to 0 + # also what would be fun is finding all sorts of robots and adding them to lerobot, see if people do the same. + # then teleop them wiuth the glove hehe + # then we get ALL THE DATA if self.is_calibrated: uncalibrated_action = action.copy() action = self.invert_calibration(action) - #if an action was 0.5 write 0 in its place + # if an action was 0.5 write 0 in its place for key, value in uncalibrated_action.items(): if value == 0.5: action[key] = 0.0 - #check if action is within bounds + # check if action is within bounds for key, value in action.items(): if value < self.calibration[key]["range_min"] or value > self.calibration[key]["range_max"]: - raise ValueError(f"Action value {value} for {key} is out of bounds, actions are not normalized") + raise ValueError( + f"Action value {value} for {key} is out of bounds, actions are not normalized" + ) if self.freeze_body: arm_joint_indices = set(range(15, 29)) # 15–28 are arms for jid in G1_29_JointIndex: @@ -556,13 +587,13 @@ class UnitreeG1(Robot): self.msg.motor_cmd[jid].tau = 0.0 action_np = np.stack([v for v in action.values()]) - #action_np is just zeros - #action_np = np.zeros(14) - #print(action_np) - #exit() + # action_np is just zeros + # action_np = np.zeros(14) + # print(action_np) + # exit() tau = np.zeros(14) - + self.ctrl_dual_arm(action_np, tau) def apply_calibration(self, action: dict[str, float]) -> dict[str, float]: @@ -585,7 +616,6 @@ class UnitreeG1(Robot): return calibrated - def invert_calibration(self, action: dict[str, float]) -> dict[str, float]: """Map [0, 1] actions back to motor ranges.""" calibrated = {} @@ -597,7 +627,7 @@ class UnitreeG1(Robot): # inverse mapping real_val = mn + value * (mx - mn) - + # Round to 5 decimal places to avoid floating point precision issues calibrated[key] = round(real_val, 5) @@ -640,25 +670,25 @@ class UnitreeG1(Robot): logger.info("Moving legs to default locomotion pos.") total_time = 2.0 num_step = int(total_time / self.config.locomotion_control_dt) - + # Only control legs, not arms dof_idx = self.config.leg_joint2motor_idx kps = self.config.locomotion_kps kds = self.config.locomotion_kds 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 for locomotion") 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 @@ -678,7 +708,7 @@ class UnitreeG1(Robot): def locomotion_default_pos_state(self): """Hold default leg position for 2 seconds (arms are not controlled).""" logger.info("Enter default pos state - holding legs for 2 seconds") - + # Only control legs, not arms for i in range(len(self.config.leg_joint2motor_idx)): motor_idx = self.config.leg_joint2motor_idx[i] @@ -687,7 +717,7 @@ class UnitreeG1(Robot): self.msg.motor_cmd[motor_idx].kp = self.config.locomotion_kps[i] self.msg.motor_cmd[motor_idx].kd = self.config.locomotion_kds[i] self.msg.motor_cmd[motor_idx].tau = 0 - + # Hold leg position for 2 seconds hold_time = 2.0 num_steps = int(hold_time / self.config.locomotion_control_dt) @@ -696,7 +726,6 @@ class UnitreeG1(Robot): self.lowcmd_publisher.Write(self.msg) time.sleep(self.config.locomotion_control_dt) logger.info("Finished holding default leg position") - class RemoteController: def __init__(self): @@ -715,21 +744,21 @@ class UnitreeG1(Robot): self.rx = struct.unpack("f", data[8:12])[0] self.ry = struct.unpack("f", data[12:16])[0] self.ly = struct.unpack("f", data[20:24])[0] - + def locomotion_get_gravity_orientation(self, quaternion): """Get gravity orientation from quaternion.""" qw = quaternion[0] qx = quaternion[1] qy = quaternion[2] qz = quaternion[3] - + gravity_orientation = np.zeros(3) gravity_orientation[0] = 2 * (-qz * qx + qw * qy) gravity_orientation[1] = -2 * (qz * qy + qw * qx) gravity_orientation[2] = 1 - 2 * (qw * qw + qz * qz) - + return gravity_orientation - + def locomotion_transform_imu_data(self, waist_yaw, waist_yaw_omega, imu_quat, imu_omega): """Transform IMU data from torso to pelvis frame.""" RzWaist = R.from_euler("z", waist_yaw).as_matrix() @@ -737,16 +766,16 @@ class UnitreeG1(Robot): R_pelvis = np.dot(R_torso, RzWaist.T) w = np.dot(RzWaist, imu_omega[0]) - np.array([0, 0, waist_yaw_omega]) return R.from_matrix(R_pelvis).as_quat()[[3, 0, 1, 2]], w - + def locomotion_run(self): """Main locomotion policy loop - runs policy and sends leg commands.""" self.locomotion_counter += 1 - + # Get current lowstate lowstate = self.lowstate_buffer.GetData() if lowstate is None: return - + # Update remote controller from lowstate if lowstate.wireless_remote is not None: self.remote_controller.set(lowstate.wireless_remote) @@ -756,22 +785,22 @@ class UnitreeG1(Robot): self.remote_controller.ly = 0.0 self.remote_controller.rx = 0.0 self.remote_controller.ry = 0.0 - + # Get the current joint position and velocity (LEGS ONLY) for i in range(len(self.config.leg_joint2motor_idx)): self.qj[i] = lowstate.motor_state[self.config.leg_joint2motor_idx[i]].q self.dqj[i] = lowstate.motor_state[self.config.leg_joint2motor_idx[i]].dq - + # Get IMU data quat = lowstate.imu_state.quaternion ang_vel = np.array([lowstate.imu_state.gyroscope], dtype=np.float32) - + if self.config.locomotion_imu_type == "torso": # Transform IMU data from torso to pelvis frame waist_yaw = lowstate.motor_state[self.config.arm_waist_joint2motor_idx[0]].q waist_yaw_omega = lowstate.motor_state[self.config.arm_waist_joint2motor_idx[0]].dq quat, ang_vel = self.locomotion_transform_imu_data(waist_yaw, waist_yaw_omega, quat, ang_vel) - + # Create observation gravity_orientation = self.locomotion_get_gravity_orientation(quat) qj_obs = self.qj.copy() @@ -779,53 +808,60 @@ class UnitreeG1(Robot): qj_obs = (qj_obs - np.array(self.config.default_leg_angles)) * self.config.dof_pos_scale dqj_obs = dqj_obs * self.config.dof_vel_scale ang_vel = ang_vel * self.config.ang_vel_scale - + # Calculate phase period = 0.8 count = self.locomotion_counter * self.config.locomotion_control_dt phase = count % period / period sin_phase = np.sin(2 * np.pi * phase) cos_phase = np.cos(2 * np.pi * phase) - + # Get velocity commands from remote controller (only if NOT in simulation mode) # In simulation mode, keyboard controls set self.locomotion_cmd directly if not self.simulation_mode: self.locomotion_cmd[0] = self.remote_controller.ly self.locomotion_cmd[1] = self.remote_controller.lx * -1 self.locomotion_cmd[2] = self.remote_controller.rx * -1 - + # Debug: print remote controller values every 50 iterations (~1 second at 50Hz) if self.locomotion_counter % 50 == 0: - logger.debug(f"Remote controller - lx:{self.remote_controller.lx:.2f}, ly:{self.remote_controller.ly:.2f}, rx:{self.remote_controller.rx:.2f}") - + logger.debug( + f"Remote controller - lx:{self.remote_controller.lx:.2f}, ly:{self.remote_controller.ly:.2f}, rx:{self.remote_controller.rx:.2f}" + ) + # Build observation vector num_actions = self.config.num_locomotion_actions self.locomotion_obs[:3] = ang_vel self.locomotion_obs[3:6] = gravity_orientation - self.locomotion_obs[6:9] = self.locomotion_cmd * np.array(self.config.cmd_scale) * np.array(self.config.max_cmd) + self.locomotion_obs[6:9] = ( + self.locomotion_cmd * np.array(self.config.cmd_scale) * np.array(self.config.max_cmd) + ) self.locomotion_obs[9 : 9 + num_actions] = qj_obs self.locomotion_obs[9 + num_actions : 9 + num_actions * 2] = dqj_obs self.locomotion_obs[9 + num_actions * 2 : 9 + num_actions * 3] = self.locomotion_action self.locomotion_obs[9 + num_actions * 3] = sin_phase self.locomotion_obs[9 + num_actions * 3 + 1] = cos_phase - + # Get action from policy network obs_tensor = torch.from_numpy(self.locomotion_obs).unsqueeze(0) - - if self.policy_type == 'torchscript': + + if self.policy_type == "torchscript": # TorchScript inference self.locomotion_action = self.policy(obs_tensor).detach().numpy().squeeze() - elif self.policy_type == 'onnx': + elif self.policy_type == "onnx": # ONNX inference ort_inputs = {self.policy.get_inputs()[0].name: obs_tensor.cpu().numpy()} ort_outs = self.policy.run(None, ort_inputs) self.locomotion_action = ort_outs[0].squeeze() else: raise ValueError(f"Unknown policy type: {self.policy_type}") - + # Transform action to target joint positions - target_dof_pos = np.array(self.config.default_leg_angles) + self.locomotion_action * self.config.locomotion_action_scale - + target_dof_pos = ( + np.array(self.config.default_leg_angles) + + self.locomotion_action * self.config.locomotion_action_scale + ) + # Send commands to LEG motors only for i in range(len(self.config.leg_joint2motor_idx)): motor_idx = self.config.leg_joint2motor_idx[i] @@ -834,7 +870,7 @@ class UnitreeG1(Robot): self.msg.motor_cmd[motor_idx].kp = self.config.locomotion_kps[i] self.msg.motor_cmd[motor_idx].kd = self.config.locomotion_kds[i] self.msg.motor_cmd[motor_idx].tau = 0 - + # Hold WAIST motors at 0 (indices 12, 13, 14 = WaistYaw, WaistRoll, WaistPitch) waist_indices = self.config.arm_waist_joint2motor_idx[:3] # First 3 are waist for i, motor_idx in enumerate(waist_indices): @@ -843,24 +879,24 @@ class UnitreeG1(Robot): self.msg.motor_cmd[motor_idx].kp = self.config.locomotion_arm_waist_kps[i] self.msg.motor_cmd[motor_idx].kd = self.config.locomotion_arm_waist_kds[i] self.msg.motor_cmd[motor_idx].tau = 0 - + # Send command self.msg.crc = self.crc.Crc(self.msg) self.lowcmd_publisher.Write(self.msg) - + def groot_locomotion_run(self): """GR00T-style locomotion policy loop for ONNX policies - reads all 29 joints, outputs 15D action.""" self.locomotion_counter += 1 - + # Get current lowstate lowstate = self.lowstate_buffer.GetData() if lowstate is None: return - + # Update remote controller from lowstate if lowstate.wireless_remote is not None: self.remote_controller.set(lowstate.wireless_remote) - + # R1/R2 buttons for height control on real robot (button indices 4 and 5) if self.remote_controller.button[0]: # R1 - raise height self.groot_height_cmd += 0.001 # Small increment per timestep (~0.05m per second at 50Hz) @@ -874,23 +910,25 @@ class UnitreeG1(Robot): self.remote_controller.ly = 0.0 self.remote_controller.rx = 0.0 self.remote_controller.ry = 0.0 - + # Get ALL 29 joint positions and velocities for i in range(29): self.groot_qj_all[i] = lowstate.motor_state[i].q self.groot_dqj_all[i] = lowstate.motor_state[i].dq - + # Get IMU data quat = lowstate.imu_state.quaternion ang_vel = np.array(lowstate.imu_state.gyroscope, dtype=np.float32) - + # Transform IMU if using torso IMU if self.config.locomotion_imu_type == "torso": waist_yaw = lowstate.motor_state[12].q # Waist yaw index waist_yaw_omega = lowstate.motor_state[12].dq - quat, ang_vel_3d = self.locomotion_transform_imu_data(waist_yaw, waist_yaw_omega, quat, np.array([ang_vel])) + quat, ang_vel_3d = self.locomotion_transform_imu_data( + waist_yaw, waist_yaw_omega, quat, np.array([ang_vel]) + ) ang_vel = ang_vel_3d.flatten() - + # Create observation gravity_orientation = self.locomotion_get_gravity_orientation(quat) joints_to_zero_obs = [12, 14, 20, 21, 27, 28] # Note: NOT 13 (waist roll exists) @@ -900,27 +938,58 @@ class UnitreeG1(Robot): # Scale joint positions and velocities qj_obs = self.groot_qj_all.copy() dqj_obs = self.groot_dqj_all.copy() - + # Subtract default angles for legs + waist (15 joints) # GR00T default_angles: [-0.1, 0.0, 0.0, 0.3, -0.2, 0.0, -0.1, 0.0, 0.0, 0.3, -0.2, 0.0, 0.0, 0.0, 0.0] - groot_default_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 - 0.0, 0.0, 0.0, # waist - 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, # left arm (zeroed) - 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float32) # right arm (zeroed) - + groot_default_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 + 0.0, + 0.0, + 0.0, # waist + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, # left arm (zeroed) + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + ], + dtype=np.float32, + ) # right arm (zeroed) + qj_obs = (qj_obs - groot_default_angles) * self.config.dof_pos_scale dqj_obs = dqj_obs * self.config.dof_vel_scale ang_vel_scaled = ang_vel * self.config.groot_ang_vel_scale # Use GR00T-specific scaling! - + # Get velocity commands (keyboard or remote) if not self.simulation_mode: self.locomotion_cmd[0] = self.remote_controller.ly self.locomotion_cmd[1] = self.remote_controller.lx * -1 self.locomotion_cmd[2] = self.remote_controller.rx * -1 - + # Build 86D single frame observation (GR00T format) - self.groot_obs_single[:3] = self.locomotion_cmd * np.array(self.config.groot_cmd_scale) # cmd - use GR00T scaling! + self.groot_obs_single[:3] = self.locomotion_cmd * np.array( + self.config.groot_cmd_scale + ) # cmd - use GR00T scaling! self.groot_obs_single[3] = self.groot_height_cmd # height_cmd self.groot_obs_single[4:7] = self.groot_orientation_cmd # roll, pitch, yaw cmd self.groot_obs_single[7:10] = ang_vel_scaled # angular velocity @@ -928,19 +997,19 @@ class UnitreeG1(Robot): self.groot_obs_single[13:42] = qj_obs # joint positions (29D) self.groot_obs_single[42:71] = dqj_obs # joint velocities (29D) self.groot_obs_single[71:86] = self.groot_action # previous actions (15D) - + # Add to history and stack observations (6 frames × 86D = 516D) self.groot_obs_history.append(self.groot_obs_single.copy()) - + # Stack all 6 frames into 516D vector for i, obs_frame in enumerate(self.groot_obs_history): start_idx = i * 86 end_idx = start_idx + 86 self.groot_obs_stacked[start_idx:end_idx] = obs_frame - + # 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) if self.policy_balance is not None and self.policy_walk is not None: # Dual-policy mode: switch between Balance and Walk @@ -954,20 +1023,22 @@ class UnitreeG1(Robot): else: # Single policy mode (fallback) selected_policy = self.policy - + 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) # This ensures action history in observations matches what's actually executed self.groot_action[12] = 0.0 # Waist yaw - self.groot_action[13] = 0.0 # Waist roll + 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, but waist actions are zeroed) - target_dof_pos_15 = groot_default_angles[:15] + self.groot_action * self.config.locomotion_action_scale - + target_dof_pos_15 = ( + groot_default_angles[:15] + self.groot_action * self.config.locomotion_action_scale + ) + # Send commands to LEG motors (0-11) for i in range(12): motor_idx = i @@ -976,17 +1047,19 @@ class UnitreeG1(Robot): self.msg.motor_cmd[motor_idx].kp = self.config.locomotion_kps[i] self.msg.motor_cmd[motor_idx].kd = self.config.locomotion_kds[i] self.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 # In the 15D action self.msg.motor_cmd[waist_roll_idx].q = target_dof_pos_15[waist_roll_action_idx] self.msg.motor_cmd[waist_roll_idx].qd = 0 - self.msg.motor_cmd[waist_roll_idx].kp = self.config.locomotion_arm_waist_kps[1] # index 1 is waist roll + self.msg.motor_cmd[waist_roll_idx].kp = self.config.locomotion_arm_waist_kps[ + 1 + ] # index 1 is waist roll self.msg.motor_cmd[waist_roll_idx].kd = self.config.locomotion_arm_waist_kds[1] self.msg.motor_cmd[waist_roll_idx].tau = 0 - + # Zero out the problematic joints (waist yaw, waist pitch, wrist pitch/yaw) problematic_joints = [12, 14, 20, 21, 27, 28] for joint_idx in problematic_joints: @@ -1001,11 +1074,10 @@ class UnitreeG1(Robot): self.msg.motor_cmd[joint_idx].kd = self.kd_wrist self.msg.motor_cmd[joint_idx].tau = 0 - # Send command self.msg.crc = self.crc.Crc(self.msg) self.lowcmd_publisher.Write(self.msg) - + def _locomotion_thread_loop(self): """Background thread that runs the locomotion policy at specified rate.""" logger.info("Locomotion thread started") @@ -1013,163 +1085,164 @@ class UnitreeG1(Robot): start_time = time.time() try: # Use different run function based on policy type - if self.policy_type == 'onnx': + if self.policy_type == "onnx": self.groot_locomotion_run() else: self.locomotion_run() except Exception as e: logger.error(f"Error in locomotion loop: {e}") - + # Sleep to maintain control rate elapsed = time.time() - start_time sleep_time = max(0, self.config.locomotion_control_dt - elapsed) time.sleep(sleep_time) logger.info("Locomotion thread stopped") - + def start_locomotion_thread(self): """Start the background locomotion control thread.""" if not self.config.locomotion_control: logger.warning("locomotion_control is False, cannot start thread") return - + 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): """Stop the background locomotion control thread.""" 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") - + # Also stop keyboard thread if running if self.keyboard_running: self.stop_keyboard_controls() - + def _keyboard_listener_thread(self): """Background thread that listens for keyboard input (sim mode only).""" - print("\n" + "="*60) + print("\n" + "=" * 60) print("KEYBOARD CONTROLS ACTIVE!") print(" W/S: Forward/Backward") print(" A/D: Left/Right") print(" Q/E: Rotate Left/Right") print(" R/F: Raise/Lower Height (±5cm)") print(" Z: Stop (zero velocity commands)") - print("="*60 + "\n") - + print("=" * 60 + "\n") + # Save terminal settings old_settings = None try: old_settings = termios.tcgetattr(sys.stdin) tty.setcbreak(sys.stdin.fileno()) - + while self.keyboard_running: if select.select([sys.stdin], [], [], 0.1)[0]: key = sys.stdin.read(1).lower() - + # Velocity commands - if key == 'w': + if key == "w": self.locomotion_cmd[0] += 0.4 # Forward - elif key == 's': + elif key == "s": self.locomotion_cmd[0] -= 0.4 # Backward - elif key == 'a': + elif key == "a": self.locomotion_cmd[1] += 0.25 # Left - elif key == 'd': + elif key == "d": self.locomotion_cmd[1] -= 0.25 # Right - elif key == 'q': + elif key == "q": self.locomotion_cmd[2] += 0.5 # Rotate left - elif key == 'e': + elif key == "e": self.locomotion_cmd[2] -= 0.5 # Rotate right - elif key == 'z': + elif key == "z": self.locomotion_cmd[:] = 0.0 # Stop - + # Height commands (only for GR00T ONNX policies) - elif key == 'r': + elif key == "r": self.groot_height_cmd += 0.05 # Raise 5cm - elif key == 'f': + elif key == "f": self.groot_height_cmd -= 0.05 # Lower 5cm - + # Clamp commands to reasonable limits self.locomotion_cmd[0] = np.clip(self.locomotion_cmd[0], -0.8, 0.8) # vx self.locomotion_cmd[1] = np.clip(self.locomotion_cmd[1], -0.5, 0.5) # vy self.locomotion_cmd[2] = np.clip(self.locomotion_cmd[2], -1.0, 1.0) # yaw_rate - + # Clamp height (reasonable range: 0.5m to 1.0m) - if hasattr(self, 'groot_height_cmd'): + if hasattr(self, "groot_height_cmd"): self.groot_height_cmd = np.clip(self.groot_height_cmd, 0.50, 1.00) - + # Print current commands - print(f"[VEL CMD] vx={self.locomotion_cmd[0]:.2f}, vy={self.locomotion_cmd[1]:.2f}, yaw={self.locomotion_cmd[2]:.2f}", end="") - if hasattr(self, 'groot_height_cmd'): + print( + f"[VEL CMD] vx={self.locomotion_cmd[0]:.2f}, vy={self.locomotion_cmd[1]:.2f}, yaw={self.locomotion_cmd[2]:.2f}", + end="", + ) + if hasattr(self, "groot_height_cmd"): print(f" | [HEIGHT] {self.groot_height_cmd:.3f}m", end="") print() # Newline - + finally: # Restore terminal settings if old_settings is not None: termios.tcsetattr(sys.stdin, termios.TCSADRAIN, old_settings) print("\nKeyboard controls stopped") - + def start_keyboard_controls(self): """Start the keyboard control thread (sim mode only).""" if not self.simulation_mode: logger.warning("Keyboard controls only available in simulation mode") return - + if self.keyboard_running: logger.warning("Keyboard controls already running") return - + self.keyboard_running = True self.keyboard_thread = threading.Thread(target=self._keyboard_listener_thread, daemon=True) self.keyboard_thread.start() logger.info("Keyboard controls started!") - + def stop_keyboard_controls(self): """Stop the keyboard control thread.""" if not self.keyboard_running: return - + logger.info("Stopping keyboard controls...") self.keyboard_running = False if self.keyboard_thread: self.keyboard_thread.join(timeout=2.0) logger.info("Keyboard controls stopped") - def init_locomotion(self): """Test locomotion control sequence: home arms -> move legs to default -> start policy thread.""" if not self.config.locomotion_control: logger.warning("locomotion_control is False, cannot run test sequence") return - - logger.info("Starting locomotion test sequence...") + logger.info("Starting locomotion test sequence...") # 2. Move legs to default position self.locomotion_move_to_default_pos() - + # 3. Wait 3 seconds time.sleep(3.0) - + # 4. Hold default leg position for 2 seconds self.locomotion_default_pos_state() - + # 5. Start locomotion policy thread (runs in background) logger.info("Starting locomotion policy control...") self.start_locomotion_thread() - + logger.info("Locomotion test sequence complete! Policy is now running in background.") logger.info("Use robot.stop_locomotion_thread() to stop the policy.") @@ -1178,88 +1251,21 @@ class UnitreeG1(Robot): if not self.config.locomotion_control: logger.warning("locomotion_control is False, cannot run GR00T init") return - + logger.info("Starting GR00T locomotion initialization...") - + # Move legs to default position (same as regular locomotion) self.locomotion_move_to_default_pos() - + # Wait 3 seconds time.sleep(3.0) - + # Hold default leg position for 2 seconds self.locomotion_default_pos_state() - + # Start locomotion policy thread (will use groot_locomotion_run) logger.info("Starting GR00T locomotion policy control...") self.start_locomotion_thread() - + logger.info("GR00T locomotion initialization complete! Policy is now running.") logger.info("516D observations (86D × 6 frames), 15D actions (legs + waist)") - - -class G1_29_JointArmIndex(IntEnum): - # Left arm - kLeftShoulderPitch = 15 - kLeftShoulderRoll = 16 - kLeftShoulderYaw = 17 - kLeftElbow = 18 - kLeftWristRoll = 19 - kLeftWristPitch = 20 - kLeftWristyaw = 21 - - # Right arm - kRightShoulderPitch = 22 - kRightShoulderRoll = 23 - kRightShoulderYaw = 24 - kRightElbow = 25 - kRightWristRoll = 26 - kRightWristPitch = 27 - kRightWristYaw = 28 - -class G1_29_JointIndex(IntEnum): - # Left leg - kLeftHipPitch = 0 - kLeftHipRoll = 1 - kLeftHipYaw = 2 - kLeftKnee = 3 - kLeftAnklePitch = 4 - kLeftAnkleRoll = 5 - - # Right leg - kRightHipPitch = 6 - kRightHipRoll = 7 - kRightHipYaw = 8 - kRightKnee = 9 - kRightAnklePitch = 10 - kRightAnkleRoll = 11 - - kWaistYaw = 12 #we're c - kWaistRoll = 13 - kWaistPitch = 14 - - # Left arm - kLeftShoulderPitch = 15 - kLeftShoulderRoll = 16 - kLeftShoulderYaw = 17 - kLeftElbow = 18 - kLeftWristRoll = 19 - kLeftWristPitch = 20 - kLeftWristyaw = 21 - - # Right arm - kRightShoulderPitch = 22 - kRightShoulderRoll = 23 - kRightShoulderYaw = 24 - kRightElbow = 25 - kRightWristRoll = 26 - kRightWristPitch = 27 - kRightWristYaw = 28 - - # not used - kNotUsedJoint0 = 29 - kNotUsedJoint1 = 30 - kNotUsedJoint2 = 31 - kNotUsedJoint3 = 32 - kNotUsedJoint4 = 33 - kNotUsedJoint5 = 34 \ No newline at end of file diff --git a/src/lerobot/robots/unitree_g1/unitree_sdk2_socket.py b/src/lerobot/robots/unitree_g1/unitree_sdk2_socket.py index bc6b23940..d7741cd2a 100644 --- a/src/lerobot/robots/unitree_g1/unitree_sdk2_socket.py +++ b/src/lerobot/robots/unitree_g1/unitree_sdk2_socket.py @@ -1,17 +1,18 @@ # unitree_sdk2_socket.py -import zmq import pickle -import time + +import zmq # you can tune these or read from env ROBOT_IP = "172.18.129.215" -LOWCMD_PORT = 6000 # laptop -> robot -LOWSTATE_PORT = 6001 # robot -> laptop +LOWCMD_PORT = 6000 # laptop -> robot +LOWSTATE_PORT = 6001 # robot -> laptop _ctx = None _lowcmd_sock = None _lowstate_sock = None + def ChannelFactoryInitialize(*args, **kwargs): global _ctx, _lowcmd_sock, _lowstate_sock if _ctx is not None: @@ -20,12 +21,12 @@ def ChannelFactoryInitialize(*args, **kwargs): # lowcmd: PUSH from laptop to robot _lowcmd_sock = _ctx.socket(zmq.PUSH) - _lowcmd_sock.setsockopt(zmq.CONFLATE, 1) + _lowcmd_sock.setsockopt(zmq.CONFLATE, 1) _lowcmd_sock.connect(f"tcp://{ROBOT_IP}:{LOWCMD_PORT}") # lowstate: SUB from robot - _lowstate_sock = _ctx.socket(zmq.SUB) # no topic filtering - _lowstate_sock.setsockopt(zmq.CONFLATE, 1) # keep only last message + _lowstate_sock = _ctx.socket(zmq.SUB) # no topic filtering + _lowstate_sock.setsockopt(zmq.CONFLATE, 1) # keep only last message _lowstate_sock.connect(f"tcp://{ROBOT_IP}:{LOWSTATE_PORT}") _lowstate_sock.setsockopt_string(zmq.SUBSCRIBE, "") # subscribe to all