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,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/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_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])
+    leg_joint2motor_idx: list = field(default_factory=lambda: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11])
-    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_kps: list = field(
-    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])
+        default_factory=lambda: [150, 150, 150, 300, 40, 40, 150, 150, 150, 300, 40, 40]
-    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_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"
+    locomotion_imu_type: str = "pelvis"  # "torso" or "pelvis"

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_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:
@@ -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 = _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