feat(unitree_g1): 64-D SONIC token interface for lerobot-rollout + GR00T waist override

Add a token-output VLA path (sonic_token_action) so a policy trained on 64-D SONIC
motion tokens (e.g. nepyope/sonic_walk) drives the decoder directly via lerobot-rollout:
the robot advertises a 64-D motion_token.{i}.pos action and echoes the last commanded
token as a 64-D observation.state (motion_token_state.{i}.pos), encoder bypassed.

Also:
- gr00t_locomotion: allow an external upper-body IK to override the 3 waist joints, and
  cap ORT to 1 intra/inter thread so the 50Hz loop doesn't stutter under contention.
- sonic_pipeline: make_ort_session_options takes optional thread caps; report the
  provider actually bound.
- unitree_g1: build the sim env with publish_images=False/cameras=[] to avoid the
  offscreen EGL context crash (we drive image policies from recorded/live frames), and
  guard the startup sim-step race (zero-norm pelvis quat) so the sim thread survives.

src/lerobot/robots/unitree_g1/config_unitree_g1.py

@@ -78,6 +78,10 @@ class UnitreeG1Config(RobotConfig):
     # ``hand_closed_pose`` (7 joints: thumb_0/1/2, middle_0/1, index_0/1). Flip signs
     # in ``hand_closed_pose`` if fingers curl the wrong way.
     publish_hands: bool = False
+    # When False, connect() does not start the background controller thread, so a
+    # caller can drive the controller synchronously (one decode per fed action),
+    # reproducing the deploy's single 50Hz control clock for faithful replay.
+    run_controller_thread: bool = True
     hand_open_grip_value: float = 1.0
     hand_closed_grip_value: float = 0.0
     hand_closed_pose: list[float] = field(default_factory=lambda: [1.0, 0.9, 0.9, 1.3, 1.3, 1.3, 1.3])
@@ -92,6 +96,16 @@ class UnitreeG1Config(RobotConfig):
     replay_camera_map: dict[str, str] = field(default_factory=dict)
     replay_camera_loop: bool = True
 
+    # Token-output VLA interface for the SONIC decoder. When True (and the controller
+    # is ``SonicWholeBodyController``), the robot advertises a 64-D latent-token action
+    # space (``motion_token.{i}.pos``) instead of the 34-D whole-body command, and
+    # exposes the last commanded token as a 64-D ``observation.state``
+    # (``motion_token_state.{i}.pos``). This lets ``lerobot-rollout`` drive a policy
+    # that was trained with 64-D SONIC motion tokens as both state and action
+    # (e.g. nepyope/sonic_walk): the decoder consumes the token directly, encoder
+    # bypassed. Ignored unless a SONIC whole-body controller is active.
+    sonic_token_action: bool = False
+
     # Compensates for gravity on the unitree's arms using the arm ik solver
     gravity_compensation: bool = False
 

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

@@ -77,9 +77,15 @@ def load_groot_policies(
         filename="GR00T-WholeBodyControl-Walk.onnx",
     )
 
-    # Load ONNX policies
-    policy_balance = ort.InferenceSession(balance_path)
-    policy_walk = ort.InferenceSession(walk_path)
+    # Load ONNX policies with a capped thread pool. GR00T runs at 50 Hz in a
+    # background thread alongside the (torch) upper-body policy, IK and sim; letting
+    # ORT grab every core starves those and makes the whole rollout stutter. These
+    # are small MLPs, so 1 thread is both enough and lowest-latency.
+    from .sonic_pipeline import make_ort_session_options
+
+    so = make_ort_session_options(intra_op_num_threads=1, inter_op_num_threads=1)
+    policy_balance = ort.InferenceSession(balance_path, sess_options=so)
+    policy_walk = ort.InferenceSession(walk_path, sess_options=so)
 
     logger.info("GR00T policies loaded successfully")
 
@@ -206,6 +212,16 @@ class GrootLocomotionController:
         # Transform action back to target joint positions
         target_dof_pos_15 = GROOT_DEFAULT_ANGLES[:15] + self.groot_action * ACTION_SCALE
 
+        # Waist override: an external upper-body IK can command the 3 waist joints
+        # (indices 12/13/14) via ``kWaist{Yaw,Roll,Pitch}.q`` in the action dict. When
+        # present, we substitute the balance policy's waist target so the torso tracks
+        # the IK while the policy keeps only the legs balanced. Single-publisher stays
+        # intact (this thread still owns joints 0-14).
+        for idx in (G1_29_JointIndex.kWaistYaw, G1_29_JointIndex.kWaistRoll, G1_29_JointIndex.kWaistPitch):
+            key = f"{idx.name}.q"
+            if key in action and action[key] is not None:
+                target_dof_pos_15[idx.value] = float(action[key])
+
         # Build action dict
         action_dict = {}
         for i in range(15):

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

@@ -277,10 +277,21 @@ def ort_providers(force_cpu: bool = False) -> list[str]:
     return ["CPUExecutionProvider"]
 
 
-def make_ort_session_options():
-    """Build ONNX Runtime SessionOptions (quiet logging, default threading)."""
+def make_ort_session_options(intra_op_num_threads: int | None = None,
+                             inter_op_num_threads: int | None = None):
+    """Build ONNX Runtime SessionOptions (quiet logging).
+
+    Pass thread counts to cap ORT's CPU pool. These tiny MLP policies are latency-
+    bound, not throughput-bound, so letting ORT grab every core just starves the
+    real-time control loop / torch policy / IK solver and causes stutter. 1 intra +
+    1 inter thread is plenty and lowest-latency for a per-step MLP inference.
+    """
     so = ort.SessionOptions()
     so.log_severity_level = 3
+    if intra_op_num_threads is not None:
+        so.intra_op_num_threads = intra_op_num_threads
+    if inter_op_num_threads is not None:
+        so.inter_op_num_threads = inter_op_num_threads
     return so
 
 

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

@@ -18,12 +18,12 @@
 
 from __future__ import annotations
 
-import logging
 from collections import deque
+import logging
 from typing import TYPE_CHECKING
 
-import numpy as np
 from huggingface_hub import hf_hub_download
+import numpy as np
 
 from lerobot.utils.import_utils import _onnxruntime_available, require_package
 
@@ -38,12 +38,33 @@ from .sonic_pipeline import (
     CONTROL_DT,
     DEFAULT_ANGLES,
     ENCODER_UPDATE_EVERY,
+    TOKEN_DIM,
     PlannerController,
     compute_kp_kd,
     make_ort_session_options,
     ort_providers,
 )
 
+# Action-feature prefix for the latent-token interface (see _extract_token_from_action).
+TOKEN_ACTION_PREFIX = "motion_token"
+# 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"
+
+
+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"
+
 if TYPE_CHECKING or _onnxruntime_available:
     import onnxruntime as ort
 else:
@@ -79,6 +100,26 @@ def _extract_wb34_from_action(action: dict | None) -> np.ndarray | None:
     )
 
 
+def _extract_token_from_action(action: dict | None) -> np.ndarray | None:
+    """Reassemble a dense (64,) latent token from ``motion_token.{i}`` keys, or None.
+
+    This is the token-only replay interface: instead of a joint reference driving the
+    encoder, the caller supplies the 64-D encoder latent directly (e.g. a recorded
+    ``action.motion_token`` column), 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,
+    )
+
+
 def _wb34_to_reference(wb: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
     """Map a 34-D OpenHLM whole-body command to a SONIC mode-0 reference.
 
@@ -121,12 +162,23 @@ class SonicRuntime:
         decoder_path = hf_hub_download(repo_id="nvidia/GEAR-SONIC", filename="model_decoder.onnx")
 
         providers = ort_providers(force_cpu=force_cpu)
-        self.use_gpu = providers[0] == "CUDAExecutionProvider"
         so = make_ort_session_options()
 
         encoder_sess = ort.InferenceSession(encoder_path, sess_options=so, providers=providers)
         decoder_sess = ort.InferenceSession(decoder_path, sess_options=so, providers=providers)
 
+        # Report the provider actually bound, not the one requested: ORT silently falls
+        # back to CPU if CUDA can't load (e.g. libcudnn not on LD_LIBRARY_PATH), and a
+        # CPU decoder drifts the closed-loop heading. Warn loudly so it can't hide.
+        self.use_gpu = decoder_sess.get_providers()[0] == "CUDAExecutionProvider"
+        if not force_cpu and not self.use_gpu:
+            print(
+                "[SONIC] WARNING: decoder bound to CPUExecutionProvider (CUDA unavailable). "
+                "Closed-loop replay/control will drift. Ensure libcudnn is on LD_LIBRARY_PATH "
+                "(site-packages/nvidia/*/lib).",
+                flush=True,
+            )
+
         self.kp, self.kd = compute_kp_kd()
         self.controller = PlannerController(encoder_sess, decoder_sess)
 
@@ -244,6 +296,20 @@ class SonicWholeBodyController:
         self._wb_step += 1
         return out
 
