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 566a9fb7c..714425139 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py +++ b/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py @@ -199,36 +199,21 @@ class SonicDecoder: assert off == 994, f"Decoder obs mismatch: {off}" return obs - def step(self, robot_obs, token): - """One control tick: read robot obs, decode the supplied token -> joint targets. + def step(self, lowstate, token): + """One control tick: read robot lowstate, decode the supplied token -> joint targets. Args: - robot_obs: dict with ``.q``/``.dq`` and ``imu.*`` fields. + lowstate: Unitree lowstate with ``motor_state[i].q/.dq`` and ``imu_state`` fields. token: 64-D latent supplied by the policy (encoder bypassed). Returns: dict of ``.q`` target positions (rad) in IsaacLab joint order. """ self.token = np.asarray(token, np.float32) - jnames = [m.name for m in G1_29_JointIndex] - q = np.array( - [ - robot_obs.get(f"{n}.q", self.default_angles[m.value]) - for m, n in zip(G1_29_JointIndex, jnames, strict=False) - ], - np.float32, - ) - dq = np.array([robot_obs.get(f"{n}.dq", 0.0) for n in jnames], np.float32) - quat = np.array( - [ - robot_obs.get("imu.quat.w", 1), - robot_obs.get("imu.quat.x", 0), - robot_obs.get("imu.quat.y", 0), - robot_obs.get("imu.quat.z", 0), - ], - np.float32, - ) - ang = np.array([robot_obs.get(f"imu.gyro.{a}", 0) for a in "xyz"], np.float32) + q = np.array([lowstate.motor_state[m.value].q for m in G1_29_JointIndex], np.float32) + dq = np.array([lowstate.motor_state[m.value].dq for m in G1_29_JointIndex], np.float32) + quat = np.array(lowstate.imu_state.quaternion, np.float32) # (w, x, y, z) + ang = np.array(lowstate.imu_state.gyroscope, np.float32) self.update_history(q, dq, ang, quat) action_mj = ( self.decoder.run(None, {self.decoder_input: self.build_decoder_obs().reshape(1, -1)})[0] @@ -294,7 +279,7 @@ class SonicWholeBodyController: token = self._last_token if self._last_token is not None else np.zeros(TOKEN_DIM, dtype=np.float32) return {token_state_key(i): float(v) for i, v in enumerate(token)} - def _startup_blend(self, obs: dict, out: dict) -> dict: + def _startup_blend(self, lowstate, 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 live commanded target, so the handoff has no snap. @@ -307,8 +292,7 @@ class SonicWholeBodyController: if self._init_step == 0: # Capture the robot's actual pose as the interpolation start point. self._start_pose = { - f"{m.name}.q": float(obs.get(f"{m.name}.q", self._default_angles[m.value])) - for m in G1_29_JointIndex + f"{m.name}.q": float(lowstate.motor_state[m.value].q) for m in G1_29_JointIndex } self._init_step += 1 ratio = min(1.0, self._init_step / self._init_ramp_steps) @@ -320,8 +304,8 @@ class SonicWholeBodyController: logger.info("SONIC startup blend complete -> full policy control") return blended - def run_step(self, action: dict, obs: dict) -> dict: - if not obs: + def run_step(self, action: dict, lowstate) -> dict: + if lowstate is None: return {} # Token-only interface (token-output VLA): a dense 64-D ``motion_token.{i}`` command @@ -335,7 +319,7 @@ class SonicWholeBodyController: self._last_token = self._neutral_token.copy() # 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.controller.step(obs, self._last_token)) + return self._startup_blend(lowstate, self.controller.step(lowstate, self._last_token)) def reset(self): self.controller.reset() diff --git a/src/lerobot/robots/unitree_g1/unitree_g1.py b/src/lerobot/robots/unitree_g1/unitree_g1.py index 2249697a2..f5b948c29 100644 --- a/src/lerobot/robots/unitree_g1/unitree_g1.py +++ b/src/lerobot/robots/unitree_g1/unitree_g1.py @@ -241,12 +241,14 @@ class UnitreeG1(Robot): if self.controller is None: return {f"{G1_29_JointIndex(motor).name}.q": float for motor in G1_29_JointIndex} - # Controllers may define their own action space (e.g. SONIC's 64-D latent token). + # Whole-body controllers (SONIC): 64-D latent token. controller_ft = getattr(self.controller, "action_features", None) if controller_ft is not None: return dict(controller_ft) # Locomotion controllers (GR00T / Holosoma): arm joint targets + joystick axes. + # TODO: have GR00T/Holosoma advertise their own action_features too, so every + # controller declares its action space and this fallthrough can be dropped. arm_features = {f"{G1_29_JointArmIndex(motor).name}.q": float for motor in G1_29_JointArmIndex} remote_features = dict.fromkeys(REMOTE_AXES, float) return {**arm_features, **remote_features} @@ -278,12 +280,8 @@ class UnitreeG1(Robot): with self._controller_action_lock: controller_input = dict(self.controller_input) - # Full-body controllers (SONIC) consume the full observation dict; others - # take the raw lowstate. get_observation() is the single lowstate -> obs builder. - controller_state = ( - self.get_observation() if getattr(self.controller, "full_body", False) else lowstate - ) - controller_action = self.controller.run_step(controller_input, controller_state) + # Run controller step + controller_action = self.controller.run_step(controller_input, lowstate) # Write controller output snapshot with self._controller_action_lock: