(refactor) token mode is default

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

@@ -44,7 +44,6 @@ from ..g1_utils import (
     G1_29_JointIndex,
     get_gravity_orientation,
 )
-from ..unitree_g1 import lowstate_to_obs
 
 logger = logging.getLogger(__name__)
 
@@ -91,17 +90,6 @@ def load_sonic_decoder(repo_id: str = DEFAULT_SONIC_REPO_ID):
     return session, arr["kp"], arr["kd"], arr["default_angles"], arr["action_scale"], arr["neutral_token"]
 
 
-def _to_mujoco(a):
-    """Apply the ``MUJOCO_TO_ISAACLAB`` gather to a 29-vector (deploy-order reorder).
-
-    NOTE: this returns ``a[MUJOCO_TO_ISAACLAB]``. The ``_mj`` suffixes and the exact
-    permutation direction are a fixed convention validated against the deployed SONIC ONNX
-    policy (the decoder consumes vectors in this order). Do not "correct" the table or
-    rename toward the opposite direction without re-validating on hardware.
-    """
-    return a[MUJOCO_TO_ISAACLAB]
-
-
 # Action-feature prefix for the latent-token interface (see _extract_token_from_action).
 TOKEN_ACTION_PREFIX = "motion_token"  # nosec B105 - feature-key prefix, not a secret
 # Proprio-state prefix for the token interface: the robot echoes the last commanded token
@@ -158,7 +146,7 @@ class SonicDecoder:
         self.decoder_input = decoder.get_inputs()[0].name
         self.default_angles = np.asarray(default_angles, np.float32)
         self.action_scale = np.asarray(action_scale, np.float32)
-        self.default_angles_mj = _to_mujoco(self.default_angles)
+        self.default_angles_mj = self.default_angles[MUJOCO_TO_ISAACLAB]
         self.token = np.zeros(TOKEN_DIM, np.float32)
         self.last_action_mj = np.zeros(29, np.float32)
         self.h_q_mj = [np.zeros(29, np.float32)] * 10
@@ -184,8 +172,10 @@ class SonicDecoder:
     def update_history(self, q, dq, ang, quat):
         """Push the latest proprioception (pos/vel/gyro/orientation) into the 10-frame buffers."""
         quat = quat / (np.linalg.norm(quat) + 1e-8)
-        q_mj = _to_mujoco(q)
-        dq_mj = _to_mujoco(dq)
+        # Reorder IsaacLab-order state into the MuJoCo order the decoder consumes. This
+        # permutation direction is validated against the deployed SONIC ONNX; don't flip it.
+        q_mj = q[MUJOCO_TO_ISAACLAB]
+        dq_mj = dq[MUJOCO_TO_ISAACLAB]
         self.h_q_mj = [q_mj - self.default_angles_mj] + self.h_q_mj[:-1]
         self.h_dq_mj = [dq_mj] + self.h_dq_mj[:-1]
         self.h_ang = [ang.copy()] + self.h_ang[:-1]
@@ -309,13 +299,9 @@ class SonicWholeBodyController:
         self._init_step = 0
         self._start_pose: dict[str, float] = {}
 
-        # Token-interface state. ``token_mode`` is set True by the robot whenever a SONIC
-        # whole-body controller is selected (token-driven deploy): the controller then holds a
-        # stable *neutral* token until the first real token arrives, and afterwards holds the
-        # *last* token received between ticks (the async controller runs ~50 Hz while a token
-        # VLA streams ~30 Hz). This lives here (not in the entry-point script) so it applies
-        # uniformly to run_g1_server, lerobot-rollout and the sim replays.
-        self.token_mode = False
+        # Token-interface state. The controller holds a stable *neutral* token until the first
+        # real token arrives, and afterwards holds the *last* token received between ticks (the
+        # async controller runs ~50 Hz while a token VLA streams ~30 Hz).
         self._last_token: np.ndarray | None = None
 
         logger.info("SONIC ready (decoder, 64-D token command path)")
@@ -346,23 +332,19 @@ class SonicWholeBodyController:
             logger.info("SONIC startup blend complete -> full policy control")
         return blended
 
-    def run_step(self, action: dict, lowstate) -> dict:
-        if lowstate is None:
+    def run_step(self, action: dict, obs: dict) -> dict:
+        if not obs:
             return {}