+    def _run_token(self, obs: dict, token: np.ndarray) -> dict:
+        """Decode a supplied 64-D latent token directly (encoder bypassed).
+
+        Token-only replay: set the pipeline's cached token to the supplied one and run
+        a decode-only step (``update_encoder=False``). The decoder still closes the loop
+        on live proprioception (history is refreshed inside ``step`` from ``obs``); only
+        the encoder — which would recompute the token from a motion reference — is
+        skipped. Returns the ``<joint>.q`` target dict.
+        """
+        c = self.controller
+        c.token = np.asarray(token, np.float32)
+        self._wb_step += 1
+        return c.step(obs, update_encoder=False, debug=False)
+
     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
@@ -275,6 +341,13 @@ class SonicWholeBodyController:
             return {}
         obs = lowstate_to_obs(lowstate)
 
+        # Token-only interface (latent replay / token-output VLA): a dense 64-D
+        # ``motion_token.{i}`` command is decoded directly, bypassing the encoder.
+        # 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))
+
         # 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
         # policy produces one, hold (no command) so the robot keeps its last target.

src/lerobot/robots/unitree_g1/unitree_g1.py

@@ -179,6 +179,16 @@ class UnitreeG1(Robot):
         if config.replay_camera_parquet and config.replay_camera_map:
             self._load_replay_frames()
 
+        # Token-mode state: last 64-D SONIC latent token commanded by the policy,
+        # echoed back as ``observation.state`` so a token-output VLA closes the loop
+        # on its own previous token (see ``sonic_token_action``). Seeded to zeros;
+        # the controller's startup blend eases joints in regardless.
+        self._last_token: np.ndarray | None = None
+        if config.sonic_token_action:
+            from .controllers.sonic_whole_body import TOKEN_DIM
+
+            self._last_token = np.zeros(TOKEN_DIM, dtype=np.float32)
+
     def _load_replay_frames(self) -> None:
         """Load only the mapped parquet columns (encoded frames); decode on demand."""
         import pyarrow.parquet as pq
@@ -223,7 +233,17 @@ class UnitreeG1(Robot):
 
             # Step simulation if in simulation mode
             if self.config.is_simulation and self.sim_env is not None:
+                try:
                     self.sim_env.step()
+                except ValueError as e:
+                    # Startup race: the sim thread can step once before reset() has
+                    # written a valid base pose, giving a zero-norm pelvis quaternion
+                    # (scipy>=1.11 raises instead of normalizing). Skip and retry so
+                    # the thread survives instead of dying and freezing the sim.
+                    if "zero norm" not in str(e).lower():
+                        raise
+                    time.sleep(self.control_dt)
+                    continue
 
             msg = self.lowstate_subscriber.Read()
             if msg is not None:
