dedup reset

src/lerobot/robots/unitree_g1/controllers/gr00t_locomotion.py

@@ -86,6 +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.cmd = np.array([0.0, 0.0, 0.0], dtype=np.float32)  # vx, vy, theta_dot
 
         # Robot state

src/lerobot/robots/unitree_g1/controllers/holosoma_locomotion.py

@@ -104,6 +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.cmd = np.zeros(3, dtype=np.float32)
 
         # Robot state

src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py

@@ -54,9 +54,6 @@ TOKEN_DIM = 64  # decoder latent size
 TOKEN_ACTION_PREFIX = "motion_token"  # nosec B105 - feature-key prefix, not a secret
 TOKEN_STATE_PREFIX = "motion_token_state"  # nosec B105 - feature-key prefix, not a secret
 
-# Startup blend duration (s): ease from the initial pose into the policy target on start.
-INIT_RAMP_S = 3.0
-
 # SONIC decoder checkpoint. Deploy constants (kp/kd, default_angles, action_scale,
 # neutral_token) are baked into the ONNX metadata; see upload_sonic_decoder.py.
 DEFAULT_SONIC_REPO_ID = "lerobot/sonic_decoder"
@@ -145,7 +142,6 @@ class SonicWholeBodyController:
     """
 
     control_dt = CONTROL_DT
-    full_body = True
 
     def __init__(self, policy_type: str = "default"):
         self.decoder, self.kp, self.kd, self.default_angles, self.action_scale, self.neutral_token = (
@@ -153,12 +149,11 @@ class SonicWholeBodyController:
         )
         self.decoder_input = self.decoder.get_inputs()[0].name
         self.default_angles_mj = self.default_angles[MUJOCO_TO_ISAACLAB]
-        self._init_ramp_steps = max(1, round(INIT_RAMP_S / CONTROL_DT))
         self.reset()
         logger.info("SonicWholeBodyController initialized")
 
     def reset(self) -> None:
-        """Reset internal state for a new episode: held token, history buffers, startup blend."""
+        """Reset internal state for a new episode: held token and 10-frame history buffers."""
         self.last_action_mj = np.zeros(29, np.float32)
         self.h_q_mj = [np.zeros(29, np.float32)] * 10
         self.h_dq_mj = [np.zeros(29, np.float32)] * 10
@@ -166,8 +161,6 @@ class SonicWholeBodyController:
         self.h_act_mj = [np.zeros(29, np.float32)] * 10
         self.h_quat = [np.array([1, 0, 0, 0], np.float32)] * 10
         self._last_token = None  # neutral token is re-seeded on the first tick
-        self._init_step = 0  # re-run the startup blend
-        self._start_pose: dict[str, float] = {}
 
     @property
     def action_features(self) -> dict[str, type]:
@@ -187,25 +180,6 @@ class SonicWholeBodyController:
         token = self._last_token if self._last_token is not None else np.zeros(TOKEN_DIM, dtype=np.float32)
         return {token_state_key(i): float(v) for i, v in enumerate(token)}
 
-    def _startup_blend(self, lowstate, out: dict) -> dict:
-        """Ease into policy control: over the first ``INIT_RAMP_S`` seconds, interpolate from
-        the pose captured on the first tick to the live policy target so the handoff has no snap."""
-        if self._init_step >= self._init_ramp_steps or not out:
-            return out
-        if self._init_step == 0:
-            self._start_pose = {
-                f"{m.name}.q": float(lowstate.motor_state[m.value].q) for m in G1_29_JointIndex
-            }
-        self._init_step += 1
-        ratio = min(1.0, self._init_step / self._init_ramp_steps)
-        blended = {
-            k: self._start_pose.get(k, float(tgt)) * (1.0 - ratio) + float(tgt) * ratio
-            for k, tgt in out.items()
-        }
-        if self._init_step >= self._init_ramp_steps:
-            logger.info("SONIC startup blend complete -> full policy control")
-        return blended
-
     def run_step(self, action: dict, lowstate) -> dict:
         if lowstate is None:
             return {}
@@ -251,5 +225,4 @@ class SonicWholeBodyController:
         )
         self.last_action_mj = action_mj.copy()
         target = self.default_angles + action_mj[ISAACLAB_TO_MUJOCO] * self.action_scale
-        out = {f"{m.name}.q": float(target[m.value]) for m in G1_29_JointIndex}
+        return {f"{m.name}.q": float(target[m.value]) for m in G1_29_JointIndex}
-        return self._startup_blend(lowstate, out)

src/lerobot/robots/unitree_g1/unitree_g1.py

@@ -359,6 +359,11 @@ class UnitreeG1(Robot):
             self.msg.motor_cmd[joint].kd = self.kd[joint.value]
             self.msg.motor_cmd[joint].q = lowstate.motor_state[joint.value].q
 
+        # Ease into the controller's home pose before it takes over, so the first commands
+        # don't snap from the connect-time pose.
+        if self.controller is not None and hasattr(self.controller, "default_angles"):
+            self.reset(default_positions=self.controller.default_angles)
+
         # Start controller thread if enabled
         if self.controller is not None:
             self._controller_thread = threading.Thread(target=self._controller_loop, daemon=True)
@@ -487,13 +492,10 @@ class UnitreeG1(Robot):
     def send_action(self, action: RobotAction) -> RobotAction:
         action_to_publish = action
         if self.controller is not None:
-            # Controller thread owns legs/waist. Here we only update joystick inputs
+            # The controller thread owns legs/waist (and for full-body controllers like SONIC,
-            # and publish arm targets from the teleoperator.
+            # the arms too). Here we only publish arm targets from the teleoperator; full-body
+            # controllers carry no <joint>.q in their action, so this is empty for them.
             self._update_controller_action(action)
-            # Full-body controllers (SONIC) own the whole 29-DoF command; the controller
-            # thread is the sole publisher, so there is nothing to publish here.
-            if getattr(self.controller, "full_body", False):
-                return action
             arm_prefixes = tuple(j.name for j in G1_29_JointArmIndex)
             action_to_publish = {
                 key: value
@@ -517,7 +519,8 @@ class UnitreeG1(Robot):
                 local_idx = joint.value - arm_start_idx
                 tau[joint.value] = arm_tau[local_idx]
 
-        self.publish_lowcmd(action_to_publish, tau=tau)
+        if action_to_publish:
+            self.publish_lowcmd(action_to_publish, tau=tau)
         return action
 
     def _update_controller_action(self, action: RobotAction) -> None:
@@ -573,6 +576,11 @@ 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
             for step in range(num_steps):
                 start_time = time.time()
@@ -584,7 +592,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.send_action(action_dict)
+                self.publish_lowcmd(action_dict, kp=ctrl_kp, kd=ctrl_kd)
 
                 # Maintain constant control rate
                 elapsed = time.time() - start_time