-        obs = lowstate_to_obs(lowstate)
 
         # Token-only interface (token-output VLA): a dense 64-D ``motion_token.{i}`` command
         # is decoded directly, encoder bypassed.
         token = _extract_token_from_action(action)
         if token is not None:
             self._last_token = token
-        elif self._last_token is None and self.token_mode:
-            # Token-driven deploy, but no token has arrived yet: hold the checkpoint's neutral
-            # token, which the decoder maps to a stable, natural standing pose.
+        elif self._last_token is None:
+            # No token has arrived yet: hold the checkpoint's neutral token, which the decoder
+            # maps to a stable, natural standing pose.
             self._last_token = self._neutral_token.copy()
-        if self._last_token is None:
-            # No token yet and not in token_mode: hold (keep last target).
-            return {}
         # Either a fresh token this tick or the last one received (held between the ~30 Hz
         # token stream and the ~50 Hz control loop).
         return self._startup_blend(obs, self.controller.step(obs, self._last_token))
@@ -371,7 +353,7 @@ class SonicWholeBodyController:
         self._runtime.reset()
         self._init_step = 0  # re-run the startup blend after a reset
         self._start_pose = {}
-        # Drop the held token so token_mode re-seeds the neutral token after a reset.
+        # Drop the held token so the neutral token is re-seeded after a reset.
         self._last_token = None
 
     def shutdown(self):

src/lerobot/robots/unitree_g1/unitree_g1.py

@@ -105,47 +105,6 @@ class G1_29_LowState:  # noqa: N801
     mode_machine: int = 0  # Robot mode
 
 
-def lowstate_to_obs(lowstate) -> dict:
-    """Build a robot observation dict from a Unitree lowstate.
-
-    Shared by ``UnitreeG1.get_observation`` and the SONIC pipeline so the
-    lowstate -> obs mapping lives in exactly one place. Keys match the
-    ``<joint>.q``/``imu.*`` schema consumed across the controllers.
-    """
-    obs: dict = {}
-
-    for motor in G1_29_JointIndex:
-        idx = motor.value
-        obs[f"{motor.name}.q"] = lowstate.motor_state[idx].q
-        obs[f"{motor.name}.dq"] = lowstate.motor_state[idx].dq
-        obs[f"{motor.name}.tau"] = lowstate.motor_state[idx].tau_est
-
-    imu = lowstate.imu_state
-    if imu.gyroscope:
-        obs["imu.gyro.x"] = imu.gyroscope[0]
-        obs["imu.gyro.y"] = imu.gyroscope[1]
-        obs["imu.gyro.z"] = imu.gyroscope[2]
-    if imu.accelerometer:
-        obs["imu.accel.x"] = imu.accelerometer[0]
-        obs["imu.accel.y"] = imu.accelerometer[1]
-        obs["imu.accel.z"] = imu.accelerometer[2]
-    if imu.quaternion:
-        obs["imu.quat.w"] = imu.quaternion[0]
-        obs["imu.quat.x"] = imu.quaternion[1]
-        obs["imu.quat.y"] = imu.quaternion[2]
-        obs["imu.quat.z"] = imu.quaternion[3]
-    if imu.rpy:
-        obs["imu.rpy.roll"] = imu.rpy[0]
-        obs["imu.rpy.pitch"] = imu.rpy[1]
-        obs["imu.rpy.yaw"] = imu.rpy[2]
-
-    wr = getattr(lowstate, "wireless_remote", None)
-    if wr:
-        obs["wireless_remote"] = bytes(wr) if not isinstance(wr, (bytes, bytearray)) else wr
-
-    return obs
-
-
 class UnitreeG1(Robot):
     config_class = UnitreeG1Config
     name = "unitree_g1"
@@ -190,12 +149,6 @@ class UnitreeG1(Robot):
 
         # Lower-body / whole-body controller loaded dynamically
         self.controller: LocomotionController | None = make_locomotion_controller(config.controller)
-
-        # A SONIC whole-body controller always runs in token mode: it holds a neutral
-        # token until the first real one arrives, then holds the last token between ticks.
-        if self.controller is not None and hasattr(self.controller, "token_mode"):
-            self.controller.token_mode = True
-
         # Controller thread state
         self._controller_thread = None
         # When set, the controller loop stops publishing low commands so reset() can
@@ -377,8 +330,12 @@ class UnitreeG1(Robot):
                 with self._controller_action_lock:
                     controller_input = dict(self.controller_input)
 
