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/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
 

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",
 ]

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(

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,8 +91,9 @@ 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
@@ -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()

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,6 +746,18 @@ class UnitreeG1(Robot):
         if default_positions is None:
             default_positions = np.array(self.config.default_positions, dtype=np.float32)
 
+        # Full-body controllers (SONIC / OpenHLM) own the whole 29-DoF command and
+        # ignore ``<joint>.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
+
+        try:
             if self.config.is_simulation and self.sim_env is not None:
                 self.sim_env.reset()
                 self.publish_lowcmd(
@@ -737,6 +786,11 @@ class UnitreeG1(Robot):
                         interp_pos = init_dof_pos[motor.value] * (1 - alpha) + target_pos * alpha
                         action_dict[f"{motor.name}.q"] = float(interp_pos)
 
+                    # 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
@@ -744,8 +798,12 @@ class UnitreeG1(Robot):
                     sleep_time = max(0, control_dt - elapsed)
                     time.sleep(sleep_time)
 
-        # Reset controller internal state (gait phase, obs history, etc.)
+            # 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")