From d1eed119a6a1b06d4b1c80b08982744132d22c9c Mon Sep 17 00:00:00 2001 From: Martino Russi Date: Fri, 31 Jul 2026 11:43:17 +0200 Subject: [PATCH] refactor --- .../controllers/gr00t_locomotion.py | 2 +- .../controllers/holosoma_locomotion.py | 2 +- src/lerobot/robots/unitree_g1/unitree_g1.py | 20 ++++++++++--------- 3 files changed, 13 insertions(+), 11 deletions(-) diff --git a/src/lerobot/robots/unitree_g1/controllers/gr00t_locomotion.py b/src/lerobot/robots/unitree_g1/controllers/gr00t_locomotion.py index e6f07d9c2..f293a0b1c 100644 --- a/src/lerobot/robots/unitree_g1/controllers/gr00t_locomotion.py +++ b/src/lerobot/robots/unitree_g1/controllers/gr00t_locomotion.py @@ -86,7 +86,7 @@ class GrootLocomotionController: # Load policies self.policy_balance, self.policy_walk = load_groot_policies() - self.default_angles = GROOT_DEFAULT_ANGLES # home pose (29,), for reset ease-in + self.default_angles = GROOT_DEFAULT_ANGLES self.cmd = np.array([0.0, 0.0, 0.0], dtype=np.float32) # vx, vy, theta_dot # Robot state diff --git a/src/lerobot/robots/unitree_g1/controllers/holosoma_locomotion.py b/src/lerobot/robots/unitree_g1/controllers/holosoma_locomotion.py index 20853e188..da9ef2ddb 100644 --- a/src/lerobot/robots/unitree_g1/controllers/holosoma_locomotion.py +++ b/src/lerobot/robots/unitree_g1/controllers/holosoma_locomotion.py @@ -104,7 +104,7 @@ class HolosomaLocomotionController: # Load policy and gains self.policy, self.kp, self.kd = load_policy() - self.default_angles = DEFAULT_ANGLES # home pose (29,), for reset ease-in + self.default_angles = DEFAULT_ANGLES self.cmd = np.zeros(3, dtype=np.float32) # Robot state diff --git a/src/lerobot/robots/unitree_g1/unitree_g1.py b/src/lerobot/robots/unitree_g1/unitree_g1.py index 6744b3b54..bfddded01 100644 --- a/src/lerobot/robots/unitree_g1/unitree_g1.py +++ b/src/lerobot/robots/unitree_g1/unitree_g1.py @@ -350,8 +350,14 @@ class UnitreeG1(Robot): logger.info("[UnitreeG1] Connected to robot.") self.msg.mode_machine = lowstate.mode_machine - self.kp = np.array(self.config.kp, dtype=np.float32) - self.kd = np.array(self.config.kd, dtype=np.float32) + # Prefer the active controller's gains (e.g. SONIC loads kp/kd from its ONNX); + # otherwise fall back to the config defaults. + if self.controller is not None and hasattr(self.controller, "kp"): + self.kp = np.array(self.controller.kp, dtype=np.float32) + self.kd = np.array(self.controller.kd, dtype=np.float32) + else: + self.kp = np.array(self.config.kp, dtype=np.float32) + self.kd = np.array(self.config.kd, dtype=np.float32) for joint in G1_29_JointIndex: self.msg.motor_cmd[joint].mode = 1 @@ -576,12 +582,8 @@ class UnitreeG1(Robot): for motor in G1_29_JointIndex: init_dof_pos[motor.value] = obs[f"{motor.name}.q"] - # Publish the whole-body pose directly (bypass send_action, which only forwards - # arm targets when a controller is active) with the controller's gains if any. - ctrl_kp = getattr(self.controller, "kp", None) - ctrl_kd = getattr(self.controller, "kd", None) - - # Interpolate to default position + # Interpolate to default position. Publish directly so the whole body moves: + # send_action() would arm-filter this to the arms while a controller is active. for step in range(num_steps): start_time = time.time() @@ -592,7 +594,7 @@ class UnitreeG1(Robot): interp_pos = init_dof_pos[motor.value] * (1 - alpha) + target_pos * alpha action_dict[f"{motor.name}.q"] = float(interp_pos) - self.publish_lowcmd(action_dict, kp=ctrl_kp, kd=ctrl_kd) + self.publish_lowcmd(action_dict) # Maintain constant control rate elapsed = time.time() - start_time