-                # Run controller step
-                controller_action = self.controller.run_step(controller_input, lowstate)
+                # Full-body controllers (SONIC) consume the full observation dict; others
+                # take the raw lowstate. get_observation() is the single lowstate -> obs builder.
+                controller_state = (
+                    self.get_observation() if getattr(self.controller, "full_body", False) else lowstate
+                )
+                controller_action = self.controller.run_step(controller_input, controller_state)
 
                 # Write controller output snapshot
                 with self._controller_action_lock:

upload_sonic_decoder.py

@@ -1,230 +0,0 @@
-#!/usr/bin/env python3
-"""Provision the SONIC decoder checkpoint at ``lerobot/sonic_decoder``.
-
-Takes NVIDIA's ``nvidia/GEAR-SONIC/model_decoder.onnx``, embeds the SONIC deploy constants
-(``kp``/``kd`` PD gains, ``default_angles`` standing pose, the residual ``action_scale``, and
-the ``neutral_token`` idle latent) into the ONNX ``metadata_props`` (the convention Holosoma
-uses for its gains), and pushes the result to ``lerobot/sonic_decoder``. After this runs, the
-runtime loads the decoder *and* every one of these constants straight from the checkpoint --
-no motor-physics math at deploy time, so ``sonic_whole_body.py`` carries none of the
-armature/bandwidth machinery nor any hardcoded deploy constants.
-
-The constants here are derived once from Unitree motor physics (armature + target bandwidth).
-That derivation is intentionally kept in this one-off provisioning script (not the runtime);
-the shared/harmonic helper is a separate PR.
-
-Build only (no network/auth needed if the source ONNX is already cached):
-    python upload_sonic_decoder.py --out ./sonic_decoder
-
-Build + upload:
-    huggingface-cli login          # or export HF_TOKEN=...
-    python upload_sonic_decoder.py --upload
-"""
-
-from __future__ import annotations
-
-import argparse
-import json
-import pathlib
-
-import numpy as np
-import onnx
-from huggingface_hub import hf_hub_download
-
-SRC_REPO_ID = "nvidia/GEAR-SONIC"
-SRC_FILENAME = "model_decoder.onnx"
-DST_REPO_ID = "lerobot/sonic_decoder"
-
-# ── SONIC deploy-constant derivation (provisioning-time only) ─────────────────
-# All constants are (29,) in IsaacLab joint order: legs, waist, arms.
-# kp = armature * w**2, kd = 4 * armature * w, with a x2 factor on the stiff joints
-# (ankles + waist). action_scale = 0.25 * effort / (armature * w**2) is the residual
-# scaling that maps decoder output to a joint-angle delta on top of default_angles.
-NATURAL_FREQ = 10.0 * 2.0 * np.pi
-MOTOR_ARMATURE = {"5020": 0.003609725, "7520_14": 0.010177520, "7520_22": 0.025101925, "4010": 0.00425}
-EFFORT = {"5020": 25.0, "7520_14": 88.0, "7520_22": 139.0, "4010": 5.0}
-MOTOR_MODELS = (
-    ["7520_22", "7520_22", "7520_14", "7520_22", "5020", "5020"] * 2
-    + ["7520_14", "5020", "5020"]
-    + ["5020", "5020", "5020", "5020", "5020", "4010", "4010"] * 2
-)
-DOUBLE_INDICES = {4, 5, 10, 11, 13, 14}  # ankles + waist
-
-# Nominal standing pose (rad), 29 joints in IsaacLab order. Decoder actions are residuals
-# added on top of this.
-DEFAULT_ANGLES = [
-    -0.312,
-    0.0,
-    0.0,
-    0.669,
-    -0.363,
-    0.0,  # left leg
-    -0.312,
-    0.0,
-    0.0,
-    0.669,
-    -0.363,
-    0.0,  # right leg
-    0.0,
-    0.0,
-    0.0,  # waist
-    0.2,
-    0.2,
-    0.0,
-    0.6,
-    0.0,
-    0.0,
-    0.0,  # left arm
-    0.2,
-    -0.2,
-    0.0,
-    0.6,
-    0.0,
-    0.0,
-    0.0,  # right arm
-]
-
-# Neutral idle token (64-D), held until the first real token arrives. Captured from the
-# encoder while the robot stood idle in sim: the encoder is an FSQ bottleneck (~5 bit/dim,
-# Div(16)), so tokens live on the 1/16 grid. We store the integer FSQ codes and rescale by
-# 1/16 -> an exact on-grid token that decodes to a stable, natural standing pose (unlike the
-# literal all-zero token, which is off-manifold and decodes to a slightly goofy stance).
-NEUTRAL_TOKEN_CODES = [
-    -1,
-    3,
-    1,
-    -1,
-    1,
-    -3,
-    6,
-    1,
-    1,
-    1,
-    -2,
-    -4,
-    -2,
-    0,
-    -3,
-    -1,
-    2,
-    -1,
-    -3,
-    -5,
-    3,
-    1,
-    1,
-    -4,
-    -1,
-    -1,
-    1,
-    -7,
-    0,
-    1,
-    2,
-    -2,
-    5,
-    -2,
-    -2,
-    -4,
-    0,
-    -1,
-    3,
-    -1,
-    0,
-    -5,
-    -1,
-    0,
-    -4,
-    0,
-    0,
-    -1,
-    -1,
-    2,
-    -2,
-    1,
-    3,
-    3,
-    1,
-    0,
-    0,
-    6,
-    0,
-    -7,
-    3,
-    0,
-    2,
-    -2,
-]
-
-
-def compute_kp_kd() -> tuple[list[float], list[float]]:
-    """Return (kp, kd) as plain float lists, (29,) in IsaacLab joint order."""
-
-    def stiffness(k):
-        return MOTOR_ARMATURE[k] * NATURAL_FREQ**2
-
-    def damping(k):
-        return 4.0 * MOTOR_ARMATURE[k] * NATURAL_FREQ
-
-    kp = [(2 if i in DOUBLE_INDICES else 1) * stiffness(k) for i, k in enumerate(MOTOR_MODELS)]
-    kd = [(2 if i in DOUBLE_INDICES else 1) * damping(k) for i, k in enumerate(MOTOR_MODELS)]
-    return kp, kd
-
-
-def compute_action_scale() -> list[float]:
-    """Return the per-joint residual action scale, (29,) in IsaacLab joint order."""
-    return [0.25 * EFFORT[k] / (MOTOR_ARMATURE[k] * NATURAL_FREQ**2) for k in MOTOR_MODELS]
-
-
-def build(out_dir: pathlib.Path) -> pathlib.Path:
-    """Download the source decoder, embed the deploy-constant metadata, save to ``out_dir``."""
-    src = hf_hub_download(repo_id=SRC_REPO_ID, filename=SRC_FILENAME)
-    model = onnx.load(src)
-
-    kp, kd = compute_kp_kd()
-    neutral_token = [c / 16.0 for c in NEUTRAL_TOKEN_CODES]  # FSQ Div(16): codes -> on-grid token
-    meta = {prop.key: prop.value for prop in model.metadata_props}
-    meta["kp"] = json.dumps(kp)
-    meta["kd"] = json.dumps(kd)
-    meta["action_scale"] = json.dumps(compute_action_scale())
-    meta["default_angles"] = json.dumps(DEFAULT_ANGLES)
-    meta["neutral_token"] = json.dumps(neutral_token)
-    # Rewrite metadata_props with the merged dict.
-    del model.metadata_props[:]
-    for key, value in meta.items():
-        model.metadata_props.add(key=key, value=value)
-
-    out_dir.mkdir(parents=True, exist_ok=True)
-    out_path = out_dir / SRC_FILENAME
-    onnx.save(model, out_path)
-    print(f"Wrote {out_path} with kp/kd/action_scale/default_angles/neutral_token metadata.")
-    return out_path
-
-
-def upload(out_path: pathlib.Path) -> None:
-    from huggingface_hub import HfApi
-
-    api = HfApi()
-    api.create_repo(repo_id=DST_REPO_ID, repo_type="model", exist_ok=True)
-    api.upload_file(
-        path_or_fileobj=str(out_path),
-        path_in_repo=SRC_FILENAME,
-        repo_id=DST_REPO_ID,
-        repo_type="model",
-    )
-    print(f"Uploaded {out_path.name} -> {DST_REPO_ID}")
-
-
-def main() -> None:
-    p = argparse.ArgumentParser()
-    p.add_argument("--out", type=pathlib.Path, default=pathlib.Path("./sonic_decoder"))
-    p.add_argument("--upload", action="store_true", help="Push the built ONNX to the hub")
-    args = p.parse_args()
-
-    out_path = build(args.out)
-    if args.upload:
-        upload(out_path)
-
-
-if __name__ == "__main__":
-    main()