(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: