diff --git a/src/lerobot/robots/unitree_g1/config_unitree_g1.py b/src/lerobot/robots/unitree_g1/config_unitree_g1.py index fc7a80a9f..298d60e03 100644 --- a/src/lerobot/robots/unitree_g1/config_unitree_g1.py +++ b/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 diff --git a/src/lerobot/robots/unitree_g1/controllers/gr00t_locomotion.py b/src/lerobot/robots/unitree_g1/controllers/gr00t_locomotion.py index 0a33e401b..b25198b42 100644 --- a/src/lerobot/robots/unitree_g1/controllers/gr00t_locomotion.py +++ b/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): diff --git a/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py b/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py index 9443f7b5f..2cda20a6d 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py +++ b/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 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 ef5e958f6..4d44d9dbe 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py +++ b/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 ``.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. diff --git a/src/lerobot/robots/unitree_g1/unitree_g1.py b/src/lerobot/robots/unitree_g1/unitree_g1.py index c7f388c46..45ed147b8 100644 --- a/src/lerobot/robots/unitree_g1/unitree_g1.py +++ b/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: - self.sim_env.step() + 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)