From 57ea6f4106528a4ea04861886e28841c3229b537 Mon Sep 17 00:00:00 2001 From: Martino Russi Date: Fri, 24 Jul 2026 12:02:49 +0200 Subject: [PATCH] feat(unitree_g1): episode reset, lazy replay decode, safe shutdown - reset(): pause the background controller and, for full-body controllers, publish the default pose directly (new _controller_paused flag) so reset and the controller loop aren't both writing low commands. - SONIC pipeline: add reset() to StandingEncoderDecoder and PlannerController (clear token/proprio history/heading, rewind motion buffer); SonicRuntime.reset() now calls controller.reset(). - sonic_whole_body: require the full dense 34-D command (no silent zero-fill of a partial action) and integrate yaw-rate (idx 33) into heading. - controllers/__init__: import the controller classes referenced in __all__. - unitree_g1: lazy replay-frame decode + small cache instead of decoding all frames up front; safer disconnect (longer controller-thread join + fail-safe that skips the graceful ramp if the thread won't stop). - lint: ruff-format config_unitree_g1 hand_closed_pose; prettier README table. --- src/lerobot/robots/unitree_g1/README.md | 6 +- .../robots/unitree_g1/config_unitree_g1.py | 4 +- .../robots/unitree_g1/controllers/__init__.py | 6 +- .../unitree_g1/controllers/sonic_pipeline.py | 49 ++++- .../controllers/sonic_whole_body.py | 33 +++- src/lerobot/robots/unitree_g1/unitree_g1.py | 168 ++++++++++++------ 6 files changed, 195 insertions(+), 71 deletions(-) diff --git a/src/lerobot/robots/unitree_g1/README.md b/src/lerobot/robots/unitree_g1/README.md index 6c2799266..45a912d1d 100644 --- a/src/lerobot/robots/unitree_g1/README.md +++ b/src/lerobot/robots/unitree_g1/README.md @@ -12,11 +12,11 @@ targets for the robot's PD controller. Selected with `--robot.controller=`: -| Controller | Purpose | +| Controller | Purpose | | ------------------------------ | ------------------------------------------------------------ | | `SonicWholeBodyController` | SONIC encoder/decoder driven by a 34-D OpenHLM/pi0.5 command | -| `GrootLocomotionController` | GR00T locomotion policy | -| `HolosomaLocomotionController` | Holosoma locomotion policy | +| `GrootLocomotionController` | GR00T locomotion policy | +| `HolosomaLocomotionController` | Holosoma locomotion policy | The rest of this document covers the SONIC whole-body path. diff --git a/src/lerobot/robots/unitree_g1/config_unitree_g1.py b/src/lerobot/robots/unitree_g1/config_unitree_g1.py index 476412930..fc7a80a9f 100644 --- a/src/lerobot/robots/unitree_g1/config_unitree_g1.py +++ b/src/lerobot/robots/unitree_g1/config_unitree_g1.py @@ -80,9 +80,7 @@ class UnitreeG1Config(RobotConfig): publish_hands: bool = False 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] - ) + hand_closed_pose: list[float] = field(default_factory=lambda: [1.0, 0.9, 0.9, 1.3, 1.3, 1.3, 1.3]) hand_kp: float = 1.5 hand_kd: float = 0.1 diff --git a/src/lerobot/robots/unitree_g1/controllers/__init__.py b/src/lerobot/robots/unitree_g1/controllers/__init__.py index 53625c634..22b02194d 100644 --- a/src/lerobot/robots/unitree_g1/controllers/__init__.py +++ b/src/lerobot/robots/unitree_g1/controllers/__init__.py @@ -16,9 +16,13 @@ """Unitree G1 locomotion controllers (Groot, Holosoma, SONIC).""" +from .gr00t_locomotion import GrootLocomotionController +from .holosoma_locomotion import HolosomaLocomotionController +from .sonic_whole_body import SonicRuntime, SonicWholeBodyController + __all__ = [ "GrootLocomotionController", "HolosomaLocomotionController", - "SonicWholeBodyController", "SonicRuntime", + "SonicWholeBodyController", ] diff --git a/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py b/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py index 8f261074f..9443f7b5f 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py +++ b/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py @@ -128,7 +128,15 @@ DEBUG_PRINT_EVERY = 100 # ticks between debug prints def _to_mujoco(a): - """Reorder a 29-vector from IsaacLab order into MuJoCo/deploy order.""" + """Apply the ``MUJOCO_TO_ISAACLAB`` gather to a 29-vector (deploy-order reorder). + + NOTE: this returns ``a[MUJOCO_TO_ISAACLAB]``. The ``_mj`` suffixes and the exact + permutation direction throughout this module are a fixed convention validated + against the deployed SONIC ONNX policy (the encoder/decoder consume vectors in + this order). Do not "correct" the table or rename toward the opposite direction + without re-validating on hardware — the labels are historical, the ordering is + load-bearing. + """ return a[MUJOCO_TO_ISAACLAB] @@ -318,6 +326,23 @@ class StandingEncoderDecoder: self.smpl_root_quat = None self.set_zero_reference() + def reset(self): + """Clear the token, 10-frame proprioception history and heading init. + + ``UnitreeG1.reset()`` relies on this so the first decoder outputs of a new + episode are not contaminated by the previous episode's state. + """ + self.token = np.zeros(TOKEN_DIM, np.float32) + self.last_action_mj = np.zeros(29, np.float32) + self.h_q_mj = [np.zeros(29, np.float32)] * 10 + self.h_dq_mj = [np.zeros(29, np.float32)] * 10 + self.h_ang = [np.zeros(3, np.float32)] * 10 + self.h_act_mj = [np.zeros(29, np.float32)] * 10 + self.h_quat = [np.array([1, 0, 0, 0], np.float32)] * 10 + self.init_base_quat = np.array([1, 0, 0, 0], np.float32) + self.init_ref_quat = np.array([1, 0, 0, 0], np.float32) + self._heading_init = False + def update_history(self, q, dq, ang, quat): """Push the latest proprioception (pos/vel/gyro/orientation) into the 10-frame buffers.""" quat = quat / (np.linalg.norm(quat) + 1e-8) @@ -499,6 +524,28 @@ class PlannerController(StandingEncoderDecoder): self.playing = self.first_motion = False self.motion_lock = threading.Lock() + def reset(self): + """Full reset: clear enc/dec state (super) plus the motion buffer and heading. + + Forces a heading re-init on the next ``step`` so the reference frame is + re-latched to the post-reset robot orientation. + """ + super().reset() + with self.motion_lock: + self.ref_cursor = 0 + self.motion_timesteps = 0 + self.motion_joint_positions[:] = 0.0 + self.motion_joint_velocities[:] = 0.0 + self.motion_body_quats[:] = 0.0 + self.motion_body_quats[:, 0] = 1.0 + self.motion_body_pos[:] = 0.0 + self.init_ref_quat = np.array([1, 0, 0, 0], np.float64) + self.heading_init_base_quat = np.array([1, 0, 0, 0], np.float64) + self.delta_heading = 0.0 + self.first_motion = False + self.playing = False + self.reinit_heading = True + def _heading_apply_delta(self): """Heading correction quaternion (init base-vs-ref heading + operator ``delta_heading``).""" delta = quat_mul( 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 921e00cb9..ef5e958f6 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py +++ b/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py @@ -65,10 +65,15 @@ def _extract_wb34_from_action(action: dict | None) -> np.ndarray | None: units. The ``.pos`` suffix lets these flow through ``lerobot-rollout`` as normal joint-position action features. """ - if not action or wb_action_key(0) not in action: + if not action: + return None + keys = [wb_action_key(i) for i in range(WB_ACTION_DIM)] + # Require the full dense command: a partial action (e.g. only ``wb.0.pos``) + # must not be silently zero-filled, which would drive most joints toward 0. + if any(key not in action for key in keys): return None return np.fromiter( - (float(action.get(wb_action_key(i), 0.0)) for i in range(WB_ACTION_DIM)), + (float(action[key]) for key in keys), dtype=np.float32, count=WB_ACTION_DIM, ) @@ -86,14 +91,15 @@ def _wb34_to_reference(wb: np.ndarray) -> tuple[np.ndarray, np.ndarray]: The 29 joints are first assembled in MuJoCo / Unitree-SDK order ([L-leg 0:6, R-leg 6:12, waist 12:15, L-arm 15:22, R-arm 22:29] — the ``G1_29_JointIndex`` grouping OpenHLM uses), then permuted to IsaacLab order via - ``MUJOCO_TO_ISAACLAB``. Grippers (7, 15) and yaw-rate (33) are not part of the - 29-DoF SONIC reference. + ``MUJOCO_TO_ISAACLAB``. Grippers (7, 15) are not part of the 29-DoF SONIC + reference, and yaw-rate (33) is integrated into the heading by the caller (it + cannot be represented in this static per-tick anchor). """ ref_mj = np.zeros(29, np.float32) # MuJoCo / Unitree-SDK grouped order - ref_mj[0:6] = wb[16:22] # left leg + ref_mj[0:6] = wb[16:22] # left leg ref_mj[6:12] = wb[22:28] # right leg ref_mj[12:15] = wb[28:31] # waist - ref_mj[15:22] = wb[0:7] # left arm + ref_mj[15:22] = wb[0:7] # left arm ref_mj[22:29] = wb[8:15] # right arm ref = ref_mj[MUJOCO_TO_ISAACLAB].astype(np.float32) # -> IsaacLab order for SONIC roll, pitch = float(wb[31]), float(wb[32]) @@ -129,8 +135,10 @@ class SonicRuntime: return self.controller def reset(self): - self.controller.reinit_heading = True - self.controller.playing = True + # Full pipeline reset: clears the encoder token, proprioception history and + # heading, and rewinds the motion buffer. reinit_heading is set so the next + # step re-latches the reference frame to the current robot orientation. + self.controller.reset() def shutdown(self): pass @@ -165,6 +173,9 @@ class SonicWholeBodyController: # stream of per-tick whole-body commands, fed to the encoder as a batch. self._wb_traj: deque[np.ndarray] = deque(maxlen=50) self._wb_quat_traj: deque[np.ndarray] = deque(maxlen=50) + # Integrated heading (rad) from the whole-body command's yaw-rate (index 33), + # forwarded to the pipeline as ``delta_heading`` so turn commands take effect. + self._heading = 0.0 logger.info("SONIC ready (encoder/decoder, 34-D whole-body command path)") @@ -181,6 +192,11 @@ class SonicWholeBodyController: if c.encode_mode != 0: c.encode_mode = 0 c.reinit_heading = True + # Index 33 is a yaw-rate (rad/s): integrate it into a heading offset and hand + # it to the pipeline as ``delta_heading`` so commanded turns are tracked rather + # than silently dropped (the anchor from _wb34_to_reference only carries r/p). + self._heading += float(wb[33]) * CONTROL_DT + c.delta_heading = self._heading # Capture the heading/anchor reference on the first whole-body tick. The # controller only latches ``init_ref_quat`` (and the base heading) inside # ``step()`` when ``first_motion or reinit_heading`` — but it already boots in @@ -282,6 +298,7 @@ class SonicWholeBodyController: self._wb_step = 0 self._wb_traj.clear() self._wb_quat_traj.clear() + self._heading = 0.0 def shutdown(self): self._runtime.shutdown() diff --git a/src/lerobot/robots/unitree_g1/unitree_g1.py b/src/lerobot/robots/unitree_g1/unitree_g1.py index 9c9caf9cb..c7f388c46 100644 --- a/src/lerobot/robots/unitree_g1/unitree_g1.py +++ b/src/lerobot/robots/unitree_g1/unitree_g1.py @@ -161,41 +161,62 @@ class UnitreeG1(Robot): # Controller thread state self._controller_thread = None + # When set, the controller loop stops publishing low commands so reset() can + # drive the joints directly without two publishers fighting (single-publisher). + self._controller_paused = threading.Event() self._controller_action_lock = threading.Lock() self.controller_input = default_remote_input() self.controller_output = {} - # Replay-camera state (decoded frames per robot camera name + play cursor). - self._replay_frames: dict[str, list[np.ndarray]] = {} + # Replay-camera state: keep the encoded (raw) cells per camera and decode + # frames lazily as the play cursor advances, with a small frame cache, so we + # don't materialize gigabytes of decoded RGB at construction time. + self._replay_raw: dict[str, list] = {} + self._replay_cache: dict[tuple[str, int], np.ndarray] = {} + self._replay_cache_cap = 8 self._replay_len = 0 self._replay_idx = 0 if config.replay_camera_parquet and config.replay_camera_map: self._load_replay_frames() def _load_replay_frames(self) -> None: - """Decode recorded episode frames from a parquet into per-camera image lists.""" - import io - + """Load only the mapped parquet columns (encoded frames); decode on demand.""" import pyarrow.parquet as pq - from PIL import Image - table = pq.read_table(self.config.replay_camera_parquet) - cols = {col: table.column(col).to_pylist() for col in self.config.replay_camera_map.values()} + cols_needed = list(dict.fromkeys(self.config.replay_camera_map.values())) + table = pq.read_table(self.config.replay_camera_parquet, columns=cols_needed) self._replay_len = table.num_rows - - def decode(cell) -> np.ndarray: - data = cell["bytes"] if isinstance(cell, dict) else cell - return np.asarray(Image.open(io.BytesIO(data)).convert("RGB"), dtype=np.uint8) - - for cam_name, column in self.config.replay_camera_map.items(): - self._replay_frames[cam_name] = [decode(c) for c in cols[column]] + self._replay_raw = { + cam_name: table.column(column).to_pylist() + for cam_name, column in self.config.replay_camera_map.items() + } logger.info( - "Loaded %d replay frames for cameras %s from %s", + "Loaded %d replay frames (lazy-decode) for cameras %s from %s", self._replay_len, list(self.config.replay_camera_map), self.config.replay_camera_parquet, ) + def _decode_replay_cell(self, cell) -> np.ndarray: + import io + + from PIL import Image + + data = cell["bytes"] if isinstance(cell, dict) else cell + return np.asarray(Image.open(io.BytesIO(data)).convert("RGB"), dtype=np.uint8) + + def _replay_frame(self, cam_name: str, idx: int) -> np.ndarray: + """Decode (and briefly cache) a single replay frame for a camera.""" + key = (cam_name, idx) + cached = self._replay_cache.get(key) + if cached is not None: + return cached + frame = self._decode_replay_cell(self._replay_raw[cam_name][idx]) + if len(self._replay_cache) >= self._replay_cache_cap: + self._replay_cache.pop(next(iter(self._replay_cache))) + self._replay_cache[key] = frame + return frame + def _subscribe_lowstate(self): # polls robot state @ 250Hz while not self._shutdown_event.is_set(): start_time = time.time() @@ -292,8 +313,10 @@ class UnitreeG1(Robot): @property def _replay_cameras_ft(self) -> dict[str, tuple]: - """Replay cameras, shaped from their first decoded frame.""" - return {name: frames[0].shape for name, frames in self._replay_frames.items() if frames} + """Replay cameras, shaped from their first (lazily decoded) frame.""" + if not self._replay_len: + return {} + return {name: self._replay_frame(name, 0).shape for name in self._replay_raw} @cached_property def observation_features(self) -> dict[str, type | tuple]: @@ -333,6 +356,11 @@ class UnitreeG1(Robot): while not self._shutdown_event.is_set(): start_time = time.time() + # Paused during reset() so the reset routine is the sole low-cmd publisher. + if self._controller_paused.is_set(): + time.sleep(control_dt) + continue + with self._lowstate_lock: lowstate = self._lowstate @@ -495,14 +523,23 @@ class UnitreeG1(Robot): def disconnect(self): # Stop the controller loop first so it isn't fighting the shutdown ramp. self._shutdown_event.set() + controller_stopped = True if self._controller_thread is not None: - self._controller_thread.join(timeout=2.0) + # Wait long enough for any in-flight inference tick to finish and the loop + # to observe the shutdown flag, so no stray low command is published while + # the ramp runs (the shutdown routine must be the single publisher). + self._controller_thread.join(timeout=5.0) if self._controller_thread.is_alive(): - logger.warning("Controller thread did not stop cleanly") + controller_stopped = False + logger.error( + "Controller thread did not stop; skipping graceful ramp to avoid " + "concurrent low commands (fail-safe: joints keep last command until exit)" + ) # Soft, damped settle instead of an instant limp (real robot only; the - # subscribe thread is still alive here to supply the current pose). - if not self.config.is_simulation: + # subscribe thread is still alive here to supply the current pose). Only ramp + # once the controller thread has definitely exited. + if not self.config.is_simulation and controller_stopped: self._graceful_stop() if self.controller is not None and hasattr(self.controller, "shutdown"): @@ -569,8 +606,8 @@ class UnitreeG1(Robot): idx = self._replay_idx if idx >= self._replay_len: idx = self._replay_len - 1 if not self.config.replay_camera_loop else idx % self._replay_len - for name, frames in self._replay_frames.items(): - obs[name] = frames[idx] + for name in self._replay_raw: + obs[name] = self._replay_frame(name, idx) self._replay_idx += 1 # Cameras - read images from ZMQ cameras @@ -709,43 +746,64 @@ class UnitreeG1(Robot): if default_positions is None: default_positions = np.array(self.config.default_positions, dtype=np.float32) - if self.config.is_simulation and self.sim_env is not None: - self.sim_env.reset() - self.publish_lowcmd( - {f"{motor.name}.q": float(default_positions[motor.value]) for motor in G1_29_JointIndex} - ) - else: - total_time = 3.0 - num_steps = int(total_time / control_dt) + # Full-body controllers (SONIC / OpenHLM) own the whole 29-DoF command and + # ignore ``.q`` in send_action(), so reset() must publish the default + # pose directly. Pause the background controller first so the two aren't both + # writing low commands while the robot moves to the default pose. + full_body = getattr(self.controller, "full_body", False) + paused = False + if full_body and self._controller_thread is not None: + self._controller_paused.set() + paused = True + time.sleep(control_dt) # let any in-flight controller tick settle - # get current state - obs = self.get_observation() + try: + if self.config.is_simulation and self.sim_env is not None: + self.sim_env.reset() + self.publish_lowcmd( + {f"{motor.name}.q": float(default_positions[motor.value]) for motor in G1_29_JointIndex} + ) + else: + total_time = 3.0 + num_steps = int(total_time / control_dt) - # record current positions - init_dof_pos = np.zeros(29, dtype=np.float32) - for motor in G1_29_JointIndex: - init_dof_pos[motor.value] = obs[f"{motor.name}.q"] + # get current state + obs = self.get_observation() - # Interpolate to default position - for step in range(num_steps): - start_time = time.time() - - alpha = step / num_steps - action_dict = {} + # record current positions + init_dof_pos = np.zeros(29, dtype=np.float32) for motor in G1_29_JointIndex: - target_pos = default_positions[motor.value] - interp_pos = init_dof_pos[motor.value] * (1 - alpha) + target_pos * alpha - action_dict[f"{motor.name}.q"] = float(interp_pos) + init_dof_pos[motor.value] = obs[f"{motor.name}.q"] - self.send_action(action_dict) + # Interpolate to default position + for step in range(num_steps): + start_time = time.time() - # Maintain constant control rate - elapsed = time.time() - start_time - sleep_time = max(0, control_dt - elapsed) - time.sleep(sleep_time) + alpha = step / num_steps + action_dict = {} + for motor in G1_29_JointIndex: + target_pos = default_positions[motor.value] + interp_pos = init_dof_pos[motor.value] * (1 - alpha) + target_pos * alpha + action_dict[f"{motor.name}.q"] = float(interp_pos) - # Reset controller internal state (gait phase, obs history, etc.) - if self.controller is not None and hasattr(self.controller, "reset"): - self.controller.reset() + # Full-body controllers no-op in send_action(); publish the pose + # directly (arm-only controllers keep the send_action() path). + if full_body: + self.publish_lowcmd(action_dict) + else: + self.send_action(action_dict) + + # Maintain constant control rate + elapsed = time.time() - start_time + sleep_time = max(0, control_dt - elapsed) + time.sleep(sleep_time) + + # Reset controller internal state (gait phase, obs history, etc.) before + # resuming so its buffers reflect the post-reset pose. + if self.controller is not None and hasattr(self.controller, "reset"): + self.controller.reset() + finally: + if paused: + self._controller_paused.clear() logger.info("Reached default position")