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.
This commit is contained in:
Martino Russi
2026-07-27 11:05:42 +02:00
parent 85f5c3606d
commit 77259f436e
2 changed files with 28 additions and 6 deletions
@@ -77,6 +77,22 @@ logger = logging.getLogger(__name__)
# eases in without a snap on the first command. # eases in without a snap on the first command.
INIT_RAMP_S = 3.0 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: 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. """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: if token is not None:
self._last_token = token self._last_token = token
elif self._last_token is None and self.token_mode: elif self._last_token is None and self.token_mode:
# Token-driven deploy, but no token has arrived yet: hold the neutral # Token-driven deploy, but no token has arrived yet: hold the captured
# (all-zero) token, which the decoder maps to a stable neutral stance. # neutral token (NEUTRAL_TOKEN), which the decoder maps to a stable, natural
self._last_token = np.zeros(TOKEN_DIM, dtype=np.float32) # standing pose (the encoder's own idle output; see NEUTRAL_TOKEN).
self._last_token = NEUTRAL_TOKEN.copy()
if self._last_token is not None: if self._last_token is not None:
# Either a fresh token this tick or the last one received (held between the # Either a fresh token this tick or the last one received (held between the
# ~30 Hz token stream and the ~50 Hz control loop). # ~30 Hz token stream and the ~50 Hz control loop).
@@ -56,7 +56,11 @@ from lerobot.cameras.zmq import ZMQCamera, ZMQCameraConfig
from lerobot.configs.policies import PreTrainedConfig from lerobot.configs.policies import PreTrainedConfig
from lerobot.policies.factory import get_policy_class, make_pre_post_processors from lerobot.policies.factory import get_policy_class, make_pre_post_processors
from lerobot.policies.utils import prepare_observation_for_inference 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) logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s", force=True)
logger = logging.getLogger("sonic_sender") logger = logging.getLogger("sonic_sender")
@@ -130,8 +134,9 @@ def main() -> None:
signal.signal(signal.SIGTERM, lambda *_: stop.__setitem__("flag", True)) signal.signal(signal.SIGTERM, lambda *_: stop.__setitem__("flag", True))
# observation.state = the token currently executing on the robot (last one we sent); # observation.state = the token currently executing on the robot (last one we sent);
# start at zeros, matching the decoder's zero-seeded initial state. # start at the neutral token the decoder holds before the first send, so the very
prev_token = np.zeros(TOKEN_DIM, dtype=np.float32) # first inference sees the true executing token (not zeros).
prev_token = NEUTRAL_TOKEN.copy()
period = 1.0 / args.fps period = 1.0 / args.fps
n = 0 n = 0
t_infer_total = 0.0 t_infer_total = 0.0