precommit

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,26 +52,38 @@ 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/motion.pt"
-    policy_path: str = "src/lerobot/robots/unitree_g1/assets/g1/locomotion/GR00T-WholeBodyControl-Walk.onnx"
+    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_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])
+    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])
+    arm_waist_joint2motor_idx: list = field(
-    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])
+        default_factory=lambda: [12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28]
-    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])
+    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

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

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.comm.motion_switcher.motion_switcher_client import MotionSwitcherClient
-from unitree_sdk2py.core.channel import ChannelPublisher, ChannelSubscriber, ChannelFactoryInitialize
+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
+LOWCMD_PORT = 6000  # laptop -> robot
-LOWSTATE_PORT = 6001    # robot -> laptop
+LOWSTATE_PORT = 6001  # robot -> laptop
 
 
 def state_forward_loop(lowstate_sub, lowstate_sock, state_period: float):

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
+import logging
-from .config_unitree_g1 import UnitreeG1Config
+import select
-
+import struct
-import numpy as np
+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,8 +77,8 @@ class DataBuffer:
         with self.lock:
             self.data = data
 
-class UnitreeG1(Robot):
 
+class UnitreeG1(Robot):
     config_class = UnitreeG1Config
     name = "unitree_g1"
 
@@ -128,8 +116,11 @@ class UnitreeG1(Robot):
 
         self.calibrated = False
 
-
+        from lerobot.robots.unitree_g1.unitree_sdk2_socket import (
-        from lerobot.robots.unitree_g1.unitree_sdk2_socket import ChannelPublisher, ChannelSubscriber, ChannelFactoryInitialize  # dds
+            ChannelFactoryInitialize,
+            ChannelPublisher,
+            ChannelSubscriber,
+        )  # dds
 
         ChannelFactoryInitialize(0)
 
@@ -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,9 +172,9 @@ 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)
@@ -217,18 +208,18 @@ class UnitreeG1(Robot):
             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():
@@ -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(
-                    logger.info(f"ONNX output: {self.policy_balance.get_outputs()[0].name}, shape: {self.policy_balance.get_outputs()[0].shape}")
+                        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,12 +242,18 @@ 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(
-                    logger.info(f"ONNX output: {self.policy.get_outputs()[0].name}, shape: {self.policy.get_outputs()[0].shape}")
+                        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'
+                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()
@@ -264,7 +265,7 @@ class UnitreeG1(Robot):
             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)
@@ -284,7 +285,7 @@ class UnitreeG1(Robot):
                 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):
@@ -323,9 +323,30 @@ class UnitreeG1(Robot):
 
             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()
@@ -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:
@@ -448,7 +470,7 @@ class UnitreeG1(Robot):
             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()
@@ -470,33 +492,40 @@ class UnitreeG1(Robot):
 
         # Extract IMU data
         imu_data = {
-            'quaternion': lowstate.imu_state.quaternion,  # [w, x, y, z]
+            "quaternion": lowstate.imu_state.quaternion,  # [w, x, y, z]
-            'gyroscope': lowstate.imu_state.gyroscope,  # [x, y, z] rad/s
+            "gyroscope": lowstate.imu_state.gyroscope,  # [x, y, z] rad/s
-            'accelerometer': lowstate.imu_state.accelerometer,  # [x, y, z] m/s²
+            "accelerometer": lowstate.imu_state.accelerometer,  # [x, y, z] m/s²
-            'rpy': lowstate.imu_state.rpy,  # [roll, pitch, yaw] rad
+            "rpy": lowstate.imu_state.rpy,  # [roll, pitch, yaw] rad
-            'temperature': lowstate.imu_state.temperature,  # °C
+            "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({
+            motors_data.append(
-                'id': i,
+                {
-                'q': motor.q,  # position (rad)
+                    "id": i,
-                'dq': motor.dq,  # velocity (rad/s)
+                    "q": motor.q,  # position (rad)
-                'tau_est': motor.tau_est,  # estimated torque (Nm)
+                    "dq": motor.dq,  # velocity (rad/s)
-                'temperature': motor.temperature[0] if isinstance(motor.temperature, (list, tuple)) else motor.temperature,  # °C
+                    "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,
+            "imu": imu_data,
-            'motors': motors_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():
@@ -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
+        # 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 
+        # 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.
+        # 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 teleop them wiuth the glove hehe
-        #then we get ALL THE DATA
+        # 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,10 +587,10 @@ 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 is just zeros
-        #action_np = np.zeros(14)
+        # action_np = np.zeros(14)
-        #print(action_np)
+        # print(action_np)
-        #exit()
+        # exit()
 
         tau = np.zeros(14)
 
@@ -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 = {}
@@ -697,7 +727,6 @@ class UnitreeG1(Robot):
             time.sleep(self.config.locomotion_control_dt)
         logger.info("Finished holding default leg position")
 
-
     class RemoteController:
         def __init__(self):
             self.lx = 0
@@ -796,13 +825,17 @@ class UnitreeG1(Robot):
 
             # 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
@@ -812,10 +845,10 @@ class UnitreeG1(Robot):
         # 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)
@@ -824,7 +857,10 @@ class UnitreeG1(Robot):
             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)):
@@ -888,7 +924,9 @@ class UnitreeG1(Robot):
         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
@@ -903,11 +941,40 @@ class UnitreeG1(Robot):
 
         # 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
+        groot_default_angles = np.array(
-                                         -0.1, 0.0, 0.0, 0.3, -0.2, 0.0,  # right leg
+            [
-                                         0.0, 0.0, 0.0,                     # waist
+                -0.1,
-                                         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, 0.0], dtype=np.float32)  # right arm (zeroed)
+                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
@@ -920,7 +987,9 @@ class UnitreeG1(Robot):
             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
@@ -966,7 +1035,9 @@ class UnitreeG1(Robot):
         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):
@@ -983,7 +1054,9 @@ class UnitreeG1(Robot):
         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
 
@@ -1001,7 +1074,6 @@ 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)
@@ -1013,7 +1085,7 @@ 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()
@@ -1059,14 +1131,14 @@ class UnitreeG1(Robot):
 
     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
@@ -1079,25 +1151,25 @@ class UnitreeG1(Robot):
                     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
@@ -1106,12 +1178,15 @@ class UnitreeG1(Robot):
                     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="")
+                    print(
-                    if hasattr(self, 'groot_height_cmd'):
+                        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
 
@@ -1147,7 +1222,6 @@ class UnitreeG1(Robot):
             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:
@@ -1156,7 +1230,6 @@ class UnitreeG1(Robot):
 
         logger.info("Starting locomotion test sequence...")
 
-
         # 2. Move legs to default position
         self.locomotion_move_to_default_pos()
 
@@ -1196,70 +1269,3 @@ class UnitreeG1(Robot):
 
         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

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
+LOWCMD_PORT = 6000  # laptop -> robot
-LOWSTATE_PORT = 6001    # robot -> laptop
+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:
@@ -24,8 +25,8 @@ def ChannelFactoryInitialize(*args, **kwargs):
     _lowcmd_sock.connect(f"tcp://{ROBOT_IP}:{LOWCMD_PORT}")
 
     # lowstate: SUB from robot
-    _lowstate_sock = _ctx.socket(zmq.SUB)      # no topic filtering
+    _lowstate_sock = _ctx.socket(zmq.SUB)  # no topic filtering
-    _lowstate_sock.setsockopt(zmq.CONFLATE, 1)          # keep only last message
+    _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