merge sonic runtime into sonicwbc

Drop the full-body pause/direct-publish reset special-casing (and the
_controller_paused plumbing) so SONIC uses the same reset path as the other
controllers. The removed logic is preserved on the better-reset branch.

Co-authored-by: Cursor <cursoragent@cursor.com>

src/lerobot/robots/unitree_g1/config_unitree_g1.py

@@ -68,6 +68,10 @@ class UnitreeG1Config(RobotConfig):
     # Compensates for gravity on the unitree's arms using the arm ik solver
     gravity_compensation: bool = False
 
-    # Lower-body controller class name, e.g. "GrootLocomotionController" or
-    # "HolosomaLocomotionController". None disables it.
+    # Locomotion controller class name, e.g. "GrootLocomotionController",
+    # "HolosomaLocomotionController", or "SonicWholeBodyController". None disables it.
+    # Selecting "SonicWholeBodyController" implicitly switches the robot to the 64-D
+    # latent-token action/observation interface (``motion_token.{i}.pos`` action and a
+    # ``motion_token_state.{i}.pos`` state echo) so ``lerobot-rollout`` can drive a
+    # policy trained on SONIC motion tokens (e.g. nepyope/sonic_walk).
     controller: str | None = None

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

@@ -0,0 +1,27 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unitree G1 locomotion controllers (Groot, Holosoma, SONIC)."""
+
+from .gr00t_locomotion import GrootLocomotionController
+from .holosoma_locomotion import HolosomaLocomotionController
+from .sonic_whole_body import SonicWholeBodyController
+
+__all__ = [
+    "GrootLocomotionController",
+    "HolosomaLocomotionController",
+    "SonicWholeBodyController",
+]

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

@@ -21,7 +21,7 @@ import numpy as np
 import onnxruntime as ort
 from huggingface_hub import hf_hub_download
 
-from .g1_utils import (
+from ..g1_utils import (
     REMOTE_AXES,
     REMOTE_BUTTONS,
     G1_29_JointIndex,

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

@@ -22,7 +22,7 @@ import onnx
 import onnxruntime as ort
 from huggingface_hub import hf_hub_download
 
-from .g1_utils import (
+from ..g1_utils import (
     REMOTE_AXES,
     G1_29_JointArmIndex,
     G1_29_JointIndex,

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

@@ -0,0 +1,359 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""SONIC decoder whole-body controller for the Unitree G1 (token-only).
+
+Pure-Python/ONNX re-implementation of the *decode* half of NVIDIA's SONIC deploy stack.
+The encoder is intentionally absent: a token-output VLA (e.g. ``nepyope/sonic_walk``)
+supplies the 64-D latent ``motion_token`` directly each tick, and the SONIC **decoder**
+maps ``token + recent proprioception history`` to a residual action that is scaled and
+added onto the standing pose (``default_angles``) to produce 50 Hz joint-position targets
+for the robot's PD controller.
+
+Index spaces: joints exist in two orderings — **IsaacLab** (policy/training order) and
+**MuJoCo** (deploy order). ``ISAACLAB_TO_MUJOCO`` / ``MUJOCO_TO_ISAACLAB`` (in g1_utils)
+convert between them. Quaternions are scalar-first ``(w, x, y, z)``.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+
+import numpy as np
+import onnx
+import onnxruntime as ort
+from huggingface_hub import hf_hub_download
+
+from ..g1_utils import (
+    ISAACLAB_TO_MUJOCO,
+    MUJOCO_TO_ISAACLAB,
+    G1_29_JointIndex,
+    get_gravity_orientation,
+)
+
+logger = logging.getLogger(__name__)
+
+# ── Constants (hardware-validated; see the NVIDIA SONIC deploy reference) ──────
+CONTROL_DT = 0.02  # 50 Hz control period (s)
+TOKEN_DIM = 64  # decoder latent size
+
+# SONIC decoder checkpoint: NVIDIA's decoder ONNX re-packaged with its deploy constants
+# (kp/kd PD gains, the standing pose default_angles, and the residual action_scale) embedded
+# in the ONNX metadata; see upload_sonic_decoder.py for provisioning. The runtime loads the
+# model *and* all of these straight from the checkpoint (the Holosoma convention), so no
+# motor-physics math happens at deploy time.
+DEFAULT_SONIC_REPO_ID = "lerobot/sonic_decoder"
+DECODER_FILENAME = "model_decoder.onnx"
+DECODER_INPUT_DIM = 994  # token(64) + 10-frame proprio history + gravity
+
+
+def load_sonic_decoder(repo_id: str = DEFAULT_SONIC_REPO_ID):
+    """Load the SONIC decoder ONNX and its baked-in deploy constants from the checkpoint.
+
+    Returns ``(decoder_session, kp, kd, default_angles, action_scale, neutral_token)``. The
+    gains/pose/scale are (29,) float32 in IsaacLab joint order and ``neutral_token`` is the
+    (64,) float32 idle latent -- all read from the ONNX ``metadata_props`` rather than
+    recomputed/hardcoded at deploy time (mirrors ``holosoma_locomotion.load_policy``).
+    """
+    decoder_path = hf_hub_download(repo_id=repo_id, filename=DECODER_FILENAME)
+    so = ort.SessionOptions()
+    so.log_severity_level = 3  # quiet ORT logs
+    session = ort.InferenceSession(decoder_path, sess_options=so)
+    dec_dim = int(session.get_inputs()[0].shape[1])
+    if dec_dim != DECODER_INPUT_DIM:
+        raise RuntimeError(f"Unexpected decoder input dim {dec_dim} (expected {DECODER_INPUT_DIM})")
+
+    meta = {p.key: p.value for p in onnx.load(decoder_path, load_external_data=False).metadata_props}
+    required = ("kp", "kd", "default_angles", "action_scale", "neutral_token")
+    missing = [k for k in required if k not in meta]
+    if missing:
+        raise ValueError(
+            f"SONIC decoder ONNX at {repo_id} is missing metadata {missing}; "
+            "re-run upload_sonic_decoder.py to (re)provision the checkpoint."
+        )
+    arr = {k: np.array(json.loads(meta[k]), dtype=np.float32) for k in required}
+    logger.info("Loaded SONIC deploy constants from %s (%d joints)", repo_id, len(arr["kp"]))
+    return session, arr["kp"], arr["kd"], arr["default_angles"], arr["action_scale"], arr["neutral_token"]
+
+
+# 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
+# here so ``lerobot-rollout`` aggregates it into a 64-D ``observation.state``.
+TOKEN_STATE_PREFIX = "motion_token_state"  # nosec B105 - feature-key prefix, not a secret
+
+
+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"
+
+
+# Startup blend duration: over the first control ticks, linearly interpolate every joint
+# from the robot's initial measured pose into the policy's commanded target, so control
+# eases in without a snap on the first command.
+INIT_RAMP_S = 3.0
+
+
+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 SonicDecoder:
+    """Runs the SONIC decoder ONNX model and owns the proprioception history.
+
+    Each tick it appends the latest robot state to 10-frame history buffers, then maps the
+    supplied 64-D ``token`` + that history to a residual action added onto ``default_angles``.
+    The encoder is bypassed entirely (token supplied by the policy). ``default_angles`` and
+    ``action_scale`` are (29,) float32 in IsaacLab order, loaded from the checkpoint.
+    """
+
+    def __init__(self, decoder, default_angles, action_scale):
+        self.decoder = decoder
+        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 = 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
+        self.h_dq_mj = [np.zeros(29, np.float32)] * 10
+        self.h_ang = [np.zeros(3, np.float32)] * 10
+        self.h_act_mj = [np.zeros(29, np.float32)] * 10
+        self.h_quat = [np.array([1, 0, 0, 0], np.float32)] * 10
+
+    def reset(self):
+        """Clear the token and 10-frame proprioception history.
+
+        ``UnitreeG1.reset()`` relies on this so the first decoder outputs of a new episode
+        are not contaminated by the previous episode's state.
+        """
+        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
+        self.h_dq_mj = [np.zeros(29, np.float32)] * 10
+        self.h_ang = [np.zeros(3, np.float32)] * 10
+        self.h_act_mj = [np.zeros(29, np.float32)] * 10
+        self.h_quat = [np.array([1, 0, 0, 0], np.float32)] * 10
+
+    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)
+        # 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]
+        self.h_act_mj = [self.last_action_mj.copy()] + self.h_act_mj[:-1]
+        self.h_quat = [quat.copy()] + self.h_quat[:-1]
+
+    def build_decoder_obs(self):
+        """Assemble the 994-D decoder input: token + 10-frame proprioception history + gravity."""
+        obs = np.zeros(994, np.float32)
+        off = 0
+        obs[off : off + 64] = self.token
+        off += 64
+        for h, sz in [
+            (list(reversed(self.h_ang)), 3),
+            (list(reversed(self.h_q_mj)), 29),
+            (list(reversed(self.h_dq_mj)), 29),
+            (list(reversed(self.h_act_mj)), 29),
+        ]:
+            for f in range(10):
+                obs[off : off + sz] = h[f]
+                off += sz
+        for q in reversed(self.h_quat):
+            obs[off : off + 3] = get_gravity_orientation(q)
+            off += 3
+        assert off == 994, f"Decoder obs mismatch: {off}"
+        return obs
+
+    def step(self, robot_obs, token, debug=False):
+        """One control tick: read robot obs, decode the supplied token -> joint targets.
+
+        Args:
+            robot_obs: dict with ``<joint>.q``/``.dq`` and ``imu.*`` fields.
+            token: 64-D latent supplied by the policy (encoder bypassed).
+            debug: log action/delta norms.
+
+        Returns:
+            dict of ``<joint>.q`` target positions (rad) in IsaacLab joint order.
+        """
+        self.token = np.asarray(token, np.float32)
+        jnames = [m.name for m in G1_29_JointIndex]
+        q = np.array(
+            [
+                robot_obs.get(f"{n}.q", self.default_angles[m.value])
+                for m, n in zip(G1_29_JointIndex, jnames, strict=False)
+            ],
+            np.float32,
+        )
+        dq = np.array([robot_obs.get(f"{n}.dq", 0.0) for n in jnames], np.float32)
+        quat = np.array(
+            [
+                robot_obs.get("imu.quat.w", 1),
+                robot_obs.get("imu.quat.x", 0),
+                robot_obs.get("imu.quat.y", 0),
+                robot_obs.get("imu.quat.z", 0),
+            ],
+            np.float32,
+        )
+        ang = np.array([robot_obs.get(f"imu.gyro.{a}", 0) for a in "xyz"], np.float32)
+        self.update_history(q, dq, ang, quat)
+        action_mj = (
+            self.decoder.run(None, {self.decoder_input: self.build_decoder_obs().reshape(1, -1)})[0]
+            .squeeze()
+            .astype(np.float32)
+        )
+        self.last_action_mj = action_mj.copy()
+        target = self.default_angles + action_mj[ISAACLAB_TO_MUJOCO] * self.action_scale
+        if debug:
+            delta = target - q
+            logger.debug(
+                "token_norm=%.4f action_norm=%.4f delta_max=%.4f delta_rms=%.4f",
+                np.linalg.norm(self.token),
+                np.linalg.norm(action_mj),
+                np.max(np.abs(delta)),
+                np.sqrt(np.mean(delta**2)),
+            )
+        return {f"{m.name}.q": float(target[m.value]) for m in G1_29_JointIndex}
+
+
+class SonicWholeBodyController:
+    """Full-body SONIC controller for UnitreeG1's background controller thread.
+
+    Token-only deploy: the encoder is bypassed; each tick the decoder consumes a 64-D
+    latent token supplied directly by the policy.
+    """
+
+    control_dt = CONTROL_DT
+    full_body = True
+
+    def __init__(self):
+        logger.info("Loading SONIC whole-body controller...")
+        decoder_sess, self.kp, self.kd, default_angles, action_scale, self._neutral_token = (
+            load_sonic_decoder()
+        )
+        self._default_angles = default_angles
+        self.controller = SonicDecoder(decoder_sess, default_angles, action_scale)
+
+        # Startup blend: ease from the robot's initial pose into the first commanded policy
+        # targets over INIT_RAMP_S (captured on the first control tick).
+        self._init_ramp_steps = max(1, round(INIT_RAMP_S / CONTROL_DT))
+        self._init_step = 0
+        self._start_pose: dict[str, float] = {}
+
+        # 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)")
+
+    @property
+    def action_features(self) -> dict[str, type]:
+        """64-D latent-token action space (``motion_token.{i}.pos``).
+
+        ``lerobot-rollout`` maps a 64-D policy output straight onto these keys, which the
+        decoder consumes with the encoder bypassed.
+        """
+        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)}
+
+    def _startup_blend(self, obs: dict, out: dict) -> dict:
+        """Ease into policy control at startup: for the first ``INIT_RAMP_S`` seconds,
+        interpolate between the robot's pose captured on the first tick and the policy's
+        live commanded target, so the handoff has no snap.
+
+        ``out`` is the policy's ``<joint>.q`` target dict for this tick; the blend ratio
+        climbs 0->1 over the ramp, after which the raw policy target passes through.
+        """
+        if self._init_step >= self._init_ramp_steps or not out:
+            return out
+        if self._init_step == 0:
+            # Capture the robot's actual pose as the interpolation start point.
+            self._start_pose = {
+                f"{m.name}.q": float(obs.get(f"{m.name}.q", self._default_angles[m.value]))
+                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, obs: dict) -> dict:
+        if not obs:
+            return {}
+
+        # 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:
+            # 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()
+        # 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))
+
+    def reset(self):
+        self.controller.reset()
+        self._init_step = 0  # re-run the startup blend after a reset
+        self._start_pose = {}
+        # Drop the held token so the neutral token is re-seeded after a reset.
+        self._last_token = None

src/lerobot/robots/unitree_g1/g1_utils.py

@@ -23,6 +23,47 @@ import numpy as np
 
 NUM_MOTORS = 29
 
+# Joint-order permutations between the two 29-DoF layouts used across the G1 stack:
+# IsaacLab (policy/training order) and MuJoCo (deploy order). ``a[ISAACLAB_TO_MUJOCO]``
+# reorders an IsaacLab-ordered vector into MuJoCo order, and vice-versa.
+ISAACLAB_TO_MUJOCO = np.array(
+    [
+        0,
+        3,
+        6,
+        9,
+        13,
+        17,
+        1,
+        4,
+        7,
+        10,
+        14,
+        18,
+        2,
+        5,
+        8,
+        11,
+        15,
+        19,
+        21,
+        23,
+        25,
+        27,
+        12,
+        16,
+        20,
+        22,
+        24,
+        26,
+        28,
+    ],
+    dtype=np.int32,
+)
+# The two orderings are inverses of each other, so derive one from the other (argsort) to
+# guarantee they can never drift out of sync.
+MUJOCO_TO_ISAACLAB = np.argsort(ISAACLAB_TO_MUJOCO).astype(np.int32)
+
 REMOTE_AXES = ("remote.lx", "remote.ly", "remote.rx", "remote.ry")
 REMOTE_BUTTONS = tuple(f"remote.button.{i}" for i in range(16))
 REMOTE_KEYS = REMOTE_AXES + REMOTE_BUTTONS
@@ -68,8 +109,9 @@ def make_locomotion_controller(name: str | None):
     if name is None:
         return None
     controllers = {
-        "GrootLocomotionController": "lerobot.robots.unitree_g1.gr00t_locomotion",
-        "HolosomaLocomotionController": "lerobot.robots.unitree_g1.holosoma_locomotion",
+        "GrootLocomotionController": "lerobot.robots.unitree_g1.controllers.gr00t_locomotion",
+        "HolosomaLocomotionController": "lerobot.robots.unitree_g1.controllers.holosoma_locomotion",
+        "SonicWholeBodyController": "lerobot.robots.unitree_g1.controllers.sonic_whole_body",
     }
     module_path = controllers.get(name)
     if module_path is None:

src/lerobot/robots/unitree_g1/unitree_g1.py

@@ -34,7 +34,6 @@ from .config_unitree_g1 import UnitreeG1Config
 from .g1_kinematics import G1_29_ArmIK
 from .g1_utils import (
     REMOTE_AXES,
-    REMOTE_KEYS,
     G1_29_JointArmIndex,
     G1_29_JointIndex,
     default_remote_input,
@@ -148,9 +147,8 @@ class UnitreeG1(Robot):
 
         self.arm_ik = G1_29_ArmIK() if config.gravity_compensation else None
 
-        # Lower-body controller loaded dynamically
+        # Lower-body / whole-body controller loaded dynamically
         self.controller: LocomotionController | None = make_locomotion_controller(config.controller)
-
         # Controller thread state
         self._controller_thread = None
         self._controller_action_lock = threading.Lock()
@@ -233,13 +231,22 @@ class UnitreeG1(Robot):
 
     @cached_property
     def observation_features(self) -> dict[str, type | tuple]:
-        return {**self._motors_ft, **self._cameras_ft}
+        # Controllers may contribute their own proprio features (e.g. SONIC's token state).
+        controller_ft = getattr(self.controller, "observation_features", {})
+        return {**self._motors_ft, **controller_ft, **self._cameras_ft}
 
     @cached_property
     def action_features(self) -> dict[str, type]:
+        # No controller configured at all: raw 29-DoF joint teleop.
         if self.controller is None:
             return {f"{G1_29_JointIndex(motor).name}.q": float for motor in G1_29_JointIndex}
 
+        # Controllers may define their own action space (e.g. SONIC's 64-D latent token).
+        controller_ft = getattr(self.controller, "action_features", None)
+        if controller_ft is not None:
+            return dict(controller_ft)
+
+        # Locomotion controllers (GR00T / Holosoma): arm joint targets + joystick axes.
         arm_features = {f"{G1_29_JointArmIndex(motor).name}.q": float for motor in G1_29_JointArmIndex}
         remote_features = dict.fromkeys(REMOTE_AXES, float)
         return {**arm_features, **remote_features}
@@ -271,8 +278,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:
@@ -461,6 +472,11 @@ class UnitreeG1(Robot):
         if lowstate.wireless_remote:
             obs["wireless_remote"] = lowstate.wireless_remote
 
+        # Controller-contributed observation (e.g. SONIC echoes its last decoded token as
+        # observation.state so a token-output VLA closes the loop on its own previous token).
+        if self.controller is not None and hasattr(self.controller, "observation_state"):
+            obs.update(self.controller.observation_state())
+
         # Cameras - read images from ZMQ cameras
         for cam_name, cam in self._cameras.items():
             if getattr(cam, "use_rgb", True):
@@ -476,6 +492,10 @@ class UnitreeG1(Robot):
             # Controller thread owns legs/waist. Here we only update joystick inputs
             # and publish arm targets from the teleoperator.
             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
@@ -503,11 +523,12 @@ class UnitreeG1(Robot):
         return action
 
     def _update_controller_action(self, action: RobotAction) -> None:
-        """Update controller input state from incoming teleop action."""
+        """Forward incoming teleop action values into ``controller_input``; each controller
+        reads only the keys it understands."""
         with self._controller_action_lock:
-            for key in REMOTE_KEYS:
-                if key in action:
-                    self.controller_input[key] = action[key]
+            for key, value in action.items():
+                if isinstance(key, str) and value is not None:
+                    self.controller_input[key] = value
 
     @property
     def is_calibrated(self) -> bool: