From 87b0dc470c6795d015bab7c1caf440f557c6238f Mon Sep 17 00:00:00 2001 From: Martino Russi Date: Fri, 31 Jul 2026 11:25:08 +0200 Subject: [PATCH] dedup reset --- .../controllers/gr00t_locomotion.py | 1 + .../controllers/holosoma_locomotion.py | 1 + .../controllers/sonic_whole_body.py | 31 ++----------------- src/lerobot/robots/unitree_g1/unitree_g1.py | 24 +++++++++----- 4 files changed, 20 insertions(+), 37 deletions(-) diff --git a/src/lerobot/robots/unitree_g1/controllers/gr00t_locomotion.py b/src/lerobot/robots/unitree_g1/controllers/gr00t_locomotion.py index c9d08c870..e6f07d9c2 100644 --- a/src/lerobot/robots/unitree_g1/controllers/gr00t_locomotion.py +++ b/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 diff --git a/src/lerobot/robots/unitree_g1/controllers/holosoma_locomotion.py b/src/lerobot/robots/unitree_g1/controllers/holosoma_locomotion.py index 0e2d1c92c..20853e188 100644 --- a/src/lerobot/robots/unitree_g1/controllers/holosoma_locomotion.py +++ b/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 diff --git a/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py b/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py index 1fcee81c6..8d6bed95b 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py +++ b/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 self._startup_blend(lowstate, out) + return {f"{m.name}.q": float(target[m.value]) for m in G1_29_JointIndex} diff --git a/src/lerobot/robots/unitree_g1/unitree_g1.py b/src/lerobot/robots/unitree_g1/unitree_g1.py index f5b948c29..6744b3b54 100644 --- a/src/lerobot/robots/unitree_g1/unitree_g1.py +++ b/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 - # and publish arm targets from the teleoperator. + # The controller thread owns legs/waist (and for full-body controllers like SONIC, + # the arms too). Here we only publish arm targets from the teleoperator; full-body + # controllers carry no .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