From bbfc4ff443901169dbaebab3a45df2d6ed2ba0c9 Mon Sep 17 00:00:00 2001 From: Martino Russi Date: Thu, 30 Jul 2026 11:28:59 +0200 Subject: [PATCH] feat(unitree_g1): make SONIC token interface implicit Drop the ``sonic_token_action`` config flag; the 64-D latent-token action/observation interface now switches on automatically whenever the SONIC whole-body controller is selected (``controller == "SonicWholeBodyController"``). Keyed via a ``_sonic_token`` property so client, onboard and sim roles agree. --- docs/source/unitree_g1.mdx | 5 +- .../robots/unitree_g1/config_unitree_g1.py | 13 ++-- .../controllers/sonic_whole_body.py | 4 +- .../robots/unitree_g1/run_g1_server.py | 1 - src/lerobot/robots/unitree_g1/unitree_g1.py | 67 +++++++++---------- 5 files changed, 39 insertions(+), 51 deletions(-) diff --git a/docs/source/unitree_g1.mdx b/docs/source/unitree_g1.mdx index d6a95ab95..7e2585afc 100644 --- a/docs/source/unitree_g1.mdx +++ b/docs/source/unitree_g1.mdx @@ -320,8 +320,8 @@ cd ~/lerobot python src/lerobot/robots/unitree_g1/run_g1_server.py --handshake --camera ``` -**From your laptop** — run the token policy; `--robot.sonic_token_action=true` switches the -robot to the 64-D latent-token action/observation interface: +**From your laptop** — run the token policy; selecting `--robot.controller=SonicWholeBodyController` +implicitly switches the robot to the 64-D latent-token action/observation interface: ```bash lerobot-rollout \ @@ -331,7 +331,6 @@ lerobot-rollout \ --robot.is_simulation=false \ --robot.robot_ip= \ --robot.controller=SonicWholeBodyController \ - --robot.sonic_token_action=true \ --robot.cameras='{"ego_view": {"type": "zmq", "server_address": "", "port": 5555, "camera_name": "head_camera", "width": 640, "height": 480, "fps": 30}}' \ --task="walk back and forth" \ --duration=1000 \ diff --git a/src/lerobot/robots/unitree_g1/config_unitree_g1.py b/src/lerobot/robots/unitree_g1/config_unitree_g1.py index 917f90dfa..adcc108d2 100644 --- a/src/lerobot/robots/unitree_g1/config_unitree_g1.py +++ b/src/lerobot/robots/unitree_g1/config_unitree_g1.py @@ -86,15 +86,10 @@ class UnitreeG1Config(RobotConfig): # Compensates for gravity on the unitree's arms using the arm ik solver gravity_compensation: bool = False - # 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``) and exposes the last commanded token as a 64-D - # ``observation.state`` (``motion_token_state.{i}.pos``). This lets - # ``lerobot-rollout`` drive a policy 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 - # 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 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 c0f27f703..968456f03 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py +++ b/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py @@ -332,8 +332,8 @@ class SonicWholeBodyController: self._init_step = 0 self._start_pose: dict[str, float] = {} - # 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 + # Token-interface state. ``token_mode`` is set True by the robot whenever a SONIC + # whole-body controller is selected (token-driven deploy): the controller then 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). This lives here (not in the entry-point script) so it applies diff --git a/src/lerobot/robots/unitree_g1/run_g1_server.py b/src/lerobot/robots/unitree_g1/run_g1_server.py index 1f134bcfa..c1b5aeb56 100644 --- a/src/lerobot/robots/unitree_g1/run_g1_server.py +++ b/src/lerobot/robots/unitree_g1/run_g1_server.py @@ -232,7 +232,6 @@ def serve_onboard_controller( dds_interface=dds_interface, release_motion_control=not sim, physical_remote=not sim, - sonic_token_action=sonic_token_action, cameras={}, ) diff --git a/src/lerobot/robots/unitree_g1/unitree_g1.py b/src/lerobot/robots/unitree_g1/unitree_g1.py index 7b3d34df2..718a216e2 100644 --- a/src/lerobot/robots/unitree_g1/unitree_g1.py +++ b/src/lerobot/robots/unitree_g1/unitree_g1.py @@ -179,9 +179,10 @@ class UnitreeG1(Robot): else: self.controller = 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"): + # Token-driven deploy: a SONIC whole-body controller always runs in token + # mode -- it holds a neutral token until the first real one arrives, then + # holds the last token between control ticks. + if hasattr(self.controller, "token_mode"): self.controller.token_mode = True # Controller thread state @@ -199,14 +200,24 @@ class UnitreeG1(Robot): # 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. + # on its own previous token. Implicit whenever the SONIC whole-body controller + # is active. Seeded to zeros; the controller's startup blend eases joints in. self._last_token: np.ndarray | None = None - if config.sonic_token_action: + if self._sonic_token: from .controllers.sonic_whole_body import TOKEN_DIM self._last_token = np.zeros(TOKEN_DIM, dtype=np.float32) + @property + def _sonic_token(self) -> bool: + """Whether the SONIC whole-body decoder is active. + + A SONIC controller consumes a 64-D latent motion token as its action and echoes + the last commanded token as ``observation.state``. Keyed purely off the selected + controller so the token interface is implicit -- no separate config flag. + """ + return self.config.controller == "SonicWholeBodyController" + def _subscribe_lowstate(self): # polls robot state @ 250Hz while not self._shutdown_event.is_set(): start_time = time.time() @@ -295,10 +306,10 @@ class UnitreeG1(Robot): 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). + Exposed only when a SONIC whole-body controller is active; aggregated by the + rollout into a 64-D ``observation.state`` (the last token the policy commanded). """ - if not self.config.sonic_token_action: + if not self._sonic_token: return {} from .controllers.sonic_whole_body import TOKEN_DIM, token_state_key @@ -314,18 +325,18 @@ class UnitreeG1(Robot): @cached_property def action_features(self) -> dict[str, type]: - # Role-agnostic: the schema is a pure function of (controller name, - # sonic_token_action). The thin client advertises the same schema as the - # onboard robot so the exact same policy output routes straight through. + # Role-agnostic: the schema is a pure function of the controller name. The thin + # client advertises the same schema as the onboard robot so the exact same + # policy output routes straight through. # No controller configured at all: raw 29-DoF joint teleop. - if self.config.controller is None and not self.config.sonic_token_action: + if self.config.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: + if self._sonic_token: from .controllers.sonic_whole_body import TOKEN_DIM, token_action_key return {token_action_key(i): float for i in range(TOKEN_DIM)} @@ -465,12 +476,12 @@ class UnitreeG1(Robot): # 1) Handshake: agree with the server on which controller it will run onboard. logger.info( "[client] handshaking with %s:%d (controller=%s, token=%s)...", - server_ip, HANDSHAKE_PORT, self.config.controller, self.config.sonic_token_action, + server_ip, HANDSHAKE_PORT, self.config.controller, self._sonic_token, ) self._client_caps = request_controller( server_ip, self.config.controller, - sonic_token_action=self.config.sonic_token_action, + sonic_token_action=self._sonic_token, port=HANDSHAKE_PORT, ) logger.info("[client] server agreed: %s", self._client_caps) @@ -555,26 +566,10 @@ class UnitreeG1(Robot): # Initialize DDS channel and simulation environment if self.config.is_simulation: - from lerobot.envs.utils import ( - _download_hub_file, - _import_hub_module, - _normalize_hub_result, - ) + from lerobot.envs import make_env self._ChannelFactoryInitialize(0, "lo") - # Call the hub env's make_env directly so we can disable the offscreen - # head_camera renderer. We drive image-conditioned policies from external - # camera frames, 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) + self._env_wrapper = make_env("lerobot/unitree-g1-mujoco", trust_remote_code=True) # Extract the actual gym env from the dict structure self.sim_env = self._env_wrapper["hub_env"][0].envs[0] elif self.config.onboard: @@ -747,7 +742,7 @@ class UnitreeG1(Robot): # Token mode: echo the last commanded latent token as observation.state so a # token-output VLA closes the loop on its own previous token. - if self.config.sonic_token_action: + if self._sonic_token: from .controllers.sonic_whole_body import token_state_key token = self._last_token if self._last_token is not None else [] @@ -769,7 +764,7 @@ class UnitreeG1(Robot): action_to_publish = action if self.controller is not None: - if self.config.sonic_token_action: + if self._sonic_token: from .controllers.sonic_whole_body import _extract_token_from_action token = _extract_token_from_action(action)