From 77259f436eb7b113ef6d58ab04cdddd3af2b2a62 Mon Sep 17 00:00:00 2001 From: Martino Russi Date: Mon, 27 Jul 2026 11:05:42 +0200 Subject: [PATCH] feat(unitree_g1): use captured neutral SONIC token instead of zeros The all-zero token is off the encoder's learned FSQ manifold and decodes to a slightly goofy stance. Replace it with a NEUTRAL_TOKEN captured from the encoder's own idle output in sim (stored as integer FSQ codes, rescaled by the encoder's 1/16 quantization step to an exact on-grid token). token_mode now seeds this neutral, and the onboard sender starts observation.state from it so the first inference sees the token the decoder is actually holding. --- .../controllers/sonic_whole_body.py | 23 ++++++++++++++++--- .../unitree_g1/infer_sonic_g1_onboard.py | 11 ++++++--- 2 files changed, 28 insertions(+), 6 deletions(-) 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 2611ec530..8320f199a 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py +++ b/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py @@ -77,6 +77,22 @@ logger = logging.getLogger(__name__) # eases in without a snap on the first command. INIT_RAMP_S = 3.0 +# Neutral ("zero pose") SONIC token, held by token_mode until the first real token +# arrives. Captured from the encoder's own output while the robot stood idle in sim +# (capture_neutral_token.py): the encoder is an FSQ bottleneck (~5 bit/dim, 15.5 half- +# width, Div(16)), so its tokens live on the 1/16 grid. We store the integer FSQ codes +# and rescale by the same 1/16 step, giving an exact on-grid token -- unlike the literal +# all-zero token, which is off the encoder's learned manifold and decodes to a slightly +# goofy stance. This one decodes to a stable, natural standing pose. +_NEUTRAL_TOKEN_CODES = np.array( + [-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], + dtype=np.float32, +) +NEUTRAL_TOKEN = _NEUTRAL_TOKEN_CODES / 16.0 # FSQ Div(16): integer codes -> on-grid token + def _extract_wb34_from_action(action: dict | None) -> np.ndarray | None: """Reassemble a dense (34,) whole-body command from ``wb.{i}.pos`` keys, or None. @@ -359,9 +375,10 @@ class SonicWholeBodyController: 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 neutral - # (all-zero) token, which the decoder maps to a stable neutral stance. - self._last_token = np.zeros(TOKEN_DIM, dtype=np.float32) + # Token-driven deploy, but no token has arrived yet: hold the captured + # neutral token (NEUTRAL_TOKEN), which the decoder maps to a stable, natural + # standing pose (the encoder's own idle output; see NEUTRAL_TOKEN). + self._last_token = NEUTRAL_TOKEN.copy() if self._last_token is not None: # Either a fresh token this tick or the last one received (held between the # ~30 Hz token stream and the ~50 Hz control loop). diff --git a/src/lerobot/robots/unitree_g1/infer_sonic_g1_onboard.py b/src/lerobot/robots/unitree_g1/infer_sonic_g1_onboard.py index 657366b40..2fbc8ebdb 100644 --- a/src/lerobot/robots/unitree_g1/infer_sonic_g1_onboard.py +++ b/src/lerobot/robots/unitree_g1/infer_sonic_g1_onboard.py @@ -56,7 +56,11 @@ from lerobot.cameras.zmq import ZMQCamera, ZMQCameraConfig from lerobot.configs.policies import PreTrainedConfig from lerobot.policies.factory import get_policy_class, make_pre_post_processors from lerobot.policies.utils import prepare_observation_for_inference -from lerobot.robots.unitree_g1.controllers.sonic_whole_body import TOKEN_DIM, token_action_key +from lerobot.robots.unitree_g1.controllers.sonic_whole_body import ( + NEUTRAL_TOKEN, + TOKEN_DIM, + token_action_key, +) logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s", force=True) logger = logging.getLogger("sonic_sender") @@ -130,8 +134,9 @@ def main() -> None: signal.signal(signal.SIGTERM, lambda *_: stop.__setitem__("flag", True)) # observation.state = the token currently executing on the robot (last one we sent); - # start at zeros, matching the decoder's zero-seeded initial state. - prev_token = np.zeros(TOKEN_DIM, dtype=np.float32) + # start at the neutral token the decoder holds before the first send, so the very + # first inference sees the true executing token (not zeros). + prev_token = NEUTRAL_TOKEN.copy() period = 1.0 / args.fps n = 0 t_infer_total = 0.0