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 8d6bed95b..bf96eaf51 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py +++ b/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py @@ -104,35 +104,6 @@ def load_policy( return decoder, arr["kp"], arr["kd"], arr["default_angles"], arr["action_scale"], arr["neutral_token"] -def token_action_key(i: int) -> str: - """Action-dict key for the i-th component of the 64-D SONIC latent token. - - The ``.pos`` suffix is required so the value flows through ``lerobot-rollout``, which - only routes ``.pos`` scalar features onto the policy action vector. - """ - return f"{TOKEN_ACTION_PREFIX}.{i}.pos" - - -def token_state_key(i: int) -> str: - """Observation key for the i-th component of the 64-D SONIC latent token state.""" - return f"{TOKEN_STATE_PREFIX}.{i}.pos" - - -def _extract_token_from_action(action: dict | None) -> np.ndarray | None: - """Reassemble a dense (64,) latent token from ``motion_token.{i}`` keys, or None. - - The token-only interface: the caller supplies the 64-D encoder latent directly (e.g. a - token-output VLA's action), which the decoder consumes with the encoder bypassed. - Requires the full dense token; a partial one is ignored (returns None). - """ - if not action: - return None - keys = [token_action_key(i) for i in range(TOKEN_DIM)] - if any(key not in action for key in keys): - return None - return np.fromiter((float(action[key]) for key in keys), dtype=np.float32, count=TOKEN_DIM) - - class SonicWholeBodyController: """Full-body SONIC decoder controller for UnitreeG1's background controller thread. @@ -149,6 +120,12 @@ class SonicWholeBodyController: ) self.decoder_input = self.decoder.get_inputs()[0].name self.default_angles_mj = self.default_angles[MUJOCO_TO_ISAACLAB] + + # 64-D latent-token action space; rollout maps the policy's 64-D output onto these keys. + self.action_ft = {f"{TOKEN_ACTION_PREFIX}.{i}.pos": float for i in range(TOKEN_DIM)} + # 64-D token proprio state, aggregated by rollout into observation.state (last token). + self.observation_ft = {f"{TOKEN_STATE_PREFIX}.{i}.pos": float for i in range(TOKEN_DIM)} + self.reset() logger.info("SonicWholeBodyController initialized") @@ -162,33 +139,23 @@ class SonicWholeBodyController: 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 - @property - def action_features(self) -> dict[str, type]: - """64-D latent-token action space (``motion_token.{i}.pos``); rollout maps the policy's - 64-D output straight onto these keys.""" - return {token_action_key(i): float for i in range(TOKEN_DIM)} - - @property - def observation_features(self) -> dict[str, type]: - """64-D latent-token proprio state (``motion_token_state.{i}.pos``), aggregated by the - rollout into a 64-D ``observation.state`` (the last token decoded).""" - return {token_state_key(i): float for i in range(TOKEN_DIM)} - def observation_state(self) -> dict[str, float]: """Echo the last decoded token as ``observation.state`` so a token-output VLA closes the loop on its own previous token.""" 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)} + return {f"{TOKEN_STATE_PREFIX}.{i}.pos": float(v) for i, v in enumerate(token)} def run_step(self, action: dict, lowstate) -> dict: if lowstate is None: return {} - # Token: fresh from the policy this tick, else hold the last one (neutral until the - # first real token arrives, which decodes to a stable standing pose). - token = _extract_token_from_action(action) - if token is not None: - self._last_token = token + # Token: reassemble the dense 64-D latent from motion_token.{i}.pos (all keys required); + # else hold the last one (neutral until the first real token, which decodes to a stand). + keys = [f"{TOKEN_ACTION_PREFIX}.{i}.pos" for i in range(TOKEN_DIM)] + if action and all(k in action for k in keys): + self._last_token = np.fromiter( + (float(action[k]) for k in keys), dtype=np.float32, count=TOKEN_DIM + ) elif self._last_token is None: self._last_token = self.neutral_token.copy() diff --git a/src/lerobot/robots/unitree_g1/unitree_g1.py b/src/lerobot/robots/unitree_g1/unitree_g1.py index bfddded01..58a9085ef 100644 --- a/src/lerobot/robots/unitree_g1/unitree_g1.py +++ b/src/lerobot/robots/unitree_g1/unitree_g1.py @@ -232,7 +232,7 @@ class UnitreeG1(Robot): @cached_property def observation_features(self) -> dict[str, type | tuple]: # Controllers may contribute their own proprio features (e.g. SONIC's token state). - controller_ft = getattr(self.controller, "observation_features", {}) + controller_ft = getattr(self.controller, "observation_ft", {}) return {**self._motors_ft, **controller_ft, **self._cameras_ft} @cached_property @@ -242,7 +242,7 @@ class UnitreeG1(Robot): return {f"{G1_29_JointIndex(motor).name}.q": float for motor in G1_29_JointIndex} # Whole-body controllers (SONIC): 64-D latent token. - controller_ft = getattr(self.controller, "action_features", None) + controller_ft = getattr(self.controller, "action_ft", None) if controller_ft is not None: return dict(controller_ft) @@ -498,9 +498,8 @@ class UnitreeG1(Robot): def send_action(self, action: RobotAction) -> RobotAction: action_to_publish = action if self.controller is not None: - # 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. + # Controller thread owns legs/waist. Here we only update joystick inputs + # and publish arm targets from the teleoperator. self._update_controller_action(action) arm_prefixes = tuple(j.name for j in G1_29_JointArmIndex) action_to_publish = { @@ -525,8 +524,7 @@ class UnitreeG1(Robot): local_idx = joint.value - arm_start_idx tau[joint.value] = arm_tau[local_idx] - if action_to_publish: - self.publish_lowcmd(action_to_publish, tau=tau) + self.publish_lowcmd(action_to_publish, tau=tau) return action def _update_controller_action(self, action: RobotAction) -> None: @@ -582,8 +580,7 @@ class UnitreeG1(Robot): for motor in G1_29_JointIndex: init_dof_pos[motor.value] = obs[f"{motor.name}.q"] - # 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. + # Interpolate to default position for step in range(num_steps): start_time = time.time()