@@ -299,12 +319,27 @@ class UnitreeG1(Robot):
         (OpenHLM / pi0.5). These ``.pos`` scalars are aggregated by the rollout
         pipeline into a single 34-D ``observation.state`` for the policy.
         """
+        if self.config.sonic_token_action:
+            return {}
         if not getattr(self.controller, "wb_action", False):
             return {}
         from .g1_utils import WB_ACTION_DIM
 
         return {f"wb_state.{i}.pos": float for i in range(WB_ACTION_DIM)}
 
+    @property
+    def _token_state_ft(self) -> dict[str, type]:
+        """64-D SONIC latent-token proprio state (``motion_token_state.{i}.pos``).
+
+        Exposed only in ``sonic_token_action`` mode; aggregated by the rollout into a
+        64-D ``observation.state`` (the last token the policy commanded).
+        """
+        if not self.config.sonic_token_action:
+            return {}
+        from .controllers.sonic_whole_body import TOKEN_DIM, token_state_key
+
+        return {token_state_key(i): float for i in range(TOKEN_DIM)}
+
     @property
     def _empty_cameras_ft(self) -> dict[str, tuple]:
         """Synthetic zero-image cameras (see ``UnitreeG1Config.empty_cameras``)."""
@@ -323,6 +358,7 @@ class UnitreeG1(Robot):
         return {
             **self._motors_ft,
             **self._wb_state_ft,
+            **self._token_state_ft,
             **self._empty_cameras_ft,
             **self._replay_cameras_ft,
             **self._cameras_ft,
@@ -333,6 +369,14 @@ class UnitreeG1(Robot):
         if self.controller is None:
             return {f"{G1_29_JointIndex(motor).name}.q": float for motor in G1_29_JointIndex}
 
+        # Token-output VLA (SONIC decoder): advertise a 64-D latent-token action space
+        # (``motion_token.{i}.pos``) so ``lerobot-rollout`` maps a 64-D policy output
+        # straight onto the decoder, bypassing the encoder.
+        if self.config.sonic_token_action:
+            from .controllers.sonic_whole_body import TOKEN_DIM, token_action_key
+
+            return {token_action_key(i): float for i in range(TOKEN_DIM)}
+
         # Dense whole-body controllers (SONIC / OpenHLM, pi0.5) consume a single
         # 34-D command per tick. Expose it as ``wb.{i}.pos`` joint-position features
         # so ``lerobot-rollout`` maps a 34-D policy output straight onto the robot.
@@ -402,10 +446,26 @@ class UnitreeG1(Robot):
     def connect(self, calibrate: bool = True) -> None:  # connect to DDS
         # Initialize DDS channel and simulation environment
         if self.config.is_simulation:
-            from lerobot.envs import make_env
+            from lerobot.envs.utils import (
+                _download_hub_file,
+                _import_hub_module,
+                _normalize_hub_result,
+            )
 
             self._ChannelFactoryInitialize(0, "lo")
-            self._env_wrapper = make_env("lerobot/unitree-g1-mujoco", trust_remote_code=True)
+            # Call the hub env's make_env directly so we can disable the offscreen
+            # head_camera renderer. We drive image-conditioned policies from recorded
+            # frames (see replay_camera_parquet / external obs), never the sim's own
+            # camera, so building a MuJoCo offscreen GL context is pure liability: it
+            # crashes with "Failed to make the EGL context current" when GLFW/SDL
+            # already own a context, killing the sim thread and hanging on
+            # "Waiting for robot state...". publish_images=False -> no renderer.
+            repo_id, _, local_file, _ = _download_hub_file(
+                "lerobot/unitree-g1-mujoco", True, None
+            )
+            hub_mod = _import_hub_module(local_file, repo_id)
+            raw = hub_mod.make_env(n_envs=1, use_async_envs=False, publish_images=False, cameras=[])
+            self._env_wrapper = _normalize_hub_result(raw)
             # Extract the actual gym env from the dict structure
             self.sim_env = self._env_wrapper["hub_env"][0].envs[0]
         else:
@@ -470,12 +530,16 @@ class UnitreeG1(Robot):
             self.msg.motor_cmd[joint].kd = self.kd[joint.value]
             self.msg.motor_cmd[joint].q = lowstate.motor_state[joint.value].q
 
-        # Start controller thread if enabled
-        if self.controller is not None:
+        # Start controller thread if enabled. Skipped when run_controller_thread is
+        # False so a caller can step the controller synchronously (faithful replay).
+        if self.controller is not None and self.config.run_controller_thread:
             self._controller_thread = threading.Thread(target=self._controller_loop, daemon=True)
             self._controller_thread.start()
             fps = int(1.0 / self.controller.control_dt)
             logger.info(f"Controller thread started ({fps}Hz)")
+        elif self.controller is not None:
+            logger.info("Controller thread disabled (run_controller_thread=False); "
+                        "caller must drive controller.run_step synchronously.")
 
     def _send_zero_torque(self) -> None:
         """Send a zero-gain command to make joints passive before shutting down."""
@@ -588,7 +652,15 @@ class UnitreeG1(Robot):
         # Dense whole-body controllers (OpenHLM / pi0.5): expose the 34-D proprio
         # state as ``wb_state.{i}.pos`` so the rollout aggregates it into
         # ``observation.state`` for the policy.
-        if getattr(self.controller, "wb_action", False):
+        if self.config.sonic_token_action:
+            # Token mode: echo the last commanded latent token as observation.state
+            # so a token-output VLA closes the loop on its own previous token.
+            from .controllers.sonic_whole_body import token_state_key
+
+            token = self._last_token if self._last_token is not None else []
+            for i, v in enumerate(token):
+                obs[token_state_key(i)] = float(v)
+        elif getattr(self.controller, "wb_action", False):
             wb_state = obs_to_wb34_state(obs)
             for i, v in enumerate(wb_state):
                 obs[f"wb_state.{i}.pos"] = float(v)
@@ -622,6 +694,12 @@ class UnitreeG1(Robot):
     def send_action(self, action: RobotAction) -> RobotAction:
         action_to_publish = action
         if self.controller is not None:
+            if self.config.sonic_token_action:
+                from .controllers.sonic_whole_body import _extract_token_from_action
+
+                token = _extract_token_from_action(action)
+                if token is not None:
+                    self._last_token = token
             self._update_controller_action(action)
             if self.config.publish_hands and getattr(self.controller, "wb_action", False):
                 self._publish_hand_cmds(action)