From 85f5c3606d71cfb513369785ead58fc33b30c471 Mon Sep 17 00:00:00 2001 From: Martino Russi Date: Mon, 27 Jul 2026 10:34:42 +0200 Subject: [PATCH] feat(unitree_g1): hold neutral SONIC token until first command Move the token-hold idle logic into SonicWholeBodyController via a token_mode flag (set by UnitreeG1 when sonic_token_action is enabled): before any real token arrives the decoder is fed the all-zero neutral token (stable neutral stance), and afterwards the last received token is held between control ticks (the ~50 Hz control loop outruns the ~30 Hz token stream). Living in the controller, this applies uniformly to run_g1_onboard, lerobot-rollout and the sim replays, so the explicit neutral seeding in run_g1_onboard is removed. --- .../controllers/sonic_whole_body.py | 23 ++++++++++++++++++- .../robots/unitree_g1/run_g1_onboard.py | 3 +++ src/lerobot/robots/unitree_g1/unitree_g1.py | 5 ++++ 3 files changed, 30 insertions(+), 1 deletion(-) 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 4d44d9dbe..2611ec530 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py +++ b/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py @@ -229,6 +229,17 @@ class SonicWholeBodyController: # forwarded to the pipeline as ``delta_heading`` so turn commands take effect. self._heading = 0.0 + # Token-interface state. ``token_mode`` is set True by the robot when the deploy + # is token-driven (``UnitreeG1Config.sonic_token_action``): the controller then + # holds a stable *neutral* (all-zero) 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_onboard, + # lerobot-rollout and the sim replays. ``token_mode`` stays False for the dense + # 34-D whole-body / OpenHLM path, which keeps its own "hold last target" idle. + self.token_mode = False + self._last_token: np.ndarray | None = None + logger.info("SONIC ready (encoder/decoder, 34-D whole-body command path)") def _run_wholebody34(self, obs: dict, wb: np.ndarray) -> dict: @@ -346,7 +357,15 @@ class SonicWholeBodyController: # Checked before the joint path so a token action takes precedence. token = _extract_token_from_action(action) if token is not None: - return self._startup_blend(obs, self._run_token(obs, token)) + 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) + 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). + return self._startup_blend(obs, self._run_token(obs, self._last_token)) # Dense 34-D whole-body command (OpenHLM / pi0.5 joint interface): a single # vector per tick drives the mode-0 encoder reference directly. Until the @@ -372,6 +391,8 @@ class SonicWholeBodyController: self._wb_traj.clear() self._wb_quat_traj.clear() self._heading = 0.0 + # Drop the held token so token_mode re-seeds the neutral token after a reset. + self._last_token = None def shutdown(self): self._runtime.shutdown() diff --git a/src/lerobot/robots/unitree_g1/run_g1_onboard.py b/src/lerobot/robots/unitree_g1/run_g1_onboard.py index bd2f99cba..905766350 100644 --- a/src/lerobot/robots/unitree_g1/run_g1_onboard.py +++ b/src/lerobot/robots/unitree_g1/run_g1_onboard.py @@ -139,6 +139,9 @@ def main() -> None: robot = UnitreeG1(cfg) logger.info("Connecting onboard robot (controller=%s, token=%s)...", args.controller, args.sonic_token_action) robot.connect() + # Note: with --sonic-token-action the SonicWholeBodyController holds a neutral + # (all-zero) token until the first laptop token arrives, then holds the last token + # between ticks -- see SonicWholeBodyController.token_mode (set from config). grippers: dict[str, Gripper] = {} if args.grippers: diff --git a/src/lerobot/robots/unitree_g1/unitree_g1.py b/src/lerobot/robots/unitree_g1/unitree_g1.py index a6ed98f32..823d7e33d 100644 --- a/src/lerobot/robots/unitree_g1/unitree_g1.py +++ b/src/lerobot/robots/unitree_g1/unitree_g1.py @@ -169,6 +169,11 @@ class UnitreeG1(Robot): # Lower-body controller loaded dynamically self.controller: LocomotionController | None = make_locomotion_controller(config.controller) + # Token-driven deploy: let a SONIC controller hold a neutral token until the + # first real one arrives, then hold the last token between control ticks. + if config.sonic_token_action 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