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.
Merge branch 'main' into feat/sonic_encoder_decoder
2026-07-30 13:09:40 +00:00 · 2026-07-30 11:30:07 +02:00 · 2026-07-30 11:15:18 +02:00 · 2026-07-30 10:53:27 +02:00 · 2026-07-29 20:24:07 +02:00 · 2026-07-29 20:11:14 +02:00
58 changed files with 2237 additions and 539 deletions
@@ -61,15 +61,19 @@ Full details in [`docs/source/so101.mdx`](./docs/source/so101.mdx) and [`docs/so
 **4.1 Install**

 ```bash
-pip install 'lerobot[feetech]'              # SO-100/SO-101 motor stack
-# pip install 'lerobot[all]'                # everything
-# pip install 'lerobot[aloha,pusht]'        # specific features
-# pip install 'lerobot[smolvla]'            # add SmolVLA deps
-git lfs install && git lfs pull
-hf auth login                               # required to push datasets/policies
-```
+# uv (recommended — see AGENTS.md and CLAUDE.md)
+uv sync --locked --extra feetech          # SO-100/SO-101 motor stack
+# uv sync --locked --extra all            # everything
+# uv sync --locked --extra smolvla        # add SmolVLA deps

-Contributors can alternatively use `uv sync --locked --extra feetech` (see `AGENTS.md`).
+# pip (alternative, e.g. when not working from source)
+# pip install 'lerobot[feetech]'
+# pip install 'lerobot[all]'
+# pip install 'lerobot[smolvla]'
+
+git lfs install && git lfs pull
+hf auth login                             # required to push datasets/policies
+```

 **4.2 Find USB ports** — run once per arm, unplug when prompted.

@@ -68,17 +68,16 @@ ENV HOME=/home/user_lerobot \
 # issues with MuJoCo and OpenGL drivers.
 RUN uv venv --python python${PYTHON_VERSION}

-# Install Python dependencies for caching
+# Install third-party dependencies separately for layer caching
 COPY --chown=user_lerobot:user_lerobot setup.py pyproject.toml uv.lock README.md MANIFEST.in ./
-COPY --chown=user_lerobot:user_lerobot src/ src/
-
-RUN uv sync --locked --extra all --no-cache
+RUN uv sync --locked --extra all --no-install-project --no-cache

 RUN chmod +x /lerobot/.venv/lib/python${PYTHON_VERSION}/site-packages/triton/backends/nvidia/bin/ptxas

-# Copy the rest of the application source code
+# Copy the application source code and install the local project
 # Make sure to have the git-LFS files for testing
 COPY --chown=user_lerobot:user_lerobot . .
+RUN uv sync --locked --extra all --no-cache

 # Set the default command
 CMD ["/bin/bash"]
@@ -60,15 +60,14 @@ ENV HOME=/home/user_lerobot \
 # run other Python projects in the same container without dependency conflicts.
 RUN uv venv

-# Install Python dependencies for caching
+# Install third-party dependencies separately for layer caching
 COPY --chown=user_lerobot:user_lerobot setup.py pyproject.toml uv.lock README.md MANIFEST.in ./
-COPY --chown=user_lerobot:user_lerobot src/ src/
+RUN uv sync --locked --extra all --no-install-project --no-cache

-RUN uv sync --locked --extra all --no-cache
-
-# Copy the rest of the application code
+# Copy the application code and install the local project
 # Make sure to have the git-LFS files for testing
 COPY --chown=user_lerobot:user_lerobot . .
+RUN uv sync --locked --extra all --no-cache

 # Set the default command
 CMD ["/bin/bash"]
@@ -58,7 +58,7 @@ final_action = postprocessor(action)

 ## Hardware API redesign

-PR [#777](https://github.com/huggingface/lerobot/pull/777) improves the LeRobot calibration but is **not backward-compatible**. Below is a overview of what changed and how you can continue to work with datasets created before this pull request.
+PR [#777](https://github.com/huggingface/lerobot/pull/777) improves the LeRobot calibration but is **not backward-compatible**. Below is an overview of what changed and how you can continue to work with datasets created before this pull request.

 ### What changed?

@@ -129,8 +129,8 @@ python examples/backward_compatibility/replay.py \

 Policies output actions in the same format as the datasets (`torch.Tensors`). Therefore, the same transformations should be applied.

-To find these transformations, we recommend to first try and and replay an episode of the dataset your policy was trained on using the section above.
-Then, add these same transformations on your inference script (shown here in the `record.py` script):
+To find these transformations, we recommend first replaying an episode of the dataset your policy was trained on using the section above.
+Then, add these same transformations to your inference script (shown here in the `record.py` script):

 ```diff
 action_values = predict_action(
@@ -88,20 +88,6 @@ policy_preprocessor = NormalizerProcessorStep(stats=dataset_stats)

 The same policy can work with different environment processors, and the same environment processor can work with different policies:

-````python
-# Use SmolVLA policy with LIBERO environment
-# Use SmolVLA policy with LIBERO environment
-libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
-    env_cfg=libero_cfg,
-    policy_cfg=smolvla_cfg,
-)
-smolvla_preprocessor, smolvla_postprocessor = make_pre_post_processors(smolvla_cfg)
-# Or use ACT policy with the same LIBERO environment
-libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
-    env_cfg=libero_cfg,
-    policy_cfg=act_cfg,
-)
-act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg)
 ```python
 # Use SmolVLA policy with LIBERO environment
 libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
@@ -116,6 +102,7 @@ libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
    policy_cfg=act_cfg,
 )
 act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg)
+```

 ### 3. **Easier Experimentation**

@@ -145,7 +132,7 @@ class LiberoVelocityProcessorStep(ObservationProcessorStep):
        state = torch.cat([eef_pos, eef_axisangle, eef_vel,
                          gripper_pos, gripper_vel], dim=-1)  # 14D
        return state
-````
+```

 ### 4. **Cleaner Environment Code**

@@ -40,10 +40,10 @@ This tutorial guides you through updating the firmware of Feetech motors using t
 For each motor you want to update:

 1. **Select the motor** from the list by clicking on it
-2. **Click on Upgrade tab**:
-3. **Click on Online button**:
-   - If an potential firmware update is found, it will be displayed in the box
-4. **Click on Upgrade button**:
+2. **Click the Upgrade tab**:
+3. **Click the Online button**:
+   - If a potential firmware update is found, it will be displayed in the box
+4. **Click the Upgrade button**:
   - The update progress will be displayed

 ## Step 6: Verify Update
@@ -211,7 +211,7 @@ Record, Replay and Train with Hope-JR is still experimental.

 ### Record

-This step records the dataset, which can be seen as an example [here](https://huggingface.co/datasets/nepyope/hand_record_test_with_video_data/settings).
+This step records the dataset, which can be seen as an example [here](https://huggingface.co/datasets/nepyope/hand_record_test_with_video_data).

 ```bash
 lerobot-record \
@@ -1,177 +1,23 @@
-# LeRobot
-
 <div class="flex justify-center">
  <a target="_blank" href="https://huggingface.co/lerobot">
    <img
-      alt="LeRobot, Hugging Face Robotics Library"
+      alt="HuggingFace Expert Acceleration Program"
      src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/lerobot/lerobot-logo-thumbnail.png"
+      style="width: 100%"
    ></img>
  </a>
 </div>

+# LeRobot
+
 **State-of-the-art machine learning for real-world robotics**

-🤗 LeRobot provides a hardware-agnostic, Python-native interface for controlling real robots - from affordable arms like the SO-ARM101 to full humanoids. Plus the tools to record, store, and share the datasets they generate. Every dataset uses the standardized **LeRobotDataset** format (synchronized video + action/state data) and can be streamed directly from the [Hugging Face Hub](https://huggingface.co/lerobot).
+🤗 LeRobot aims to provide models, datasets, and tools for real-world robotics in PyTorch. The goal is to lower the barrier for entry to robotics so that everyone can contribute and benefit from sharing datasets and pretrained models.

-🤗 On top of that data, LeRobot implements state-of-the-art policies - from lightweight imitation-learning models like ACT to large vision-language-action models like π₀ and SmolVLA - all trainable, shareable, and deployable with the same handful of CLI commands.
+🤗 LeRobot contains state-of-the-art approaches that have been shown to transfer to the real-world with a focus on imitation learning and reinforcement learning.

-The goal: lower the barrier to entry for robotics, so that everyone can contribute to, and benefit from, shared datasets and pretrained models.
+🤗 LeRobot already provides a set of pretrained models, datasets with human collected demonstrations, and simulated environments so that everyone can get started.

-<div align="center" style="display: flex; justify-content: center; gap: 8px; flex-wrap: wrap; margin: 20px 0;">
-  <a href="https://discord.gg/s3KuuzsPFb" target="_blank">
-    <img alt="Discord" src="https://img.shields.io/badge/Discord-Join_the_Community-5865F2?style=flat&logo=discord&logoColor=white">
-  </a>
-  <a href="https://x.com/LeRobotHF" target="_blank">
-    <img alt="X (Twitter)" src="https://img.shields.io/badge/X-Follow_%40LeRobotHF-black?style=flat&logo=x&logoColor=white">
-  </a>
-  <a href="https://huggingface.co/lerobot" target="_blank">
-    <img alt="Hugging Face Hub" src="https://img.shields.io/badge/HF_Hub-Models_%26_Datasets-FFD21E?style=flat">
-  </a>
-</div>
+🤗 LeRobot hosts pretrained models and datasets on the LeRobot HuggingFace page.

-<div align="center">
-  <img src="../../media/readme/robots_control_video.webp" width="640px" alt="Reachy 2 Demo">
-</div>
-
-## How It Works
-
-**Teleoperate → Record → Train → Deploy**
-
-1. **Teleoperate** - control the robot yourself (with a leader arm, keyboard, or phone) so it can learn from your movements.
-2. **Record** - each demonstration is saved as a dataset: synchronized camera video plus the actions you took.
-3. **Train** - a policy (the neural network that will control the robot) learns to imitate your demonstrations.
-4. **Deploy** - run the trained policy on the robot and watch it complete the task on its own.
-
-## Get Started
-
-New here? [Install LeRobot](./installation), then pick your path:
-
-<div class="grid grid-cols-1 md:grid-cols-3 gap-4 my-6">
-  <div class="border dark:border-gray-700 rounded-lg p-4 shadow">
-    <div class="text-lg font-semibold mb-2">🔧 I have a robot</div>
-    <p class="text-gray-700 dark:text-gray-300 text-sm">
-      LeRobot supports a wide range of arms and mobile robots. Popular picks:
-    </p>
-    <ul class="text-gray-700 dark:text-gray-300 text-sm list-disc pl-5 mb-2">
-      <li>
-        <a href="./so101">SO-101</a> - our flagship, low-cost arm
-      </li>
-      <li>
-        <a href="./lekiwi">LeKiwi</a> - a mobile base with an arm on top
-      </li>
-      <li>
-        <a href="./koch">Koch v1.1</a> - a long-time community favorite
-      </li>
-      <li>
-        or find yours under <strong>Robots</strong> in the sidebar
-      </li>
-    </ul>
-    <p class="text-gray-700 dark:text-gray-300 text-sm">
-      Once it's assembled and calibrated, record a dataset and train your first
-      policy with the <a href="./il_robots">imitation learning tutorial</a> - or
-      skip the CLI entirely with <a href="./lelab">LeLab</a>, a browser GUI for
-      the same workflow.
-    </p>
-  </div>
-  <div class="border dark:border-gray-700 rounded-lg p-4 shadow">
-    <div class="text-lg font-semibold mb-2">💻 No hardware yet</div>
-    <p class="text-gray-700 dark:text-gray-300 text-sm">
-      You can still train and evaluate policies without owning a robot:
-    </p>
-    <ul class="text-gray-700 dark:text-gray-300 text-sm list-disc pl-5 mb-2">
-      <li>
-        train on an existing
-        <a href="https://huggingface.co/datasets?other=LeRobot">
-          LeRobot dataset
-        </a>
-        from the Hub
-      </li>
-      <li>
-        evaluate in <a href="./envhub">simulation</a>, against benchmarks like
-        LIBERO or Meta-World
-      </li>
-      <li>
-        try the free <a href="./notebooks">Colab notebooks</a> - nothing to
-        install
-      </li>
-    </ul>
-  </div>
-  <div class="border dark:border-gray-700 rounded-lg p-4 shadow">
-    <div class="text-lg font-semibold mb-2">🤝 I want to contribute</div>
-    <p class="text-gray-700 dark:text-gray-300 text-sm">
-      Start with the <a href="./contributing">Contributing guide</a>, then
-      <a href="./bring_your_own_policies">add a new policy</a> or
-      <a href="./integrate_hardware">bring your own hardware</a>.
-    </p>
-  </div>
-</div>
-
-## Explore the Docs
-
-<div class="grid grid-cols-1 md:grid-cols-3 gap-4 my-6">
-  <a
-    class="!no-underline border dark:border-gray-700 rounded-lg p-4 shadow hover:shadow-lg"
-    href="./cheat-sheet"
-  >
-    <div class="font-semibold mb-1">📋 Cheat Sheet</div>
-    <p class="text-gray-700 dark:text-gray-300 text-sm">
-      Every LeRobot CLI command, copy-paste ready.
-    </p>
-  </a>
-  <a
-    class="!no-underline border dark:border-gray-700 rounded-lg p-4 shadow hover:shadow-lg"
-    href="./hardware_guide"
-  >
-    <div class="font-semibold mb-1">🖥️ Compute & Hardware Guide</div>
-    <p class="text-gray-700 dark:text-gray-300 text-sm">
-      Which policy fits your GPU, and how long training takes.
-    </p>
-  </a>
-  <a
-    class="!no-underline border dark:border-gray-700 rounded-lg p-4 shadow hover:shadow-lg"
-    href="./lerobot-dataset-v3"
-  >
-    <div class="font-semibold mb-1">🗂️ LeRobotDataset</div>
-    <p class="text-gray-700 dark:text-gray-300 text-sm">
-      Load, stream, and visualize robot datasets from the Hub.
-    </p>
-  </a>
-  <a
-    class="!no-underline border dark:border-gray-700 rounded-lg p-4 shadow hover:shadow-lg"
-    href="./lelab"
-  >
-    <div class="font-semibold mb-1">🖼 LeLab</div>
-    <p class="text-gray-700 dark:text-gray-300 text-sm">
-      A browser GUI for calibrating, recording, and training - no CLI required.
-    </p>
-  </a>
-  <a
-    class="!no-underline border dark:border-gray-700 rounded-lg p-4 shadow hover:shadow-lg"
-    href="./act"
-  >
-    <div class="font-semibold mb-1">🧠 Policies</div>
-    <p class="text-gray-700 dark:text-gray-300 text-sm">
-      Start with ACT, our recommended first policy - or browse SmolVLA, π₀, and
-      more in the sidebar.
-    </p>
-  </a>
-  <a
-    class="!no-underline border dark:border-gray-700 rounded-lg p-4 shadow hover:shadow-lg"
-    href="./envhub"
-  >
-    <div class="font-semibold mb-1">🎮 Simulation & Benchmarks</div>
-    <p class="text-gray-700 dark:text-gray-300 text-sm">
-      Train and evaluate in simulated environments before touching real
-      hardware.
-    </p>
-  </a>
-</div>
-
-## Common Problems
-
-Running into issues? A few of the most frequent ones:
-
- **Blurry or unusable camera footage** - lighting matters more than resolution. See the [Cameras](./cameras) guide.
- **Build or install errors** (`cmake`, `ffmpeg`, CUDA) - see the Troubleshooting section of the [Installation guide](./installation#troubleshooting).
- **Not sure which policy fits your GPU** - check the [Compute & Hardware Guide](./hardware_guide).
- **Still stuck?** Ask on [Discord](https://discord.gg/s3KuuzsPFb) - the community (and the LeRobot team) is there to help.
+Join the LeRobot community on [Discord](https://discord.gg/s3KuuzsPFb)
@@ -18,7 +18,7 @@ If you're using Feetech or Dynamixel motors, LeRobot provides built-in bus inter
 - [`DynamixelMotorsBus`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/motors/dynamixel/dynamixel.py) – for controlling Dynamixel servos

 Please refer to the [`MotorsBus`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/motors/motors_bus.py) abstract class to learn about its API.
-For a good example of how it can be used, you can have a look at our own [SO101 follower implementation](https://github.com/huggingface/lerobot/blob/main/src/lerobot/robots/so_follower/so101_follower/so101_follower.py)
+For a good example of how it can be used, you can have a look at our own [SO101 follower implementation](https://github.com/huggingface/lerobot/blob/main/src/lerobot/robots/so_follower/so_follower.py)

 Use these if compatible. Otherwise, you'll need to find or write a Python interface (not covered in this tutorial):

@@ -51,7 +51,7 @@ In addition to these instructions, you need to install the Feetech SDK & ZeroMQ
 pip install -e ".[lekiwi]"
 ```

-Great :hugs:! You are now done installing LeRobot, and we can begin assembling the SO100/SO101 arms and the mobile base :robot:.
+Great 🤗! You are now done installing LeRobot, and we can begin assembling the SO100/SO101 arms and the mobile base 🤖.
 Every time you now want to use LeRobot, you can go to the `~/lerobot` folder where we installed LeRobot and run one of the commands.

 # Step-by-Step Assembly Instructions
@@ -174,7 +174,7 @@ The model takes images, text instructions, and robot state as input, and outputs

 ## Reproducing π₀Fast results

-We reproduce the results of π₀Fast on the LIBERO benchmark using the LeRobot implementation. We take the LeRobot PiFast base model [lerobot/pi0fast-base](https://huggingface.co/lerobot/pi0fast-base) and finetune for an additional 40kk steps in bfloat16, with batch size of 256 on 8 H100 GPUs using the [HuggingFace LIBERO dataset](https://huggingface.co/datasets/HuggingFaceVLA/libero).
+We reproduce the results of π₀Fast on the LIBERO benchmark using the LeRobot implementation. We take the LeRobot PiFast base model [lerobot/pi0fast-base](https://huggingface.co/lerobot/pi0fast-base) and finetune for an additional 40k steps in bfloat16, with batch size of 256 on 8 H100 GPUs using the [HuggingFace LIBERO dataset](https://huggingface.co/datasets/HuggingFaceVLA/libero).

 The finetuned model can be found here:

@@ -22,7 +22,7 @@ With processors, you choose the learning features you want to use for your polic
 ## Three pipelines

 We often compose three pipelines. Depending on your setup, some can be empty if action and observation spaces already match.
-Each of these pipelines handle different conversions between different action and observation spaces. Below is a quick explanation of each pipeline.
+Each of these pipelines handles different conversions between different action and observation spaces. Below is a quick explanation of each pipeline.

 1. Pipeline 1: Teleop action space → dataset action space (phone pose → EE targets)
 2. Pipeline 2: Dataset action space → robot command space (EE targets → joints)
@@ -74,15 +74,15 @@ In the phone to SO-100 follower examples we use the following adapters:
 - `robot_action_to_transition`: transforms the teleop action dict to a pipeline transition.
 - `transition_to_robot_action`: transforms the pipeline transition to a robot action dict.
 - `observation_to_transition`: transforms the robot observation dict to a pipeline transition.
- `transition_to_observation`: transforms the pipeline transition to a observation dict.
+- `transition_to_observation`: transforms the pipeline transition to an observation dict.

-Checkout [src/lerobot/processor/converters.py](https://github.com/huggingface/lerobot/blob/main/src/lerobot/processor/converters.py) for more details.
+Check out [src/lerobot/processor/converters.py](https://github.com/huggingface/lerobot/blob/main/src/lerobot/processor/converters.py) for more details.

 ## Dataset feature contracts

 Dataset features are determined by the keys saved in the dataset. Each step can declare what features it modifies in a contract called `transform_features(...)`. Once you build a processor, the processor can then aggregate all of these features with `aggregate_pipeline_dataset_features()` and merge multiple feature dicts with `combine_feature_dicts(...)`.

-Below is and example of how we declare features with the `transform_features` method in the phone to SO-100 follower examples:
+Below is an example of how we declare features with the `transform_features` method in the phone to SO-100 follower examples:

 ```python
    def transform_features(
@@ -57,7 +57,7 @@ policy_cfg.rtc_config = RTCConfig(
 policy = PI0Policy.from_pretrained("lerobot/pi0_base", policy_cfg=policy_cfg, device="cuda")

 # Now use predict_action_chunk with RTC parameters
-inference_delay = 4  # How many steps of inference latency, this values should be calculated based on the inference latency of the policy
+inference_delay = 4  # How many steps of inference latency, this value should be calculated based on the inference latency of the policy

 # Initialize the action queue
 action_queue = ActionQueue(policy_cfg.rtc_config)
@@ -100,7 +100,7 @@ Typical values: 8-12 steps
 RTCConfig(execution_horizon=10)
 ```

-**`max_guidance_weight`**: How strongly to enforce consistency with the previous chunk. This is a hyperparameter that can be tuned to balance the smoothness of the transitions and the reactivity of the policy. For 10 steps flow matching (SmolVLA, Pi0, Pi0.5), a value of 10.0 is a optimal value.
+**`max_guidance_weight`**: How strongly to enforce consistency with the previous chunk. This is a hyperparameter that can be tuned to balance the smoothness of the transitions and the reactivity of the policy. For 10 steps flow matching (SmolVLA, Pi0, Pi0.5), a value of 10.0 is an optimal value.

 **`prefix_attention_schedule`**: How to weight consistency across the overlap region.

@@ -93,7 +93,7 @@ lerobot-train --help

 ## Evaluate the finetuned model and run it in real-time

-Similarly for when recording an episode, it is recommended that you are logged in to the HuggingFace Hub. You can follow the corresponding steps: [Record a dataset](./il_robots).
+Similarly for when recording an episode, it is recommended that you are logged in to the HuggingFace Hub. You can follow the corresponding steps: [Record a dataset](./il_robots#record-a-dataset).
 Once you are logged in, you can run inference in your setup by doing:

 ```bash
@@ -8,6 +8,15 @@

 The Unitree G1 humanoid is now supported in LeRobot! You can teleoperate, train locomanipulation policies, test in sim, and more. Both 29 and 23 DoF variants are supported.

+<Tip>
+  **New: SONIC whole-body control.** The `SonicWholeBodyController` runs NVIDIA's
+  [GEAR-SONIC](https://huggingface.co/nvidia/GEAR-SONIC) decoder on the G1, turning a
+  64-D latent motion token into full-body joint targets at 50 Hz. This lets you drive
+  the robot from a VLA policy trained on SONIC motion tokens (token in → whole-body
+  motion out) with `lerobot-rollout`, in sim or on the physical robot. See
+  [Whole-body control with SONIC](#whole-body-control-with-sonic) below.
+</Tip>
+
 ---

 ## Part 1: Getting Started
@@ -59,7 +68,7 @@ lerobot-teleoperate \
  --robot.controller=GrootLocomotionController
 ```

-This will launch a [MuJoCo sim instance](https://huggingface.co/lerobot/unitree-g1-mujoco/tree/main) for the G1. You can connect a gamepad to your machine before launching in order to control the robot's locomotion in sim. We support both [HolosomaLocomotionController](https://github.com/amazon-far/holosoma) and [GrootLocomotionController](https://github.com/NVlabs/GR00T-WholeBodyControl) via `--robot.controller`.
+This will launch a [MuJoCo sim instance](https://huggingface.co/lerobot/unitree-g1-mujoco/tree/main) for the G1. You can connect a gamepad to your machine before launching in order to control the robot's locomotion in sim. We support [HolosomaLocomotionController](https://github.com/amazon-far/holosoma), [GrootLocomotionController](https://github.com/NVlabs/GR00T-WholeBodyControl), and [SonicWholeBodyController](https://huggingface.co/nvidia/GEAR-SONIC) via `--robot.controller`.

 - Press `9` to release the robot
 - Press `7` / `8` to increase / decrease waist height
@@ -290,6 +299,52 @@ lerobot-rollout \

 ---

+## Whole-body control with SONIC
+
+The `SonicWholeBodyController` runs NVIDIA's [GEAR-SONIC](https://huggingface.co/nvidia/GEAR-SONIC)
+decoder on the G1. Each 50 Hz tick it consumes a **64-D latent motion token** and emits
+full-body joint targets — the encoder is bypassed, so a policy feeds tokens in and the
+decoder turns them into motion. Before the first token arrives the controller holds a
+neutral (idle) pose.
+
+This makes the G1 drivable by a VLA policy trained to output SONIC motion tokens (token
+as both `observation.state` and `action`, e.g. [`nepyope/sonic_walk`](https://huggingface.co/nepyope/sonic_walk))
+using the standard `lerobot-rollout`. The controller always runs **onboard** the robot;
+the laptop is a thin client that streams tokens and receives camera frames over ZMQ.
+
+**On the robot** — start the server in handshake mode so it instantiates and runs the
+controller onboard against local DDS at full rate:
+
+```bash
+cd ~/lerobot
+python src/lerobot/robots/unitree_g1/run_g1_server.py --handshake --camera
+```
+
+**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 \
+  --policy.path=nepyope/sonic_walk \
+  --policy.device=cuda \
+  --robot.type=unitree_g1 \
+  --robot.is_simulation=false \
+  --robot.robot_ip=<ROBOT_IP> \
+  --robot.controller=SonicWholeBodyController \
+  --robot.cameras='{"ego_view": {"type": "zmq", "server_address": "<ROBOT_IP>", "port": 5555, "camera_name": "head_camera", "width": 640, "height": 480, "fps": 30}}' \
+  --task="walk back and forth" \
+  --duration=1000 \
+  --fps=30
+```
+
+<Tip>
+  SONIC is a token-only decoder in LeRobot: the only input path is the 64-D latent
+  vector. To train your own token policy, expose the 64-D token as the action (a config
+  choice, e.g. `pi05` with a 64-D action dim) — no policy code changes are needed.
+</Tip>
+
+---
+
 ## Additional Resources

 - [Unitree SDK Documentation](https://github.com/unitreerobotics/unitree_sdk2_python)
@@ -50,11 +50,11 @@ lerobot-edit-dataset \
 Divide a dataset into multiple subsets.

 ```bash
-# Split by fractions (e.g. 80% train, 20% test, 20% val)
+# Split by fractions (e.g. 60% train, 20% val, 20% test)
 lerobot-edit-dataset \
    --repo_id lerobot/pusht \
    --operation.type split \
-    --operation.splits '{"train": 0.8, "test": 0.2, "val": 0.2}'
+    --operation.splits '{"train": 0.6, "val": 0.2, "test": 0.2}'

 # Split by specific episode indices
 lerobot-edit-dataset \
@@ -494,6 +494,19 @@ ignore_errors = true
 module = "lerobot.envs.*"
 ignore_errors = false

+[[tool.mypy.overrides]]
+module = "lerobot.annotations.*"
+ignore_errors = false
+disallow_untyped_defs = true
+disallow_incomplete_defs = true
+check_untyped_defs = true
+
+[[tool.mypy.overrides]]
+module = "lerobot.transforms.*"
+ignore_errors = false
+disallow_untyped_defs = true
+disallow_incomplete_defs = true
+check_untyped_defs = true

 # [[tool.mypy.overrides]]
 # module = "lerobot.utils.*"
@@ -19,6 +19,7 @@ import copy
 import logging
 import shutil
 from pathlib import Path
+from typing import Any, NotRequired, TypedDict

 import datasets
 import pandas as pd
@@ -49,8 +50,32 @@ from .utils import (
 )
 from .video_utils import concatenate_video_files, get_video_duration_in_s

+logger = logging.getLogger(__name__)

-def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMetadata]) -> dict[str, dict]:
+type FeatureDict = dict[str, dict[str, Any]]
+type ChunkFile = tuple[int, int]
+
+
+class IndexState(TypedDict):
+    chunk: int
+    file: int
+    src_to_dst: NotRequired[dict[ChunkFile, ChunkFile]]
+
+
+class VideoIndex(TypedDict):
+    chunk: int
+    file: int
+    latest_duration: float
+    episode_duration: float
+    src_to_offset: NotRequired[dict[ChunkFile, float]]
+    src_to_dst: NotRequired[dict[ChunkFile, ChunkFile]]
+    dst_file_durations: NotRequired[dict[ChunkFile, float]]
+
+
+type VideoIndexState = dict[str, VideoIndex]
+
+
+def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMetadata]) -> FeatureDict:
    """Create a merged video feature info dictionary for aggregation. The video encoder info is merged field-by-field: each key is kept only when every source agrees; otherwise that key is set to ``null`` (or ``{}`` for ``video.extra_options``) and a warning is logged.

    Args:
@@ -59,14 +84,14 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
    Returns:
        dict: A dictionary of merged video feature info.
    """
-    merged_info = copy.deepcopy(all_metadata[0].features)
+    merged_info: FeatureDict = copy.deepcopy(all_metadata[0].features)
    video_keys = [k for k in merged_info if merged_info[k].get("dtype") == "video"]

    for vk in video_keys:
        video_infos = [m.features.get(vk, {}).get("info") or {} for m in all_metadata]
        base_video_info = video_infos[0]

-        merged_encoder_info: dict = {}
+        merged_encoder_info: dict[str, Any] = {}
        fallback_keys: list[str] = []
        for info_key in VIDEO_ENCODER_INFO_KEYS:
            values = [info.get(info_key, None) for info in video_infos]
@@ -80,7 +105,7 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
                merged_encoder_info[info_key] = {} if info_key == "video.extra_options" else None

        if fallback_keys:
-            logging.warning(
+            logger.warning(
                f"Merging heterogeneous or incomplete video encoder metadata for feature {vk}. "
                f"Setting these keys to null: {fallback_keys}.",
            )
@@ -92,7 +117,7 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
    return merged_info


-def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]):
+def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]) -> tuple[int, str | None, FeatureDict]:
    """Validates that all dataset metadata have consistent properties.

    Ensures all datasets have the same fps, robot_type, and features to guarantee
@@ -129,7 +154,9 @@ def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]):
    return fps, robot_type, features


-def update_data_df(df, src_meta, dst_meta):
+def update_data_df(
+    df: pd.DataFrame, src_meta: LeRobotDatasetMetadata, dst_meta: LeRobotDatasetMetadata
+) -> pd.DataFrame:
    """Updates a data DataFrame with new indices and task mappings for aggregation.

    Adjusts episode indices, frame indices, and task indices to account for
@@ -154,12 +181,12 @@ def update_data_df(df, src_meta, dst_meta):


 def update_meta_data(
-    df,
-    dst_meta,
-    meta_idx,
-    data_idx,
-    videos_idx,
-):
+    df: pd.DataFrame,
+    dst_meta: LeRobotDatasetMetadata,
+    meta_idx: IndexState,
+    data_idx: IndexState,
+    videos_idx: VideoIndexState,
+) -> pd.DataFrame:
    """Updates metadata DataFrame with new chunk, file, and timestamp indices.

    Adjusts all indices and timestamps to account for previously aggregated
@@ -289,7 +316,7 @@ def aggregate_datasets(
    chunk_size: int | None = None,
    concatenate_videos: bool = True,
    concatenate_data: bool = True,
-):
+) -> None:
    """Aggregates multiple LeRobot datasets into a single unified dataset.

    This is the main function that orchestrates the aggregation process by:
@@ -309,7 +336,7 @@ def aggregate_datasets(
        concatenate_videos: When False, keep one mp4 per source file instead of packing into shards.
        concatenate_data: When False, keep one parquet per source file instead of packing into shards.
    """
-    logging.info("Start aggregate_datasets")
+    logger.info("Start aggregate_datasets")

    if data_files_size_in_mb is None:
        data_files_size_in_mb = DEFAULT_DATA_FILE_SIZE_IN_MB
@@ -341,15 +368,15 @@ def aggregate_datasets(
        video_files_size_in_mb=video_files_size_in_mb,
    )

-    logging.info("Find all tasks")
+    logger.info("Find all tasks")
    unique_tasks = pd.concat([m.tasks for m in all_metadata]).index.unique()
    dst_meta.tasks = pd.DataFrame(
        {"task_index": range(len(unique_tasks))}, index=pd.Index(unique_tasks, name="task")
    )

-    meta_idx = {"chunk": 0, "file": 0}
-    data_idx = {"chunk": 0, "file": 0}
-    videos_idx = {
+    meta_idx: IndexState = {"chunk": 0, "file": 0}
+    data_idx: IndexState = {"chunk": 0, "file": 0}
+    videos_idx: VideoIndexState = {
        key: {"chunk": 0, "file": 0, "latest_duration": 0, "episode_duration": 0} for key in video_keys
    }

@@ -373,12 +400,17 @@ def aggregate_datasets(
        dst_meta.info.total_frames += src_meta.total_frames

    finalize_aggregation(dst_meta, all_metadata)
-    logging.info("Aggregation complete.")
+    logger.info("Aggregation complete.")


 def aggregate_videos(
-    src_meta, dst_meta, videos_idx, video_files_size_in_mb, chunk_size, concatenate_videos=True
-):
+    src_meta: LeRobotDatasetMetadata,
+    dst_meta: LeRobotDatasetMetadata,
+    videos_idx: VideoIndexState,
+    video_files_size_in_mb: float,
+    chunk_size: int,
+    concatenate_videos: bool = True,
+) -> VideoIndexState:
    """Aggregates video chunks from a source dataset into the destination dataset.

    Handles video file concatenation and rotation based on file size limits.
@@ -406,15 +438,16 @@ def aggregate_videos(
            videos_idx[key]["dst_file_durations"] = {}

    for key, video_idx in videos_idx.items():
-        unique_chunk_file_pairs = {
-            (chunk, file)
-            for chunk, file in zip(
-                src_meta.episodes[f"videos/{key}/chunk_index"],
-                src_meta.episodes[f"videos/{key}/file_index"],
-                strict=False,
-            )
-        }
-        unique_chunk_file_pairs = sorted(unique_chunk_file_pairs)
+        unique_chunk_file_pairs: list[ChunkFile] = sorted(
+            {
+                (chunk, file)
+                for chunk, file in zip(
+                    src_meta.episodes[f"videos/{key}/chunk_index"],
+                    src_meta.episodes[f"videos/{key}/file_index"],
+                    strict=False,
+                )
+            }
+        )

        chunk_idx = video_idx["chunk"]
        file_idx = video_idx["file"]
@@ -489,7 +522,14 @@ def aggregate_videos(
    return videos_idx


-def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_size, concatenate_data=True):
+def aggregate_data(
+    src_meta: LeRobotDatasetMetadata,
+    dst_meta: LeRobotDatasetMetadata,
+    data_idx: IndexState,
+    data_files_size_in_mb: float,
+    chunk_size: int,
+    concatenate_data: bool = True,
+) -> IndexState:
    """Aggregates data chunks from a source dataset into the destination dataset.

    Reads source data files, updates indices to match the aggregated dataset,
@@ -510,14 +550,16 @@ def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_si
    Returns:
        dict: Updated data_idx with current chunk and file indices.
    """
-    unique_chunk_file_ids = {
-        (c, f)
-        for c, f in zip(
-            src_meta.episodes["data/chunk_index"], src_meta.episodes["data/file_index"], strict=False
-        )
-    }
-
-    unique_chunk_file_ids = sorted(unique_chunk_file_ids)
+    unique_chunk_file_ids: list[ChunkFile] = sorted(
+        {
+            (c, f)
+            for c, f in zip(
+                src_meta.episodes["data/chunk_index"],
+                src_meta.episodes["data/file_index"],
+                strict=False,
+            )
+        }
+    )
    contains_images = len(dst_meta.image_keys) > 0

    # retrieve features schema for proper image typing in parquet
@@ -525,7 +567,7 @@ def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_si

    # Track source to destination file mapping for metadata update
    # This is critical for handling datasets that are already results of a merge
-    src_to_dst: dict[tuple[int, int], tuple[int, int]] = {}
+    src_to_dst: dict[ChunkFile, ChunkFile] = {}

    for src_chunk_idx, src_file_idx in unique_chunk_file_ids:
        src_path = src_meta.root / DEFAULT_DATA_PATH.format(
@@ -564,7 +606,13 @@ def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_si
    return data_idx


-def aggregate_metadata(src_meta, dst_meta, meta_idx, data_idx, videos_idx):
+def aggregate_metadata(
+    src_meta: LeRobotDatasetMetadata,
+    dst_meta: LeRobotDatasetMetadata,
+    meta_idx: IndexState,
+    data_idx: IndexState,
+    videos_idx: VideoIndexState,
+) -> IndexState:
    """Aggregates metadata from a source dataset into the destination dataset.

    Reads source metadata files, updates all indices and timestamps,
@@ -580,16 +628,16 @@ def aggregate_metadata(src_meta, dst_meta, meta_idx, data_idx, videos_idx):
    Returns:
        dict: Updated meta_idx with current chunk and file indices.
    """
-    chunk_file_ids = {
-        (c, f)
-        for c, f in zip(
-            src_meta.episodes["meta/episodes/chunk_index"],
-            src_meta.episodes["meta/episodes/file_index"],
-            strict=False,
-        )
-    }
-
-    chunk_file_ids = sorted(chunk_file_ids)
+    chunk_file_ids: list[ChunkFile] = sorted(
+        {
+            (c, f)
+            for c, f in zip(
+                src_meta.episodes["meta/episodes/chunk_index"],
+                src_meta.episodes["meta/episodes/file_index"],
+                strict=False,
+            )
+        }
+    )
    for chunk_idx, file_idx in chunk_file_ids:
        src_path = src_meta.root / DEFAULT_EPISODES_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
        df = pd.read_parquet(src_path)
@@ -622,16 +670,16 @@ def aggregate_metadata(src_meta, dst_meta, meta_idx, data_idx, videos_idx):
 def append_or_create_parquet_file(
    df: pd.DataFrame,
    src_path: Path,
-    idx: dict[str, int],
+    idx: IndexState,
    max_mb: float,
    chunk_size: int,
    default_path: str,
    contains_images: bool = False,
-    aggr_root: Path = None,
+    aggr_root: Path | None = None,
    hf_features: datasets.Features | None = None,
    concatenate: bool = True,
    one_row_group_per_episode: bool = False,
-) -> tuple[dict[str, int], tuple[int, int]]:
+) -> tuple[IndexState, ChunkFile]:
    """Appends data to an existing parquet file or creates a new one based on size constraints.

    Manages file rotation when size limits are exceeded to prevent individual files
@@ -654,7 +702,13 @@ def append_or_create_parquet_file(
    Returns:
        tuple: (updated_idx, (dst_chunk, dst_file)) where updated_idx is the index dict
               and (dst_chunk, dst_file) is the actual destination file the data was written to.
+
+    Raises:
+        ValueError: If aggr_root is not provided.
    """
+    if aggr_root is None:
+        raise ValueError("aggr_root must be provided.")
+
    dst_chunk, dst_file = idx["chunk"], idx["file"]
    dst_path = aggr_root / default_path.format(chunk_index=dst_chunk, file_index=dst_file)

@@ -698,7 +752,9 @@ def append_or_create_parquet_file(
    return idx, (dst_chunk, dst_file)


-def finalize_aggregation(aggr_meta, all_metadata):
+def finalize_aggregation(
+    aggr_meta: LeRobotDatasetMetadata, all_metadata: list[LeRobotDatasetMetadata]
+) -> None:
    """Finalizes the dataset aggregation by writing summary files and statistics.

    Writes the tasks file, info file with total counts and splits, and
@@ -708,16 +764,16 @@ def finalize_aggregation(aggr_meta, all_metadata):
        aggr_meta: Aggregated dataset metadata.
        all_metadata: List of all source dataset metadata objects.
    """
-    logging.info("write tasks")
+    logger.info("write tasks")
    write_tasks(aggr_meta.tasks, aggr_meta.root)

-    logging.info("write info")
+    logger.info("write info")
    aggr_meta.info.total_tasks = len(aggr_meta.tasks)
    aggr_meta.info.total_episodes = sum(m.total_episodes for m in all_metadata)
    aggr_meta.info.total_frames = sum(m.total_frames for m in all_metadata)
    aggr_meta.info.splits = {"train": f"0:{sum(m.total_episodes for m in all_metadata)}"}
    write_info(aggr_meta.info, aggr_meta.root)

-    logging.info("write stats")
+    logger.info("write stats")
    aggr_meta.stats = aggregate_stats([m.stats for m in all_metadata])
    write_stats(aggr_meta.stats, aggr_meta.root)
@@ -188,8 +188,8 @@ class LeRobotDatasetMetadata:
    def _load_metadata(self):
        self.info = load_info(self.root)
        check_version_compatibility(self.repo_id, self._version, CODEBASE_VERSION)
-        self.tasks = load_tasks(self.root)
-        self.episodes = load_episodes(self.root)
+        self.tasks = load_tasks(self.root) if self.total_tasks > 0 else None
+        self.episodes = load_episodes(self.root) if self.total_episodes > 0 else None
        self.stats = load_stats(self.root)

    def ensure_readable(self) -> None:
@@ -384,7 +384,12 @@ class LiberoEnv(gym.Env):

    def close(self):
        if self._env is not None:
-            self._env.close()
+            try:
+                self._env.close()
+            finally:
+                # LIBERO deletes its inner env on close, so this wrapper must
+                # be recreated before the next reset.
+                self._env = None


 def _make_env_fns(
@@ -384,7 +384,9 @@ class RoboTwinEnv(gym.Env):

        self._env: Any | None = None  # deferred — created on first reset() inside worker
        self._step_count: int = 0
-        self._black_frame = np.zeros((self.observation_height, self.observation_width, 3), dtype=np.uint8)
+        self._black_frame: np.ndarray = np.zeros(
+            (self.observation_height, self.observation_width, 3), dtype=np.uint8
+        )

        image_spaces = {
            cam: spaces.Box(
@@ -373,7 +373,7 @@ class VLABenchEnv(gym.Env):

        if action.shape[0] != 7:
            # Unknown layout — fall back to zero-pad so the sim doesn't crash.
-            padded = np.zeros(ctrl_dim, dtype=np.float64)
+            padded: np.ndarray = np.zeros(ctrl_dim, dtype=np.float64)
            padded[: min(action.shape[0], ctrl_dim)] = action[:ctrl_dim]
            return padded

@@ -302,6 +302,33 @@ def _pad_evo1_stats(
    return padded_stats


+def _refresh_evo1_normalization_steps(
+    config: Evo1Config,
+    preprocessor: PolicyProcessorPipeline,
+    postprocessor: PolicyProcessorPipeline,
+) -> None:
+    """Re-pad checkpoint-loaded (un)normalizer stats/features to EVO1's fixed widths.
+
+    Loading a checkpoint injects the raw dataset stats (unpadded to max_state_dim/max_action_dim)
+    into the (un)normalizer via the generic override path in make_pre_post_processors. Those stats
+    and their declared features must be re-padded/reshaped to EVO1's fixed widths, otherwise
+    normalization fails against the padded state/action tensors (e.g. state padded to 24 vs. 8-dim
+    LIBERO stats). Padding is a no-op when stats are already at the target width.
+    """
+    normalization_features = _evo1_normalization_features(config)
+    action_features = _evo1_action_features(config)
+    for step in preprocessor.steps:
+        if isinstance(step, NormalizerProcessorStep):
+            step.features = normalization_features
+            step.stats = _pad_evo1_stats(config, step.stats)
+            step.to(device=step.device, dtype=step.dtype)
+    for step in postprocessor.steps:
+        if isinstance(step, UnnormalizerProcessorStep):
+            step.features = action_features
+            step.stats = _pad_evo1_stats(config, step.stats)
+            step.to(device=step.device, dtype=step.dtype)
+
+
 def reconcile_evo1_processors(
    config: Evo1Config,
    preprocessor: PolicyProcessorPipeline,
@@ -309,16 +336,19 @@ def reconcile_evo1_processors(
 ) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]:
    """Reconcile checkpoint-loaded pipelines with the current EVO1 config.

-    Two things cannot be restored from a serialized pipeline alone: the EVO1 batch converter
-    (converters are plain functions and are never serialized), and eval-time CLI overrides of the
-    action postprocessing flags (`postprocess_action_dim`, `binarize_gripper`, `gripper_*`). This
-    restores the converter and rebuilds the action step from the current config so those overrides
-    take effect.
+    Three things cannot be restored from a serialized pipeline alone: the EVO1 batch converter
+    (converters are plain functions and are never serialized), eval-time CLI overrides of the
+    action postprocessing flags (`postprocess_action_dim`, `binarize_gripper`, `gripper_*`), and the
+    (un)normalizer stats/features when the generic override path injects raw, unpadded dataset
+    stats. This restores the converter, re-pads the normalization stats to EVO1's fixed widths, and
+    rebuilds the action step from the current config so those overrides take effect.
    """
    # Pipelines reloaded from a checkpoint come back with the default batch converter, which drops
    # non-observation extras (embodiment_id, state_mask, custom task fields) needed by EVO1.
    preprocessor.to_transition = evo1_batch_to_transition

+    _refresh_evo1_normalization_steps(config, preprocessor, postprocessor)
+
    action_step = Evo1ActionProcessorStep(
        action_dim=_evo1_action_dim(config),
        binarize_gripper=config.binarize_gripper,
@@ -18,7 +18,7 @@ import functools
 import threading
 from collections.abc import Callable, Sequence
 from contextlib import suppress
-from typing import TypedDict
+from typing import NotRequired, TypedDict

 import torch
 import torch.nn.functional as F  # noqa: N812
@@ -36,7 +36,7 @@ class BatchTransition(TypedDict):
    next_state: dict[str, torch.Tensor]
    done: torch.Tensor
    truncated: torch.Tensor
-    complementary_info: dict[str, torch.Tensor | float | int] | None = None
+    complementary_info: NotRequired[dict[str, torch.Tensor | float | int] | None]


 def random_crop_vectorized(images: torch.Tensor, output_size: tuple) -> torch.Tensor:
@@ -46,6 +46,12 @@ class SOFollowerConfig:
    position_i_coefficient: int = 0
    position_d_coefficient: int = 32

+    # Number of extra attempts when a `sync_read` of the motors fails. Feetech buses can occasionally
+    # return a corrupted status packet ("Incorrect status packet!"), especially when several joints move
+    # at once, which otherwise aborts the control loop. Retries are immediate (no sleep) and only happen on
+    # failure, so the steady-state read cost is unchanged.
+    num_read_retries: int = 2
+

@RobotConfig.register_subclass("so101_follower")
@RobotConfig.register_subclass("so100_follower")
@@ -510,10 +510,10 @@ class ForwardKinematicsJointsToEEAction(RobotActionProcessorStep):
        # We only use the ee pose in the dataset, so we don't need the joint positions
        for n in self.motor_names:
            features[PipelineFeatureType.ACTION].pop(f"{n}.pos", None)
-        # We specify the dataset features of this step that we want to be stored in the dataset
+        # Store end-effector features as actions in the dataset schema
        for k in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]:
            features[PipelineFeatureType.ACTION][f"ee.{k}"] = PolicyFeature(
-                type=FeatureType.STATE, shape=(1,)
+                type=FeatureType.ACTION, shape=(1,)
            )
        return features

@@ -180,7 +180,7 @@ class SOFollower(Robot):
    def get_observation(self) -> RobotObservation:
        # Read arm position
        start = time.perf_counter()
-        obs_dict = self.bus.sync_read("Present_Position")
+        obs_dict = self.bus.sync_read("Present_Position", num_retry=self.config.num_read_retries)
        obs_dict = {f"{motor}.pos": val for motor, val in obs_dict.items()}
        dt_ms = (time.perf_counter() - start) * 1e3
        logger.debug(f"{self} read state: {dt_ms:.1f}ms")
@@ -221,7 +221,7 @@ class SOFollower(Robot):
        # Cap goal position when too far away from present position.
        # /!\ Slower fps expected due to reading from the follower.
        if self.config.max_relative_target is not None:
-            present_pos = self.bus.sync_read("Present_Position")
+            present_pos = self.bus.sync_read("Present_Position", num_retry=self.config.num_read_retries)
            goal_present_pos = {key: (g_pos, present_pos[key]) for key, g_pos in goal_pos.items()}
            goal_pos = ensure_safe_goal_position(goal_present_pos, self.config.max_relative_target)

@@ -62,12 +62,34 @@ class UnitreeG1Config(RobotConfig):
    # Socket config for ZMQ bridge
    robot_ip: str = "192.168.123.164"  # default G1 IP

+    # Run the locomotion / whole-body controller ONBOARD the robot (policy on the G1
+    # itself, against local DDS at full rate) instead of on the laptop over the ZMQ
+    # socket bridge. In this mode the robot object uses the real Unitree SDK channels
+    # and expects high-level actions (arm targets + joystick axes, or 64-D SONIC
+    # tokens) fed via send_action -- e.g. by run_g1_server's serve_onboard_controller,
+    # which receives them from the laptop over ZMQ. Mutually exclusive with is_simulation.
+    onboard: bool = False
+    # DDS network interface for onboard mode (None = SDK default, matching
+    # run_g1_server.py's ChannelFactoryInitialize(0)).
+    dds_interface: str | None = None
+    # Onboard sub-flags. On a real G1 both are True: the built-in motion services
+    # must be released before we can write lowcmd, and locomotion axes are read from
+    # the physical wireless remote. Against a DDS sim neither applies (no
+    # MotionSwitcher, no physical remote), so set both False so the controller takes
+    # its locomotion axes purely from send_action (ZMQ) input.
+    release_motion_control: bool = True
+    physical_remote: bool = True
+
    # Cameras (ZMQ-based remote cameras)
    cameras: dict[str, CameraConfig] = field(default_factory=dict)

    # Compensates for gravity on the unitree's arms using the arm ik solver
    gravity_compensation: bool = False

-    # Lower-body controller class name, e.g. "GrootLocomotionController" or
-    # "HolosomaLocomotionController". None disables it.
+    # 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
@@ -0,0 +1,27 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unitree G1 locomotion controllers (Groot, Holosoma, SONIC)."""
+
+from .gr00t_locomotion import GrootLocomotionController
+from .holosoma_locomotion import HolosomaLocomotionController
+from .sonic_whole_body import SonicWholeBodyController
+
+__all__ = [
+    "GrootLocomotionController",
+    "HolosomaLocomotionController",
+    "SonicWholeBodyController",
+]
@@ -14,6 +14,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.

+from __future__ import annotations
+
 import logging
 from collections import deque

@@ -21,7 +23,7 @@ import numpy as np
 import onnxruntime as ort
 from huggingface_hub import hf_hub_download

-from .g1_utils import (
+from ..g1_utils import (
    REMOTE_AXES,
    REMOTE_BUTTONS,
    G1_29_JointIndex,
@@ -68,9 +70,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 ..g1_utils 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")

@@ -196,6 +204,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):
@@ -14,18 +14,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.

-import json
+from __future__ import annotations
+
 import logging

 import numpy as np
-import onnx
 import onnxruntime as ort
 from huggingface_hub import hf_hub_download

-from .g1_utils import (
+from ..g1_utils import (
    REMOTE_AXES,
    G1_29_JointArmIndex,
    G1_29_JointIndex,
+    compute_pd_gains,
    get_gravity_orientation,
 )

@@ -57,12 +58,23 @@ POLICY_FILES = {
    "ppo": "ppo_g1_29dof.onnx",
 }

+# Per-joint motor model in Holosoma's joint order, plus the joints that get a x2
+# stiffness/damping factor. These reproduce the kp/kd that used to be read from the
+# policy's ONNX metadata exactly (both fastsac and ppo), so gains are now derived from
+# the shared motor model (see g1_utils.compute_pd_gains) instead.
+HOLOSOMA_MOTOR_MODELS = (
+    ["7520_14", "7520_22", "7520_14", "7520_22", "5020", "5020"] * 2
+    + ["7520_14", "5020", "5020"]
+    + ["5020", "5020", "5020", "5020", "5020", "4010", "4010"] * 2
+)
+HOLOSOMA_DOUBLE = {4, 5, 10, 11, 13, 14}
+

 def load_policy(
    repo_id: str = DEFAULT_HOLOSOMA_REPO_ID,
    policy_type: str = "fastsac",
 ) -> tuple[ort.InferenceSession, np.ndarray, np.ndarray]:
-    """Load Holosoma locomotion policy and extract KP/KD from metadata.
+    """Load the Holosoma locomotion policy and its motor-model-derived PD gains.

    Args:
        repo_id: Hugging Face Hub repo ID
@@ -81,16 +93,7 @@ def load_policy(
    policy = ort.InferenceSession(policy_path)
    logger.info(f"Policy loaded: {policy.get_inputs()[0].shape} → {policy.get_outputs()[0].shape}")

-    # Extract KP/KD from ONNX metadata
-    model = onnx.load(policy_path, load_external_data=False)
-    metadata = {prop.key: prop.value for prop in model.metadata_props}
-
-    if "kp" not in metadata or "kd" not in metadata:
-        raise ValueError("ONNX model must contain 'kp' and 'kd' in metadata")
-
-    kp = np.array(json.loads(metadata["kp"]), dtype=np.float32)
-    kd = np.array(json.loads(metadata["kd"]), dtype=np.float32)
-    logger.info(f"Loaded KP/KD from ONNX ({len(kp)} joints)")
+    kp, kd = compute_pd_gains(HOLOSOMA_MOTOR_MODELS, HOLOSOMA_DOUBLE)

    return policy, kp, kd

@@ -0,0 +1,401 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""SONIC decoder whole-body controller for the Unitree G1 (token-only).
+
+Pure-Python/ONNX re-implementation of the *decode* half of NVIDIA's SONIC deploy stack.
+The encoder is intentionally absent: a token-output VLA (e.g. ``nepyope/sonic_walk``)
+supplies the 64-D latent ``motion_token`` directly each tick, and the SONIC **decoder**
+maps ``token + recent proprioception history`` to a residual action that is scaled and
+added onto ``DEFAULT_ANGLES`` to produce 50 Hz joint-position targets for the robot's PD
+controller.
+
+Index spaces: joints exist in two orderings — **IsaacLab** (policy/training order) and
+**MuJoCo** (deploy order). ``ISAACLAB_TO_MUJOCO`` / ``MUJOCO_TO_ISAACLAB`` (in g1_utils)
+convert between them. Quaternions are scalar-first ``(w, x, y, z)``.
+"""
+
+from __future__ import annotations
+
+import logging
+
+import numpy as np
+import onnxruntime as ort
+from huggingface_hub import hf_hub_download
+
+from ..g1_utils import (
+    ISAACLAB_TO_MUJOCO,
+    MOTOR_ARMATURE,
+    MUJOCO_TO_ISAACLAB,
+    NATURAL_FREQ,
+    G1_29_JointIndex,
+    compute_pd_gains,
+    get_gravity_orientation,
+    lowstate_to_obs,
+    make_ort_session_options,
+)
+
+logger = logging.getLogger(__name__)
+
+# ── Constants (hardware-validated; see the NVIDIA SONIC deploy reference) ──────
+CONTROL_DT = 0.02  # 50 Hz control period (s)
+TOKEN_DIM = 64  # decoder latent size
+
+# Nominal standing pose (rad), 29 joints in IsaacLab order. Decoder actions are residuals
+# added on top of this.
+DEFAULT_ANGLES = np.array(
+    [
+        -0.312, 0.0, 0.0, 0.669, -0.363, 0.0,
+        -0.312, 0.0, 0.0, 0.669, -0.363, 0.0,
+        0.0, 0.0, 0.0,
+        0.2, 0.2, 0.0, 0.6, 0.0, 0.0, 0.0,
+        0.2, -0.2, 0.0, 0.6, 0.0, 0.0, 0.0,
+    ],
+    dtype=np.float32,
+)
+
+# Per-motor torque limits (N·m), used only for SONIC's residual-action scaling. The
+# armature / bandwidth constants and the PD-gain formula are shared (see g1_utils).
+EFFORT = {"5020": 25.0, "7520_14": 88.0, "7520_22": 139.0, "4010": 5.0}
+
+
+def _action_scale(k):
+    """Per-motor residual-action scale (maps policy output to joint-angle delta)."""
+    return 0.25 * EFFORT[k] / (MOTOR_ARMATURE[k] * NATURAL_FREQ**2)
+
+
+# Per-joint motor model (IsaacLab order): legs, waist, then arms. Single source of truth
+# for both ACTION_SCALE and compute_kp_kd().
+MOTOR_MODELS = (
+    ["7520_22", "7520_22", "7520_14", "7520_22", "5020", "5020"] * 2
+    + ["7520_14", "5020", "5020"]
+    + ["5020", "5020", "5020", "5020", "5020", "4010", "4010"] * 2
+)
+ACTION_SCALE = np.array([_action_scale(k) for k in MOTOR_MODELS], dtype=np.float32)  # (29,) IsaacLab
+
+
+def _to_mujoco(a):
+    """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 are a fixed convention validated against the deployed SONIC ONNX
+    policy (the decoder consumes vectors in this order). Do not "correct" the table or
+    rename toward the opposite direction without re-validating on hardware.
+    """
+    return a[MUJOCO_TO_ISAACLAB]
+
+
+DEFAULT_ANGLES_MUJOCO = _to_mujoco(DEFAULT_ANGLES)
+
+
+# Ankle + waist joint indices (IsaacLab order) that get a x2 stiffness/damping factor.
+_SONIC_DOUBLE = {4, 5, 10, 11, 13, 14}
+
+
+def compute_kp_kd():
+    """SONIC per-joint PD gains (kp, kd), (29,) float32 in IsaacLab joint order."""
+    return compute_pd_gains(MOTOR_MODELS, _SONIC_DOUBLE)
+
+
+# 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"
+
+
+# Startup blend duration: over the first control ticks, linearly interpolate every joint
+# from the robot's initial measured pose into the policy's commanded target, so control
+# eases in without a snap on the first command.
+INIT_RAMP_S = 3.0
+
+# Neutral ("zero pose") SONIC token, held by token_mode until the first real token arrives.
+# Captured from the encoder's own output while the robot stood idle in sim: the encoder is
+# an FSQ bottleneck (~5 bit/dim, Div(16)), so its tokens live on the 1/16 grid. We store the
+# integer FSQ codes and rescale by 1/16, giving an exact on-grid token -- unlike the literal
+# all-zero token, which is off the learned manifold and decodes to a slightly goofy stance.
+# This one decodes to a stable, natural standing pose.
+_NEUTRAL_TOKEN_CODES = np.array(
+    [-1, 3, 1, -1, 1, -3, 6, 1, 1, 1, -2, -4, -2, 0, -3, -1,
+     2, -1, -3, -5, 3, 1, 1, -4, -1, -1, 1, -7, 0, 1, 2, -2,
+     5, -2, -2, -4, 0, -1, 3, -1, 0, -5, -1, 0, -4, 0, 0, -1,
+     -1, 2, -2, 1, 3, 3, 1, 0, 0, 6, 0, -7, 3, 0, 2, -2],
+    dtype=np.float32,
+)
+NEUTRAL_TOKEN = _NEUTRAL_TOKEN_CODES / 16.0  # FSQ Div(16): integer codes -> on-grid token
+
+
+def _extract_token_from_action(action: dict | None) -> np.ndarray | None:
+    """Reassemble a dense (64,) latent token from ``motion_token.{i}`` keys, or None.
+
+    The token-only interface: the caller supplies the 64-D encoder latent directly (e.g. a
+    token-output VLA's action), 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)
+
+
+class SonicDecoder:
+    """Runs the SONIC decoder ONNX model and owns the proprioception history.
+
+    Each tick it appends the latest robot state to 10-frame history buffers, then maps the
+    supplied 64-D ``token`` + that history to a residual action added onto
+    ``DEFAULT_ANGLES``. The encoder is bypassed entirely (token supplied by the policy).
+    """
+
+    def __init__(self, decoder):
+        self.decoder = decoder
+        self.decoder_input = decoder.get_inputs()[0].name
+        dec_dim = int(decoder.get_inputs()[0].shape[1])
+        if dec_dim != 994:
+            raise RuntimeError(f"Unexpected decoder input dim {dec_dim} (expected 994)")
+        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
+
+    def reset(self):
+        """Clear the token and 10-frame proprioception history.
+
+        ``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
+
+    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)
+        q_mj = _to_mujoco(q)
+        dq_mj = _to_mujoco(dq)
+        self.h_q_mj = [q_mj - DEFAULT_ANGLES_MUJOCO] + self.h_q_mj[:-1]
+        self.h_dq_mj = [dq_mj] + self.h_dq_mj[:-1]
+        self.h_ang = [ang.copy()] + self.h_ang[:-1]
+        self.h_act_mj = [self.last_action_mj.copy()] + self.h_act_mj[:-1]
+        self.h_quat = [quat.copy()] + self.h_quat[:-1]
+
+    def build_decoder_obs(self):
+        """Assemble the 994-D decoder input: token + 10-frame proprioception history + gravity."""
+        obs = np.zeros(994, np.float32)
+        off = 0
+        obs[off : off + 64] = self.token
+        off += 64
+        for h, sz in [
+            (list(reversed(self.h_ang)), 3),
+            (list(reversed(self.h_q_mj)), 29),
+            (list(reversed(self.h_dq_mj)), 29),
+            (list(reversed(self.h_act_mj)), 29),
+        ]:
+            for f in range(10):
+                obs[off : off + sz] = h[f]
+                off += sz
+        for q in reversed(self.h_quat):
+            obs[off : off + 3] = get_gravity_orientation(q)
+            off += 3
+        assert off == 994, f"Decoder obs mismatch: {off}"
+        return obs
+
+    def step(self, robot_obs, token, debug=False):
+        """One control tick: read robot obs, decode the supplied token -> joint targets.
+
+        Args:
+            robot_obs: dict with ``<joint>.q``/``.dq`` and ``imu.*`` fields.
+            token: 64-D latent supplied by the policy (encoder bypassed).
+            debug: log action/delta norms.
+
+        Returns:
+            dict of ``<joint>.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", 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)
+        self.update_history(q, dq, ang, quat)
+        action_mj = (
+            self.decoder.run(None, {self.decoder_input: self.build_decoder_obs().reshape(1, -1)})[0]
+            .squeeze()
+            .astype(np.float32)
+        )
+        self.last_action_mj = action_mj.copy()
+        target = DEFAULT_ANGLES + action_mj[ISAACLAB_TO_MUJOCO] * ACTION_SCALE
+        if debug:
+            delta = target - q
+            logger.debug(
+                "token_norm=%.4f action_norm=%.4f delta_max=%.4f delta_rms=%.4f",
+                np.linalg.norm(self.token),
+                np.linalg.norm(action_mj),
+                np.max(np.abs(delta)),
+                np.sqrt(np.mean(delta**2)),
+            )
+        return {f"{m.name}.q": float(target[m.value]) for m in G1_29_JointIndex}
+
+
+class SonicRuntime:
+    """Loads the SONIC decoder ONNX model and owns the decode controller.
+
+    Token-only deploy: the encoder is bypassed; each tick the decoder consumes a 64-D
+    latent token supplied directly by the policy.
+    """
+
+    def __init__(self):
+        decoder_path = hf_hub_download(repo_id="nvidia/GEAR-SONIC", filename="model_decoder.onnx")
+
+        so = make_ort_session_options()
+        decoder_sess = ort.InferenceSession(decoder_path, sess_options=so)
+
+        self.kp, self.kd = compute_kp_kd()
+        self.controller = SonicDecoder(decoder_sess)
+
+    @property
+    def pipeline(self):
+        return self.controller
+
+    def reset(self):
+        self.controller.reset()
+
+    def shutdown(self):
+        pass
+
+
+class SonicWholeBodyController:
+    """Full-body SONIC controller for UnitreeG1's background controller thread."""
+
+    control_dt = CONTROL_DT
+    full_body = True
+
+    def __init__(self):
+        logger.info("Loading SONIC whole-body controller...")
+        self._runtime = SonicRuntime()
+        self.kp = self._runtime.kp
+        self.kd = self._runtime.kd
+        self.controller = self._runtime.controller
+
+        # Startup blend: ease from the robot's initial pose into the first commanded policy
+        # targets over INIT_RAMP_S (captured on the first control tick).
+        self._init_ramp_steps = max(1, round(INIT_RAMP_S / CONTROL_DT))
+        self._init_step = 0
+        self._start_pose: dict[str, float] = {}
+
+        # 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
+        # uniformly to run_g1_server, lerobot-rollout and the sim replays.
+        self.token_mode = False
+        self._last_token: np.ndarray | None = None
+
+        logger.info("SONIC ready (decoder, 64-D token command path)")
+
+    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
+        live commanded target, so the handoff has no snap.
+
+        ``out`` is the policy's ``<joint>.q`` target dict for this tick; the blend ratio
+        climbs 0->1 over the ramp, after which the raw policy target passes through.
+        """
+        if self._init_step >= self._init_ramp_steps or not out:
+            return out
+        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", DEFAULT_ANGLES[m.value]))
+                for m in G1_29_JointIndex
+            }
+        self._init_step += 1
+        ratio = min(1.0, self._init_step / self._init_ramp_steps)
+        blended = {
+            k: self._start_pose.get(k, float(tgt)) * (1.0 - ratio) + float(tgt) * ratio
+            for k, tgt in out.items()
+        }
+        if self._init_step >= self._init_ramp_steps:
+            logger.info("SONIC startup blend complete -> full policy control")
+        return blended
+
+    def run_step(self, action: dict, lowstate) -> dict:
+        if lowstate is None:
+            return {}
+        obs = lowstate_to_obs(lowstate)
+
+        # Token-only interface (token-output VLA): a dense 64-D ``motion_token.{i}`` command
+        # is decoded directly, encoder bypassed.
+        token = _extract_token_from_action(action)
+        if token is not None:
+            self._last_token = token
+        elif self._last_token is None and self.token_mode:
+            # Token-driven deploy, but no token has arrived yet: hold the captured neutral
+            # token (NEUTRAL_TOKEN), which the decoder maps to a stable, natural standing pose.
+            self._last_token = NEUTRAL_TOKEN.copy()
+        if self._last_token is None:
+            # No token yet and not in token_mode: hold (keep last target).
+            return {}
+        # 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))
+
+    def reset(self):
+        self._runtime.reset()
+        self._init_step = 0  # re-run the startup blend after a reset
+        self._start_pose = {}
+        # Drop the held token so token_mode re-seeds the neutral token after a reset.
+        self._last_token = None
+
+    def shutdown(self):
+        self._runtime.shutdown()
@@ -23,11 +23,82 @@ import numpy as np

 NUM_MOTORS = 29

+# Joint-order permutations between the two 29-DoF layouts used across the G1 stack:
+# IsaacLab (policy/training order) and MuJoCo (deploy order). ``a[ISAACLAB_TO_MUJOCO]``
+# reorders an IsaacLab-ordered vector into MuJoCo order, and vice-versa.
+ISAACLAB_TO_MUJOCO = np.array(
+    [
+        0,
+        3,
+        6,
+        9,
+        13,
+        17,
+        1,
+        4,
+        7,
+        10,
+        14,
+        18,
+        2,
+        5,
+        8,
+        11,
+        15,
+        19,
+        21,
+        23,
+        25,
+        27,
+        12,
+        16,
+        20,
+        22,
+        24,
+        26,
+        28,
+    ],
+    dtype=np.int32,
+)
+MUJOCO_TO_ISAACLAB = np.array(
+    [
+        0,
+        6,
+        12,
+        1,
+        7,
+        13,
+        2,
+        8,
+        14,
+        3,
+        9,
+        15,
+        22,
+        4,
+        10,
+        16,
+        23,
+        5,
+        11,
+        17,
+        24,
+        18,
+        25,
+        19,
+        26,
+        20,
+        27,
+        21,
+        28,
+    ],
+    dtype=np.int32,
+)
+
 REMOTE_AXES = ("remote.lx", "remote.ly", "remote.rx", "remote.ry")
 REMOTE_BUTTONS = tuple(f"remote.button.{i}" for i in range(16))
 REMOTE_KEYS = REMOTE_AXES + REMOTE_BUTTONS

-
 def default_remote_input() -> dict[str, float]:
    """Return a zeroed-out remote input dict (axes + buttons)."""
    return dict.fromkeys(REMOTE_KEYS, 0.0)
@@ -43,6 +114,53 @@ def get_gravity_orientation(quaternion: list[float] | np.ndarray) -> np.ndarray:
    return gravity_orientation


+# Unitree motor-model parameters shared by the controllers that derive their PD gains
+# from motor physics rather than hand-tuning (SONIC decoder, Holosoma). NATURAL_FREQ is
+# the target closed-loop stiffness bandwidth (rad/s); MOTOR_ARMATURE is per-model rotor
+# inertia (keys are Unitree motor model names). From these: kp = armature * w**2 and
+# kd = 4 * armature * w, with an optional x2 factor on stiff joints (ankles/waist).
+NATURAL_FREQ = 10.0 * 2.0 * np.pi
+MOTOR_ARMATURE = {"5020": 0.003609725, "7520_14": 0.010177520, "7520_22": 0.025101925, "4010": 0.00425}
+
+
+def compute_pd_gains(motor_models, double_indices=()) -> tuple[np.ndarray, np.ndarray]:
+    """Derive per-joint PD gains (kp, kd) from motor armature and target bandwidth.
+
+    ``motor_models`` is a per-joint sequence of Unitree motor model names (in the
+    controller's own joint order); joints whose index is in ``double_indices`` get a
+    x2 stiffness/damping factor. Returns two (N,) float32 arrays in that same order.
+    """
+    double = set(double_indices)
+
+    def s(k):
+        return MOTOR_ARMATURE[k] * NATURAL_FREQ**2
+
+    def d(k):
+        return 4.0 * MOTOR_ARMATURE[k] * NATURAL_FREQ
+
+    kp = np.array([2 * s(k) if i in double else s(k) for i, k in enumerate(motor_models)], dtype=np.float32)
+    kd = np.array([2 * d(k) if i in double else d(k) for i, k in enumerate(motor_models)], dtype=np.float32)
+    return kp, kd
+
+
+def make_ort_session_options(intra_op_num_threads: int | None = None, inter_op_num_threads: int | None = None):
+    """Build quiet ONNX Runtime SessionOptions, optionally capping the CPU thread pool.
+
+    These tiny MLP policies are latency-bound, not throughput-bound, so letting ORT grab
+    every core starves the real-time control loop / torch policy and causes stutter. Pass
+    1 intra + 1 inter thread for lowest-latency per-step inference.
+    """
+    import onnxruntime as ort
+
+    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
+
+
 class G1_29_JointArmIndex(IntEnum):
    # Left arm
    kLeftShoulderPitch = 15
@@ -63,13 +181,55 @@ class G1_29_JointArmIndex(IntEnum):
    kRightWristYaw = 28


+def lowstate_to_obs(lowstate) -> dict:
+    """Build a robot observation dict from a Unitree lowstate.
+
+    Shared by ``UnitreeG1.get_observation`` and the SONIC pipeline so the
+    lowstate -> obs mapping lives in exactly one place. Keys match the
+    ``<joint>.q``/``imu.*`` schema consumed across the controllers.
+    """
+    obs: dict = {}
+
+    for motor in G1_29_JointIndex:
+        idx = motor.value
+        obs[f"{motor.name}.q"] = lowstate.motor_state[idx].q
+        obs[f"{motor.name}.dq"] = lowstate.motor_state[idx].dq
+        obs[f"{motor.name}.tau"] = lowstate.motor_state[idx].tau_est
+
+    imu = lowstate.imu_state
+    if imu.gyroscope:
+        obs["imu.gyro.x"] = imu.gyroscope[0]
+        obs["imu.gyro.y"] = imu.gyroscope[1]
+        obs["imu.gyro.z"] = imu.gyroscope[2]
+    if imu.accelerometer:
+        obs["imu.accel.x"] = imu.accelerometer[0]
+        obs["imu.accel.y"] = imu.accelerometer[1]
+        obs["imu.accel.z"] = imu.accelerometer[2]
+    if imu.quaternion:
+        obs["imu.quat.w"] = imu.quaternion[0]
+        obs["imu.quat.x"] = imu.quaternion[1]
+        obs["imu.quat.y"] = imu.quaternion[2]
+        obs["imu.quat.z"] = imu.quaternion[3]
+    if imu.rpy:
+        obs["imu.rpy.roll"] = imu.rpy[0]
+        obs["imu.rpy.pitch"] = imu.rpy[1]
+        obs["imu.rpy.yaw"] = imu.rpy[2]
+
+    wr = getattr(lowstate, "wireless_remote", None)
+    if wr:
+        obs["wireless_remote"] = bytes(wr) if not isinstance(wr, (bytes, bytearray)) else wr
+
+    return obs
+
+
 def make_locomotion_controller(name: str | None):
    """Instantiate a locomotion controller by class name. Returns None if name is None."""
    if name is None:
        return None
    controllers = {
-        "GrootLocomotionController": "lerobot.robots.unitree_g1.gr00t_locomotion",
-        "HolosomaLocomotionController": "lerobot.robots.unitree_g1.holosoma_locomotion",
+        "GrootLocomotionController": "lerobot.robots.unitree_g1.controllers.gr00t_locomotion",
+        "HolosomaLocomotionController": "lerobot.robots.unitree_g1.controllers.holosoma_locomotion",
+        "SonicWholeBodyController": "lerobot.robots.unitree_g1.controllers.sonic_whole_body",
    }
    module_path = controllers.get(name)
    if module_path is None:
@@ -22,16 +22,33 @@ This server runs on the robot and forwards:
 - Robot commands (LowCmd) from ZMQ to DDS (from remote clients)

 Uses JSON for secure serialization instead of pickle.
+
+Controller-negotiation handshake
+--------------------------------
+The first message from a client agrees on which controller the server will run onboard
+(``serve_onboard_controller``); the controller NEVER runs on the laptop client.
+Test the handshake in isolation (no DDS, runs on a laptop) in two terminals::
+
+    # terminal A: handshake-only server
+    python -m lerobot.robots.unitree_g1.run_g1_server --handshake-only
+
+    # terminal B: client proposes a controller
+    python -m lerobot.robots.unitree_g1.run_g1_server \\
+        --handshake-client SonicWholeBodyController --sonic-token-action --server-ip 127.0.0.1
+
+On the real robot, add ``--handshake`` to the normal bridge to require agreement first.
 """

 import argparse
 import base64
 import contextlib
 import json
+import signal
 import threading
 import time
 from typing import Any

+import numpy as np
 import zmq
 from unitree_sdk2py.comm.motion_switcher.motion_switcher_client import MotionSwitcherClient
 from unitree_sdk2py.core.channel import ChannelFactoryInitialize, ChannelPublisher, ChannelSubscriber
@@ -50,6 +67,253 @@ LOWCMD_PORT = 6000
 LOWSTATE_PORT = 6001
 NUM_MOTORS = 35

+# Onboard high-level channels (serve_onboard_controller): compact actions in, state out.
+ACTION_PORT = 6004
+STATE_PORT = 6005
+
+# Controller-negotiation handshake (REQ/REP). The client's first message agrees on
+# which controller the server will run before any control data flows.
+HANDSHAKE_PORT = 6002
+PROTOCOL_VERSION = 1
+
+# Controllers that can run ONBOARD (must match g1_utils.make_locomotion_controller).
+# ``None`` (a.k.a. "bridge") means no onboard controller: the laptop owns control and
+# streams raw lowcmd over the ZMQ DDS bridge (the legacy run_g1_server behavior).
+VALID_CONTROLLERS = (
+    "GrootLocomotionController",
+    "HolosomaLocomotionController",
+    "SonicWholeBodyController",
+)
+# SONIC latent-token dimensionality (mirrors sonic_whole_body.TOKEN_DIM; kept local so
+# the handshake can run without importing the heavy controller / onnxruntime).
+TOKEN_DIM = 64
+_BRIDGE_ALIASES = {"", "none", "null", "bridge", "raw"}
+
+
+def _normalize_controller(name: str | None) -> str | None:
+    """Map a requested controller name to a canonical value (or None for raw bridge)."""
+    if name is None:
+        return None
+    low = str(name).strip().lower()
+    if low in _BRIDGE_ALIASES:
+        return None
+    for c in VALID_CONTROLLERS:
+        if c.lower() == low:
+            return c
+    raise ValueError(f"Unknown controller {name!r}. Available: {list(VALID_CONTROLLERS)} or 'bridge'")
+
+
+def _capabilities(controller: str | None, sonic_token_action: bool) -> dict[str, Any]:
+    """The interface the server advertises for an agreed controller."""
+    caps: dict[str, Any] = {
+        "controller": controller,
+        "sonic_token_action": bool(sonic_token_action),
+        "protocol": PROTOCOL_VERSION,
+    }
+    if controller is None:
+        # Raw DDS bridge: the laptop runs the controller and streams lowcmd.
+        caps["mode"] = "bridge"
+        caps["lowcmd_port"] = LOWCMD_PORT
+        caps["lowstate_port"] = LOWSTATE_PORT
+    else:
+        # Onboard: the controller runs here; the laptop ships compact high-level actions.
+        caps["mode"] = "onboard"
+        caps["action_port"] = ACTION_PORT
+        caps["state_port"] = STATE_PORT
+        if sonic_token_action:
+            caps["action_space"] = "motion_token"
+            caps["action_dim"] = TOKEN_DIM
+    return caps
+
+
+def negotiate_controller(sock: zmq.Socket, shutdown_event: threading.Event) -> dict[str, Any]:
+    """Server side of the handshake: block on one REP socket until a client sends a
+    valid ``hello``, then reply with the negotiated capabilities and return them.
+
+    Rejects malformed / unknown-controller requests with an error reply and keeps
+    waiting (a rejected client can retry). Honors ``shutdown_event`` so Ctrl-C works.
+    """
+    poller = zmq.Poller()
+    poller.register(sock, zmq.POLLIN)
+    while not shutdown_event.is_set():
+        if not dict(poller.poll(timeout=200)):
+            continue
+        raw = sock.recv()
+        try:
+            hello = json.loads(raw.decode("utf-8"))
+        except (ValueError, UnicodeDecodeError) as e:
+            sock.send_json({"type": "error", "ok": False, "error": f"bad hello: {e}"})
+            continue
+        try:
+            controller = _normalize_controller(hello.get("controller"))
+        except ValueError as e:
+            sock.send_json(
+                {"type": "error", "ok": False, "error": str(e), "available": list(VALID_CONTROLLERS)}
+            )
+            continue
+        reply = {"type": "welcome", "ok": True, **_capabilities(controller, hello.get("sonic_token_action", False))}
+        sock.send_json(reply)
+        return reply
+    raise KeyboardInterrupt
+
+
+def request_controller(
+    server_ip: str,
+    controller: str | None,
+    *,
+    sonic_token_action: bool = False,
+    port: int = HANDSHAKE_PORT,
+    timeout_s: float = 5.0,
+) -> dict[str, Any]:
+    """Client side of the handshake: propose a controller, return the server's agreed
+    capabilities (or raise on rejection / timeout)."""
+    ctx = zmq.Context.instance()
+    sock = ctx.socket(zmq.REQ)
+    sock.setsockopt(zmq.LINGER, 0)
+    sock.setsockopt(zmq.RCVTIMEO, int(timeout_s * 1000))
+    sock.setsockopt(zmq.SNDTIMEO, int(timeout_s * 1000))
+    sock.connect(f"tcp://{server_ip}:{port}")
+    hello = {
+        "type": "hello",
+        "controller": controller,
+        "sonic_token_action": bool(sonic_token_action),
+        "protocol": PROTOCOL_VERSION,
+    }
+    try:
+        sock.send_json(hello)
+        reply = sock.recv_json()
+    except zmq.Again as e:
+        raise TimeoutError(f"no handshake reply from {server_ip}:{port} within {timeout_s}s") from e
+    finally:
+        sock.close(linger=0)
+    if not reply.get("ok"):
+        raise RuntimeError(f"handshake rejected: {reply.get('error')} (available: {reply.get('available')})")
+    return reply
+
+
+def serve_onboard_controller(
+    *,
+    controller: str,
+    sonic_token_action: bool,
+    dds_interface: str | None = None,
+    sim: bool = False,
+    cameras: dict | None = None,
+    camera_fps: int = 30,
+    camera_port: int = 5555,
+    action_port: int = ACTION_PORT,
+    state_port: int = STATE_PORT,
+    state_fps: float = 30.0,
+    stop: threading.Event | None = None,
+) -> None:
+    """Run the negotiated controller ONBOARD -- the single control path on the robot.
+
+    Builds ``UnitreeG1(onboard=True, controller=...)`` so the controller/balance loop runs
+    locally against DDS at full rate (the 50 Hz ``_controller_loop`` thread lives in
+    UnitreeG1), then receives compact high-level actions from the laptop over ZMQ
+    (:action_port), decodes them via the controller, publishes ``observation.state``
+    (:state_port), and optionally serves the ego camera. The controller NEVER runs on the
+    laptop; the laptop (lerobot-rollout thin-client) only ships tokens/axes and reads back
+    state + camera frames.
+    """
+    # Imported lazily: UnitreeG1 imports request_controller from this module, so a
+    # top-level import here would be circular.
+    from lerobot.robots.unitree_g1.config_unitree_g1 import UnitreeG1Config
+    from lerobot.robots.unitree_g1.unitree_g1 import UnitreeG1
+
+    if stop is None:
+        stop = threading.Event()
+        signal.signal(signal.SIGINT, lambda *_: stop.set())
+        signal.signal(signal.SIGTERM, lambda *_: stop.set())
+
+    cfg = UnitreeG1Config(
+        is_simulation=False,
+        onboard=True,
+        controller=controller,
+        dds_interface=dds_interface,
+        release_motion_control=not sim,
+        physical_remote=not sim,
+        cameras={},
+    )
+
+    # Optional camera server (background daemon thread; independent of DDS).
+    if cameras:
+        camera_server = ImageServer({"fps": camera_fps, "cameras": cameras}, port=camera_port)
+        threading.Thread(target=camera_server.run, daemon=True).start()
+        cam_summary = ", ".join(f"{name}(dev {c['device_id']})" for name, c in cameras.items())
+        print(f"Camera server started on :{camera_port}: {cam_summary}")
+
+    robot = UnitreeG1(cfg)
+    print(f"Connecting onboard robot (controller={controller}, token={sonic_token_action})...")
+    robot.connect()
+
+    ctx = zmq.Context.instance()
+    sock = ctx.socket(zmq.PULL)
+    sock.setsockopt(zmq.CONFLATE, 1)  # only ever act on the freshest command
+    sock.setsockopt(zmq.RCVTIMEO, 200)  # keeps the loop responsive to the stop event
+    sock.bind(f"tcp://0.0.0.0:{action_port}")
+    print(f"Onboard controller live. Waiting for laptop actions on :{action_port} ...")
+    print("Ctrl-C for graceful shutdown.")
+
+    state_sock = None
+    if state_fps > 0:
+        state_sock = ctx.socket(zmq.PUB)
+        state_sock.setsockopt(zmq.SNDHWM, 2)
+        state_sock.setsockopt(zmq.LINGER, 0)
+        state_sock.bind(f"tcp://0.0.0.0:{state_port}")
+        print(f"Publishing observation.state on :{state_port} at {state_fps:.0f} Hz")
+
+        def publish_state() -> None:
+            period = 1.0 / state_fps
+            while not stop.is_set():
+                t0 = time.time()
+                obs = robot.get_observation()
+                if obs:
+                    # Forward every scalar proprio key the robot exposes (29 joint .q, IMU,
+                    # and the SONIC token echo: 64-D motion_token_state.*). Camera arrays are
+                    # streamed separately by the ImageServer, so drop ndarrays here. This
+                    # makes the laptop thin-client a pure relay.
+                    state = {
+                        k: float(v)
+                        for k, v in obs.items()
+                        if isinstance(v, (bool, int, float, np.floating, np.integer))
+                    }
+                    with contextlib.suppress(zmq.Again):
+                        state_sock.send_json(state, zmq.NOBLOCK)
+                time.sleep(max(0.0, period - (time.time() - t0)))
+
+        threading.Thread(target=publish_state, daemon=True).start()
+    else:
+        print("observation.state PUB disabled (state_fps<=0)")
+
+    n = 0
+    try:
+        while not stop.is_set():
+            try:
+                payload = sock.recv()
+            except zmq.Again:
+                continue
+            except zmq.ContextTerminated:
+                break
+
+            try:
+                action = json.loads(payload.decode("utf-8"))
+            except (ValueError, UnicodeDecodeError) as e:
+                print(f"Dropping malformed action: {e}")
+                continue
+
+            robot.send_action(action)
+
+            n += 1
+            if n % 60 == 0:
+                print(f"Applied {n} actions")
+    finally:
+        print("Shutting down onboard controller...")
+        stop.set()
+        if state_sock is not None:
+            with contextlib.suppress(Exception):
+                state_sock.close(linger=0)
+        robot.disconnect()
+

 def lowstate_to_dict(msg: hg_LowState) -> dict[str, Any]:
    """Convert LowState SDK message to a JSON-serializable dictionary."""
@@ -160,8 +424,86 @@ def main() -> None:
    parser.add_argument("--camera-width", type=int, default=640, help="Camera width (default: 640)")
    parser.add_argument("--camera-height", type=int, default=480, help="Camera height (default: 480)")
    parser.add_argument("--camera-port", type=int, default=5555, help="Camera ZMQ port (default: 5555)")
+    # Controller-negotiation handshake (first message agrees on the controller).
+    parser.add_argument("--handshake", action="store_true",
+                        help="Wait for a client to negotiate the controller before bridging")
+    parser.add_argument("--handshake-port", type=int, default=HANDSHAKE_PORT,
+                        help=f"Handshake REQ/REP port (default: {HANDSHAKE_PORT})")
+    parser.add_argument("--handshake-only", action="store_true",
+                        help="Run ONLY the handshake server (no DDS/cameras) to test negotiation")
+    parser.add_argument("--handshake-client", default=None, metavar="CONTROLLER",
+                        help="Act as a client: propose CONTROLLER (or 'bridge') to --server-ip and print the reply")
+    parser.add_argument("--server-ip", default="127.0.0.1", help="[--handshake-client] server IP")
+    parser.add_argument("--sonic-token-action", action="store_true",
+                        help="[handshake] negotiate the 64-D SONIC token action interface")
    args = parser.parse_args()

+    # --- Isolated handshake test paths (no DDS, safe to run on a laptop) ---
+    if args.handshake_client is not None:
+        controller = None if args.handshake_client.strip().lower() in _BRIDGE_ALIASES else args.handshake_client
+        reply = request_controller(
+            args.server_ip, controller,
+            sonic_token_action=args.sonic_token_action, port=args.handshake_port,
+        )
+        print(json.dumps(reply, indent=2))
+        return
+
+    if args.handshake_only:
+        ctx = zmq.Context.instance()
+        rep = ctx.socket(zmq.REP)
+        rep.bind(f"tcp://0.0.0.0:{args.handshake_port}")
+        print(f"[handshake] server listening on :{args.handshake_port} (no DDS). Ctrl-C to stop.")
+        shutdown = threading.Event()
+        try:
+            while True:
+                reply = negotiate_controller(rep, shutdown)
+                print(f"[handshake] agreed: controller={reply['controller']} mode={reply['mode']} "
+                      f"sonic_token_action={reply['sonic_token_action']}")
+        except KeyboardInterrupt:
+            print("\n[handshake] stopping")
+        finally:
+            rep.close(linger=0)
+            ctx.term()
+        return
+
+    # Controller-negotiation handshake: the client's first message agrees on the
+    # controller, which we then run ONBOARD (the controller NEVER runs on the laptop).
+    # Bridge/None falls through to the legacy raw DDS forward (deprecated laptop control).
+    if args.handshake:
+        ctx = zmq.Context.instance()
+        hs = ctx.socket(zmq.REP)
+        hs.bind(f"tcp://0.0.0.0:{args.handshake_port}")
+        print(f"[handshake] waiting for client controller agreement on :{args.handshake_port} ...")
+        shutdown = threading.Event()
+        try:
+            agreed = negotiate_controller(hs, shutdown)
+        except KeyboardInterrupt:
+            print("[handshake] interrupted before agreement; exiting")
+            hs.close(linger=0)
+            ctx.term()
+            return
+        hs.close(linger=0)
+        if agreed["controller"] is not None:
+            print(f"[handshake] running controller ONBOARD: {agreed['controller']} "
+                  f"(sonic_token_action={agreed['sonic_token_action']})")
+            cameras = None
+            if args.camera:
+                cameras = {
+                    "head_camera": {
+                        "device_id": args.camera_device,
+                        "shape": [args.camera_height, args.camera_width],
+                    }
+                }
+            serve_onboard_controller(
+                controller=agreed["controller"],
+                sonic_token_action=bool(agreed["sonic_token_action"]),
+                cameras=cameras,
+                camera_fps=args.camera_fps,
+                camera_port=args.camera_port,
+            )
+            return
+        print("[handshake] client selected raw DDS bridge (laptop owns control) -> legacy forward.")
+
    # Optionally start camera server in background thread
    camera_thread = None
    if args.camera:
@@ -205,6 +547,7 @@ def main() -> None:

    # initialize ZMQ
    ctx = zmq.Context.instance()
+    shutdown_event = threading.Event()

    # receive commands from remote client
    lowcmd_sock = ctx.socket(zmq.PULL)
@@ -215,7 +558,6 @@ def main() -> None:
    lowstate_sock.bind(f"tcp://0.0.0.0:{LOWSTATE_PORT}")

    state_period = 0.002  # ~500 hz
-    shutdown_event = threading.Event()

    # start observation forwarding in background thread
    t_state = threading.Thread(
@@ -16,6 +16,8 @@

 from __future__ import annotations

+import contextlib
+import json
 import logging
 import threading
 import time
@@ -26,6 +28,7 @@ from typing import TYPE_CHECKING, Protocol, runtime_checkable
 import numpy as np

 from lerobot.cameras import make_cameras_from_configs
+from lerobot.utils.errors import DeviceNotConnectedError
 from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.import_utils import _unitree_sdk_available, require_package

@@ -34,10 +37,10 @@ from .config_unitree_g1 import UnitreeG1Config
 from .g1_kinematics import G1_29_ArmIK
 from .g1_utils import (
    REMOTE_AXES,
-    REMOTE_KEYS,
    G1_29_JointArmIndex,
    G1_29_JointIndex,
    default_remote_input,
+    lowstate_to_obs,
    make_locomotion_controller,
 )

@@ -47,7 +50,9 @@ if TYPE_CHECKING or _unitree_sdk_available:
        ChannelPublisher as _SDKChannelPublisher,
        ChannelSubscriber as _SDKChannelSubscriber,
    )
-    from unitree_sdk2py.idl.default import unitree_hg_msg_dds__LowCmd_
+    from unitree_sdk2py.idl.default import (
+        unitree_hg_msg_dds__LowCmd_,
+    )
    from unitree_sdk2py.idl.unitree_hg.msg.dds_ import (
        LowCmd_ as hg_LowCmd,
        LowState_ as hg_LowState,
@@ -79,6 +84,14 @@ class LocomotionController(Protocol):
 kTopicLowCommand_Debug = "rt/lowcmd"
 kTopicLowState = "rt/lowstate"

+# Wireless-remote button byte layout, mapped to the positional button indices the
+# locomotion controllers expect. Used in onboard mode to read the physical Unitree
+# remote from lowstate (mirrors the exo teleoperator's RemoteController).
+_REMOTE_BUTTON_MAP: list[str] = [
+    "RB", "LB", "start", "back", "RT", "LT", "", "",
+    "A", "B", "X", "Y", "up", "right", "down", "left",
+]
+

@dataclass
 class MotorState:
@@ -119,24 +132,37 @@ class UnitreeG1(Robot):
        self.config = config
        self.control_dt = config.control_dt

+        # Three mutually-exclusive roles:
+        #   * simulation : local DDS + controller run in-process against a MuJoCo world.
+        #   * onboard    : local DDS + controller run in-process on the robot NX.
+        #   * client     : thin laptop client. No DDS, no controller. It negotiates a
+        #                  controller with ``run_g1_server`` (which runs it onboard),
+        #                  PUSHes high-level actions and reads back state + cameras over
+        #                  ZMQ. The controller *always* runs on the robot, never here.
+        self._client = not config.is_simulation and not config.onboard
+
        # Initialize cameras config (ZMQ-based) - actual connection in connect()
        self._cameras = make_cameras_from_configs(config.cameras)

-        # Import channel classes based on mode
-        if config.is_simulation:
+        # DDS channel classes are only needed by the in-process control roles. The thin
+        # client never touches DDS, so we don't import the socket shim at all.
+        if config.is_simulation or config.onboard:
            self._ChannelFactoryInitialize = _SDKChannelFactoryInitialize
            self._ChannelPublisher = _SDKChannelPublisher
            self._ChannelSubscriber = _SDKChannelSubscriber
        else:
-            from .unitree_sdk2_socket import (
-                ChannelFactoryInitialize,
-                ChannelPublisher,
-                ChannelSubscriber,
-            )
+            self._ChannelFactoryInitialize = None
+            self._ChannelPublisher = None
+            self._ChannelSubscriber = None

-            self._ChannelFactoryInitialize = ChannelFactoryInitialize
-            self._ChannelPublisher = ChannelPublisher
-            self._ChannelSubscriber = ChannelSubscriber
+        # Client-side ZMQ handles / negotiated capabilities (populated in connect()).
+        self._client_action_sock = None
+        self._client_state_sock = None
+        self._client_state_latest: dict[str, float] = {}
+        self._client_caps: dict | None = None
+
+        # Optional arm gravity compensation (feed-forward torque via the arm IK solver).
+        self.arm_ik = G1_29_ArmIK() if config.gravity_compensation else None

        # Initialize state variables
        self.sim_env = None
@@ -146,24 +172,69 @@ class UnitreeG1(Robot):
        self._shutdown_event = threading.Event()
        self.subscribe_thread = None

-        self.arm_ik = G1_29_ArmIK() if config.gravity_compensation else None
+        # Lower-body controller loaded dynamically. GUARDRAIL: the controller must never
+        # be built or run on the laptop client -- it always runs onboard (or in sim).
+        if self._client:
+            self.controller: LocomotionController | None = None
+        else:
+            self.controller = make_locomotion_controller(config.controller)

-        # Lower-body controller loaded dynamically
-        self.controller: LocomotionController | None = make_locomotion_controller(config.controller)
+            # 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
        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 = {}

+        # Onboard-only: parser for the physical Unitree wireless remote (read straight
+        # from local lowstate so joystick locomotion works without a laptop round-trip).
+        self._joystick = None
+
+        # 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. 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 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()

            # 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:
@@ -231,15 +302,46 @@ class UnitreeG1(Robot):
                features[f"{cam}_depth"] = (cfg.height, cfg.width, 1)
        return features

+    @property
+    def _token_state_ft(self) -> dict[str, type]:
+        """64-D SONIC latent-token proprio state (``motion_token_state.{i}.pos``).
+
+        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._sonic_token:
+            return {}
+        from .controllers.sonic_whole_body import TOKEN_DIM, token_state_key
+
+        return {token_state_key(i): float for i in range(TOKEN_DIM)}
+
    @cached_property
    def observation_features(self) -> dict[str, type | tuple]:
-        return {**self._motors_ft, **self._cameras_ft}
+        return {
+            **self._motors_ft,
+            **self._token_state_ft,
+            **self._cameras_ft,
+        }

    @cached_property
    def action_features(self) -> dict[str, type]:
-        if self.controller is None:
+        # 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:
            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._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)}
+
+        # Locomotion controllers (GR00T / Holosoma): arm joint targets + joystick axes.
        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}
@@ -255,6 +357,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

@@ -271,6 +378,13 @@ class UnitreeG1(Robot):
                with self._controller_action_lock:
                    controller_input = dict(self.controller_input)

+                # Onboard: the physical Unitree remote (in local lowstate) takes
+                # priority for locomotion when active; otherwise laptop/ZMQ axes stand.
+                if self.config.onboard:
+                    wl = self._wireless_remote_input(lowstate)
+                    if wl is not None:
+                        controller_input.update(wl)
+
                # Run controller step
                controller_action = self.controller.run_step(controller_input, lowstate)

@@ -293,7 +407,163 @@ class UnitreeG1(Robot):
    def configure(self) -> None:
        pass

+    def _wireless_remote_input(self, lowstate) -> dict | None:
+        """Parse the physical Unitree remote from lowstate into controller inputs.
+
+        Onboard only. Returns None when the remote is idle so the laptop-provided
+        (ZMQ) axes keep control; otherwise the physical remote takes priority.
+        """
+        js = self._joystick
+        if js is None:
+            return None
+        wr = getattr(lowstate, "wireless_remote", None)
+        if not wr or len(wr) < 24:
+            return None
+        try:
+            js.extract(wr)
+        except Exception:  # noqa: BLE001
+            return None
+
+        axes = {
+            "remote.lx": float(js.lx.data),
+            "remote.ly": float(js.ly.data),
+            "remote.rx": float(js.rx.data),
+            "remote.ry": float(js.ry.data),
+        }
+        active = any(abs(v) > 1e-2 for v in axes.values())
+        out = dict(axes)
+        for i, name in enumerate(_REMOTE_BUTTON_MAP):
+            if name:
+                val = float(getattr(js, name).data)
+                out[f"remote.button.{i}"] = val
+                if val:
+                    active = True
+        return out if active else None
+
+    def _release_motion_control(self) -> None:
+        """Release the robot's built-in motion services so we can send raw lowcmd.
+
+        Onboard-only. Mirrors run_g1_server.py: on the real robot the factory
+        locomotion/hand services must relinquish control before our controller can
+        write to ``rt/lowcmd``, otherwise commands are ignored or fought.
+        """
+        from unitree_sdk2py.comm.motion_switcher.motion_switcher_client import MotionSwitcherClient
+
+        msc = MotionSwitcherClient()
+        msc.SetTimeout(5.0)
+        msc.Init()
+        _, result = msc.CheckMode()
+        while result is not None and "name" in result and result["name"]:
+            logger.info("[UnitreeG1] Releasing built-in mode '%s'...", result["name"])
+            msc.ReleaseMode()
+            _, result = msc.CheckMode()
+            time.sleep(1.0)
+
+    # ------------------------------------------------------------------ #
+    # Thin-client role (laptop): no DDS, no controller. Talks to run_g1_server
+    # over ZMQ. The controller ALWAYS runs onboard; we only relay high-level
+    # actions and read back the state echo + camera frames.
+    # ------------------------------------------------------------------ #
+    def _connect_client(self) -> None:
+        import zmq
+
+        from .run_g1_server import ACTION_PORT, HANDSHAKE_PORT, STATE_PORT, request_controller
+
+        server_ip = self.config.robot_ip
+        if not server_ip:
+            raise ValueError("client mode requires config.robot_ip (the G1 running run_g1_server)")
+
+        # 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._sonic_token,
+        )
+        self._client_caps = request_controller(
+            server_ip,
+            self.config.controller,
+            sonic_token_action=self._sonic_token,
+            port=HANDSHAKE_PORT,
+        )
+        logger.info("[client] server agreed: %s", self._client_caps)
+
+        ctx = zmq.Context.instance()
+
+        # 2) Action PUSH: ship compact high-level actions to the onboard controller.
+        self._client_action_sock = ctx.socket(zmq.PUSH)
+        self._client_action_sock.setsockopt(zmq.SNDHWM, 2)
+        self._client_action_sock.setsockopt(zmq.LINGER, 0)
+        self._client_action_sock.connect(f"tcp://{server_ip}:{ACTION_PORT}")
+
+        # 3) State SUB: read the onboard observation.state echo (last token / joints).
+        self._client_state_sock = ctx.socket(zmq.SUB)
+        self._client_state_sock.setsockopt(zmq.CONFLATE, 1)
+        self._client_state_sock.setsockopt_string(zmq.SUBSCRIBE, "")
+        self._client_state_sock.connect(f"tcp://{server_ip}:{STATE_PORT}")
+
+        # 4) Cameras (ZMQ ImageServer served by run_g1_server) - same as any client.
+        for cam in self._cameras.values():
+            if not cam.is_connected:
+                cam.connect()
+        logger.info("[client] connected: actions ->:%d, state <-:%d, %d camera(s).",
+                    ACTION_PORT, STATE_PORT, len(self._cameras))
+
+    def _recv_client_state(self) -> None:
+        """Drain the state SUB (CONFLATE keeps only the freshest) into the latest cache."""
+        import zmq
+
+        if self._client_state_sock is None:
+            return
+        while True:
+            try:
+                state = self._client_state_sock.recv_json(flags=zmq.NOBLOCK)
+            except zmq.Again:
+                break
+            except (ValueError, zmq.ZMQError):
+                break
+            if isinstance(state, dict):
+                self._client_state_latest = {k: float(v) for k, v in state.items()}
+
+    def _get_observation_client(self) -> RobotObservation:
+        self._recv_client_state()
+        obs: dict = dict(self._client_state_latest)
+        for cam_name, cam in self._cameras.items():
+            if getattr(cam, "use_rgb", True):
+                obs[cam_name] = cam.read_latest()
+            if getattr(cam, "use_depth", False):
+                obs[f"{cam_name}_depth"] = cam.read_latest_depth()
+        return obs
+
+    def _send_action_client(self, action: RobotAction) -> RobotAction:
+        """Relay the raw action straight to the onboard controller. NO processing here:
+        the controller negotiated in the handshake interprets it (token / wb / arm)."""
+        import zmq
+
+        if self._client_action_sock is None:
+            raise DeviceNotConnectedError("UnitreeG1 client is not connected")
+        payload = json.dumps({k: float(v) for k, v in action.items()}).encode("utf-8")
+        with contextlib.suppress(zmq.Again):
+            self._client_action_sock.send(payload, zmq.NOBLOCK)
+        return action
+
+    def _disconnect_client(self) -> None:
+        for sock in (self._client_action_sock, self._client_state_sock):
+            if sock is not None:
+                with contextlib.suppress(Exception):
+                    sock.close(linger=0)
+        self._client_action_sock = None
+        self._client_state_sock = None
+        for cam in self._cameras.values():
+            with contextlib.suppress(Exception):
+                cam.disconnect()
+
    def connect(self, calibrate: bool = True) -> None:  # connect to DDS
+        # Thin-client role: no DDS, no controller. Negotiate the controller with
+        # run_g1_server (which runs it onboard), then open the high-level ZMQ links:
+        # PUSH actions on :ACTION_PORT, SUB state echo on :STATE_PORT, cameras via ZMQ.
+        if self._client:
+            self._connect_client()
+            return
+
        # Initialize DDS channel and simulation environment
        if self.config.is_simulation:
            from lerobot.envs import make_env
@@ -302,6 +572,28 @@ class UnitreeG1(Robot):
            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:
+            # Real robot, controller running onboard against local DDS. Initialize the
+            # real SDK channel factory on the robot's DDS interface and take low-level
+            # control from the built-in services before we start writing lowcmd.
+            if self.config.dds_interface:
+                self._ChannelFactoryInitialize(0, self.config.dds_interface)
+            else:
+                self._ChannelFactoryInitialize(0)
+            # Real robot: hand low-level control over from the built-in services.
+            # A DDS sim has no MotionSwitcher, so this is skipped there.
+            if self.config.release_motion_control:
+                self._release_motion_control()
+            # Real robot: read the physical wireless remote from lowstate for
+            # locomotion. A sim has no physical remote, so leave _joystick=None and
+            # let send_action (ZMQ) drive the locomotion axes instead.
+            if self.config.physical_remote:
+                from unitree_sdk2py.utils.joystick import Joystick
+
+                self._joystick = Joystick()
+                for axis in (self._joystick.lx, self._joystick.ly, self._joystick.rx, self._joystick.ry):
+                    axis.smooth = 1.0
+                    axis.deadzone = 0.0
        else:
            self._ChannelFactoryInitialize(0, config=self.config)

@@ -343,6 +635,9 @@ class UnitreeG1(Robot):

        self.kp = np.array(self.config.kp, dtype=np.float32)
        self.kd = np.array(self.config.kd, dtype=np.float32)
+        if self.controller is not None and hasattr(self.controller, "kp"):
+            self.kp = np.array(self.controller.kp, dtype=np.float32)
+            self.kd = np.array(self.controller.kd, dtype=np.float32)

        for joint in G1_29_JointIndex:
            self.msg.motor_cmd[joint].mode = 1
@@ -350,7 +645,8 @@ 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
+        # Start the 50 Hz controller thread (runs the locomotion/whole-body policy and
+        # publishes low commands to DDS).
        if self.controller is not None:
            self._controller_thread = threading.Thread(target=self._controller_loop, daemon=True)
            self._controller_thread.start()
@@ -372,12 +668,34 @@ class UnitreeG1(Robot):
            logger.warning(f"Failed to send zero-torque on disconnect: {e}")

    def disconnect(self):
-        # Put robot in passive mode before stopping threads
-        if not self.config.is_simulation:
+        if self._client:
+            self._disconnect_client()
+            return
+
+        # 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:
+            # 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():
+                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)"
+                )
+
+        # Put the robot in passive mode (zero-torque) before stopping the rest (real
+        # robot only; the subscribe thread is still alive here to supply the current
+        # pose). Only publish once the controller thread has definitely exited so the
+        # two aren't publishing at once.
+        if not self.config.is_simulation and controller_stopped:
            self._send_zero_torque()

-        # Signal thread to stop and unblock any waits
-        self._shutdown_event.set()
+        if self.controller is not None and hasattr(self.controller, "shutdown"):
+            self.controller.shutdown()

        # Wait for subscribe thread to finish
        if self.subscribe_thread is not None:
@@ -385,12 +703,6 @@ class UnitreeG1(Robot):
            if self.subscribe_thread.is_alive():
                logger.warning("Subscribe thread did not stop cleanly")

-        # Wait for controller thread to finish
-        if self._controller_thread is not None:
-            self._controller_thread.join(timeout=2.0)
-            if self._controller_thread.is_alive():
-                logger.warning("Controller thread did not stop cleanly")
-
        # Close simulation environment
        if self.config.is_simulation and self.sim_env is not None:
            try:
@@ -417,49 +729,25 @@ class UnitreeG1(Robot):
            cam.disconnect()

    def get_observation(self) -> RobotObservation:
+        if self._client:
+            return self._get_observation_client()
+
        with self._lowstate_lock:
            lowstate = self._lowstate
        if lowstate is None:
            return {}

-        obs = {}
+        # Motors + IMU + wireless remote (shared lowstate -> obs mapping)
+        obs = lowstate_to_obs(lowstate)

-        # Motors - q, dq, tau for all joints
-        for motor in G1_29_JointIndex:
-            name = motor.name
-            idx = motor.value
-            obs[f"{name}.q"] = lowstate.motor_state[idx].q
-            obs[f"{name}.dq"] = lowstate.motor_state[idx].dq
-            obs[f"{name}.tau"] = lowstate.motor_state[idx].tau_est
+        # 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._sonic_token:
+            from .controllers.sonic_whole_body import token_state_key

-        # IMU - gyroscope
-        if lowstate.imu_state.gyroscope:
-            obs["imu.gyro.x"] = lowstate.imu_state.gyroscope[0]
-            obs["imu.gyro.y"] = lowstate.imu_state.gyroscope[1]
-            obs["imu.gyro.z"] = lowstate.imu_state.gyroscope[2]
-
-        # IMU - accelerometer
-        if lowstate.imu_state.accelerometer:
-            obs["imu.accel.x"] = lowstate.imu_state.accelerometer[0]
-            obs["imu.accel.y"] = lowstate.imu_state.accelerometer[1]
-            obs["imu.accel.z"] = lowstate.imu_state.accelerometer[2]
-
-        # IMU - quaternion
-        if lowstate.imu_state.quaternion:
-            obs["imu.quat.w"] = lowstate.imu_state.quaternion[0]
-            obs["imu.quat.x"] = lowstate.imu_state.quaternion[1]
-            obs["imu.quat.y"] = lowstate.imu_state.quaternion[2]
-            obs["imu.quat.z"] = lowstate.imu_state.quaternion[3]
-
-        # IMU - rpy
-        if lowstate.imu_state.rpy:
-            obs["imu.rpy.roll"] = lowstate.imu_state.rpy[0]
-            obs["imu.rpy.pitch"] = lowstate.imu_state.rpy[1]
-            obs["imu.rpy.yaw"] = lowstate.imu_state.rpy[2]
-
-        # Wireless remote (raw bytes for teleoperator)
-        if lowstate.wireless_remote:
-            obs["wireless_remote"] = lowstate.wireless_remote
+            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)

        # Cameras - read images from ZMQ cameras
        for cam_name, cam in self._cameras.items():
@@ -471,11 +759,22 @@ class UnitreeG1(Robot):
        return obs

    def send_action(self, action: RobotAction) -> RobotAction:
+        if self._client:
+            return self._send_action_client(action)
+
        action_to_publish = action
        if self.controller is not None:
+            if self._sonic_token:
+                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 getattr(self.controller, "full_body", False):
+                return action
            # Controller thread owns legs/waist. Here we only update joystick inputs
            # and publish arm targets from the teleoperator.
-            self._update_controller_action(action)
            arm_prefixes = tuple(j.name for j in G1_29_JointArmIndex)
            action_to_publish = {
                key: value
@@ -503,11 +802,17 @@ class UnitreeG1(Robot):
        return action

    def _update_controller_action(self, action: RobotAction) -> None:
-        """Update controller input state from incoming teleop action."""
+        """Update controller input state from an incoming teleop action.
+
+        Controller-agnostic: every value-carrying key (e.g. locomotion ``remote.*``
+        axes/buttons) is forwarded verbatim into ``controller_input`` and each
+        controller extracts only the keys it understands. The robot deliberately does
+        not enumerate any controller's key schema here.
+        """
        with self._controller_action_lock:
-            for key in REMOTE_KEYS:
-                if key in action:
-                    self.controller_input[key] = action[key]
+            for key, value in action.items():
+                if isinstance(key, str) and value is not None:
+                    self.controller_input[key] = value

    @property
    def is_calibrated(self) -> bool:
@@ -515,6 +820,8 @@ class UnitreeG1(Robot):

    @property
    def is_connected(self) -> bool:
+        if self._client:
+            return self._client_action_sock is not None
        with self._lowstate_lock:
            return self._lowstate is not None

@@ -537,43 +844,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 ``<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

-            # 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")
@@ -36,6 +36,7 @@ python src/lerobot/scripts/augment_dataset_quantile_stats.py \
 import argparse
 import concurrent.futures
 import logging
+import os
 from pathlib import Path

 import numpy as np
@@ -52,6 +53,7 @@ from lerobot.datasets import (
    get_feature_stats,
    write_stats,
 )
+from lerobot.datasets.compute_stats import sample_indices
 from lerobot.utils.utils import init_logging


@@ -77,12 +79,14 @@ def has_quantile_stats(stats: dict[str, dict] | None, quantile_list_keys: list[s
    return False


-def process_single_episode(dataset: LeRobotDataset, episode_idx: int) -> dict:
+def process_single_episode(dataset: LeRobotDataset, episode_idx: int, use_sampling: bool = True) -> dict:
    """Process a single episode and return its statistics.

    Args:
        dataset: The LeRobot dataset
        episode_idx: Index of the episode to process
+        use_sampling: If True, sub-sample image/video frames per episode to bound
+            memory. If False, use every frame (exact, higher memory).

    Returns:
        Dictionary containing episode statistics
@@ -92,16 +96,31 @@ def process_single_episode(dataset: LeRobotDataset, episode_idx: int) -> dict:
    start_idx = dataset.meta.episodes[episode_idx]["dataset_from_index"]
    end_idx = dataset.meta.episodes[episode_idx]["dataset_to_index"]

-    collected_data: dict[str, list] = {}
-    for idx in range(start_idx, end_idx):
-        item = dataset[idx]
-        for key, value in item.items():
-            if key not in dataset.features:
-                continue
+    episode_len = end_idx - start_idx

-            if key not in collected_data:
-                collected_data[key] = []
-            collected_data[key].append(value)
+    # Images/video are the memory hog, so sub-sample those frames per episode;
+    # numeric columns are cheap, so read them in full (exact).
+    image_keys = [k for k in dataset.features if dataset.features[k]["dtype"] in ("image", "video")]
+    numeric_keys = [
+        k for k in dataset.features if dataset.features[k]["dtype"] not in ("image", "video", "string")
+    ]
+
+    collected_data: dict[str, list] = {}
+
+    # Numeric features: every frame, read directly from the underlying table.
+    if numeric_keys:
+        numeric_cols = dataset.hf_dataset.select_columns(numeric_keys)[start_idx:end_idx]
+        for key in numeric_keys:
+            collected_data[key] = [torch.as_tensor(v) for v in numeric_cols[key]]
+
+    # Image/video features: decode only a sampled subset of frames.
+    if image_keys:
+        sampled_offsets = sample_indices(episode_len) if use_sampling else list(range(episode_len))
+        for offset in sampled_offsets:
+            item = dataset[start_idx + offset]
+            for key in image_keys:
+                if key in item:
+                    collected_data.setdefault(key, []).append(item[key])

    ep_stats = {}
    for key, data_list in collected_data.items():
@@ -131,11 +150,13 @@ def process_single_episode(dataset: LeRobotDataset, episode_idx: int) -> dict:
    return ep_stats


-def compute_quantile_stats_for_dataset(dataset: LeRobotDataset) -> dict[str, dict]:
+def compute_quantile_stats_for_dataset(dataset: LeRobotDataset, use_sampling: bool = True) -> dict[str, dict]:
    """Compute quantile statistics for all episodes in the dataset.

    Args:
        dataset: The LeRobot dataset to compute statistics for
+        use_sampling: If True, sub-sample image/video frames per episode to bound
+            memory. If False, use every frame (exact, higher memory).

    Returns:
        Dictionary containing aggregated statistics with quantiles
@@ -153,15 +174,15 @@ def compute_quantile_stats_for_dataset(dataset: LeRobotDataset) -> dict[str, dic
    if has_videos:
        logging.info("Dataset contains video keys - using sequential processing for thread safety")
        for episode_idx in tqdm(range(dataset.num_episodes), desc="Processing episodes"):
-            ep_stats = process_single_episode(dataset, episode_idx)
+            ep_stats = process_single_episode(dataset, episode_idx, use_sampling)
            episode_stats_list.append(ep_stats)
    else:
        logging.info("Dataset has no video keys - using parallel processing for better performance")
-        max_workers = min(dataset.num_episodes, 16)
+        max_workers = min(dataset.num_episodes, int(os.environ.get("LEROBOT_STATS_MAX_WORKERS", 16)))

        with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
            future_to_episode = {
-                executor.submit(process_single_episode, dataset, episode_idx): episode_idx
+                executor.submit(process_single_episode, dataset, episode_idx, use_sampling): episode_idx
                for episode_idx in range(dataset.num_episodes)
            }

@@ -188,6 +209,7 @@ def augment_dataset_with_quantile_stats(
    repo_id: str,
    root: str | Path | None = None,
    overwrite: bool = False,
+    use_sampling: bool = True,
 ) -> None:
    """Augment a dataset with quantile statistics if they are missing.

@@ -195,6 +217,8 @@ def augment_dataset_with_quantile_stats(
        repo_id: Repository ID of the dataset
        root: Local root directory for the dataset
        overwrite: Overwrite existing quantile statistics if they already exist
+        use_sampling: If True, sub-sample image/video frames per episode to bound
+            memory. If False, use every frame (exact, higher memory).
    """
    logging.info(f"Loading dataset: {repo_id}")
    dataset = LeRobotDataset(
@@ -208,7 +232,7 @@ def augment_dataset_with_quantile_stats(

    logging.info("Dataset does not contain quantile statistics. Computing them now...")

-    new_stats = compute_quantile_stats_for_dataset(dataset)
+    new_stats = compute_quantile_stats_for_dataset(dataset, use_sampling=use_sampling)

    logging.info("Updating dataset metadata with new quantile statistics")
    dataset.meta.stats = new_stats
@@ -248,6 +272,14 @@ def main():
        action="store_true",
        help="Overwrite existing quantile statistics if they already exist",
    )
+    parser.add_argument(
+        "--no-sampling",
+        action="store_true",
+        help=(
+            "Compute stats over every frame (exact, higher memory). By default, "
+            "image/video frames are sub-sampled per episode to bound memory."
+        ),
+    )

    args = parser.parse_args()
    root = Path(args.root) if args.root else None
@@ -258,6 +290,7 @@ def main():
        repo_id=args.repo_id,
        root=root,
        overwrite=args.overwrite,
+        use_sampling=not args.no_sampling,
    )


@@ -564,7 +564,7 @@ def eval_policy(
        if seeds:
            all_seeds.extend(seeds)
        else:
-            all_seeds.append(None)
+            all_seeds.extend([None] * env.num_envs)

        # FIXME: episode_data is either None or it doesn't exist
        if return_episode_data:
@@ -28,7 +28,6 @@ lerobot-find-cameras
 # NOTE(Steven): macOS cameras sometimes report different FPS at init time, not an issue here as we don't specify FPS when opening the cameras, but the information displayed might not be truthful.

 import argparse
-import concurrent.futures
 import logging
 import time
 from pathlib import Path
@@ -133,7 +132,7 @@ def save_image(
    camera_identifier: str | int,
    images_dir: Path,
    camera_type: str,
-):
+) -> None:
    """
    Saves a single image to disk using Pillow. Handles color conversion if necessary.
    """
@@ -152,7 +151,7 @@ def save_image(
        logger.error(f"Failed to save image for camera {camera_identifier} (type {camera_type}): {e}")


-def create_camera_instance(cam_meta: dict[str, Any]) -> dict[str, Any] | None:
+def create_camera_instance(cam_meta: dict[str, Any], *, warmup_s: int = 1) -> dict[str, Any] | None:
    """Create and connect to a camera instance based on metadata."""
    cam_type = cam_meta.get("type")
    cam_id = cam_meta.get("id")
@@ -165,12 +164,14 @@ def create_camera_instance(cam_meta: dict[str, Any]) -> dict[str, Any] | None:
            cv_config = OpenCVCameraConfig(
                index_or_path=cam_id,
                color_mode=ColorMode.RGB,
+                warmup_s=warmup_s,
            )
            instance = OpenCVCamera(cv_config)
        elif cam_type == "RealSense":
            rs_config = RealSenseCameraConfig(
                serial_number_or_name=cam_id,
                color_mode=ColorMode.RGB,
+                warmup_s=warmup_s,
            )
            instance = RealSenseCamera(rs_config)
        else:
@@ -188,9 +189,7 @@ def create_camera_instance(cam_meta: dict[str, Any]) -> dict[str, Any] | None:
        return None


-def process_camera_image(
-    cam_dict: dict[str, Any], output_dir: Path, current_time: float
-) -> concurrent.futures.Future | None:
+def process_camera_image(cam_dict: dict[str, Any], output_dir: Path, current_time: float) -> None:
    """Capture and process an image from a single camera."""
    cam = cam_dict["instance"]
    meta = cam_dict["meta"]
@@ -200,7 +199,7 @@ def process_camera_image(
    try:
        image_data = cam.read()

-        return save_image(
+        save_image(
            image_data,
            cam_id_str,
            output_dir,
@@ -215,21 +214,21 @@ def process_camera_image(
    return None


-def cleanup_cameras(cameras_to_use: list[dict[str, Any]]):
+def cleanup_camera(cam_dict: dict[str, Any]) -> None:
    """Disconnect all cameras."""
-    logger.info(f"Disconnecting {len(cameras_to_use)} cameras...")
-    for cam_dict in cameras_to_use:
-        try:
-            if cam_dict["instance"] and cam_dict["instance"].is_connected:
-                cam_dict["instance"].disconnect()
-        except Exception as e:
-            logger.error(f"Error disconnecting camera {cam_dict['meta'].get('id')}: {e}")
+    logger.info(f"Disconnecting camera with ID {cam_dict['meta'].get('id')}...")
+    try:
+        if cam_dict["instance"] and cam_dict["instance"].is_connected:
+            cam_dict["instance"].disconnect()
+    except Exception as e:
+        logger.error(f"Error disconnecting camera {cam_dict['meta'].get('id')}: {e}")


 def save_images_from_all_cameras(
    output_dir: Path,
    record_time_s: float = 2.0,
    camera_type: str | None = None,
+    warmup_s: int = 1,
 ):
    """
    Connects to detected cameras (optionally filtered by type) and saves images from each.
@@ -240,6 +239,7 @@ def save_images_from_all_cameras(
        record_time_s: Duration in seconds to record images.
        camera_type: Optional string to filter cameras ("realsense" or "opencv").
                            If None, uses all detected cameras.
+        warmup_s: Duration in seconds to warmup camera before recording images.
    """
    output_dir.mkdir(parents=True, exist_ok=True)
    logger.info(f"Saving images to {output_dir}")
@@ -249,40 +249,24 @@ def save_images_from_all_cameras(
        logger.warning("No cameras detected matching the criteria. Cannot save images.")
        return

-    cameras_to_use = []
-    for cam_meta in all_camera_metadata:
-        camera_instance = create_camera_instance(cam_meta)
-        if camera_instance:
-            cameras_to_use.append(camera_instance)
+    logger.info(
+        f"Starting image capture for {record_time_s} seconds from {len(all_camera_metadata)} cameras."
+    )

-    if not cameras_to_use:
-        logger.warning("No cameras could be connected. Aborting image save.")
-        return
-
-    logger.info(f"Starting image capture for {record_time_s} seconds from {len(cameras_to_use)} cameras.")
-    start_time = time.perf_counter()
-
-    with concurrent.futures.ThreadPoolExecutor(max_workers=len(cameras_to_use) * 2) as executor:
-        try:
+    try:
+        for cam_meta in all_camera_metadata:
+            cam_dict = create_camera_instance(cam_meta, warmup_s=warmup_s)
+            if cam_dict is None:
+                continue
+            start_time = time.perf_counter()
            while time.perf_counter() - start_time < record_time_s:
-                futures = []
                current_capture_time = time.perf_counter()
-
-                for cam_dict in cameras_to_use:
-                    future = process_camera_image(cam_dict, output_dir, current_capture_time)
-                    if future:
-                        futures.append(future)
-
-                if futures:
-                    concurrent.futures.wait(futures)
-
-        except KeyboardInterrupt:
-            logger.info("Capture interrupted by user.")
-        finally:
-            print("\nFinalizing image saving...")
-            executor.shutdown(wait=True)
-            cleanup_cameras(cameras_to_use)
-            print(f"Image capture finished. Images saved to {output_dir}")
+                process_camera_image(cam_dict, output_dir, current_capture_time)
+            cleanup_camera(cam_dict)
+    except KeyboardInterrupt:
+        logger.info("Capture interrupted by user.")
+    finally:
+        print(f"Image capture finished. Images saved to {output_dir}")


 def main():
@@ -291,7 +275,6 @@ def main():
    parser = argparse.ArgumentParser(
        description="Unified camera utility script for listing cameras and capturing images."
    )
-
    parser.add_argument(
        "camera_type",
        type=str,
@@ -309,8 +292,14 @@ def main():
    parser.add_argument(
        "--record-time-s",
        type=float,
-        default=6.0,
-        help="Time duration to attempt capturing frames. Default: 6 seconds.",
+        default=2.0,
+        help="Time duration to attempt capturing frames. Default: 2 seconds.",
+    )
+    parser.add_argument(
+        "--warmup-s",
+        type=int,
+        default=1,
+        help="Time duration to warmup camera before attempting to capture frames. Default: 1 second.",
    )
    args = parser.parse_args()
    save_images_from_all_cameras(**vars(args))
@@ -22,7 +22,8 @@ import dataclasses
 import logging
 import sys
 import time
-from contextlib import nullcontext
+from collections.abc import Iterator
+from contextlib import contextmanager, nullcontext
 from pprint import pformat
 from typing import TYPE_CHECKING, Any

@@ -76,6 +77,20 @@ else:
 from .lerobot_eval import eval_policy_all


+@contextmanager
+def _make_eval_envs(cfg: TrainPipelineConfig) -> Iterator[dict[str, dict[int, Any]]]:
+    """Create evaluation environments for one run and always dispose of them."""
+    envs = make_env(
+        cfg.env,
+        n_envs=cfg.eval.batch_size,
+        use_async_envs=cfg.eval.use_async_envs,
+    )
+    try:
+        yield envs
+    finally:
+        close_envs(envs)
+
+
 def _dataloader_worker_kwargs(cfg: TrainPipelineConfig) -> dict[str, Any]:
    """Return worker-only DataLoader options, disabling them for single-process loading."""
    workers_enabled = cfg.num_workers > 0
@@ -280,14 +295,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
    if not is_main_process:
        dataset, eval_dataset = make_train_eval_datasets(cfg)

-    # Create environment used for evaluating checkpoints during training on simulation data.
-    # On real-world data, no need to create an environment as evaluations are done outside train.py,
-    # using the eval.py instead, with gym_dora environment and dora-rs.
-    eval_env = None
-    if cfg.env_eval_freq > 0 and cfg.env is not None and is_main_process:
-        logging.info("Creating env")
-        eval_env = make_env(cfg.env, n_envs=cfg.eval.batch_size, use_async_envs=cfg.eval.use_async_envs)
-
    if cfg.is_reward_model_training:
        if is_main_process:
            logging.info("Creating reward model")
@@ -695,7 +702,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
            if is_main_process:
                step_id = get_step_identifier(step, cfg.steps)
                logging.info(f"Eval policy at step {step}")
-                with torch.no_grad(), accelerator.autocast():
+                with _make_eval_envs(cfg) as eval_env, torch.no_grad(), accelerator.autocast():
                    eval_info = eval_policy_all(
                        envs=eval_env,  # dict[suite][task_id] -> vec_env
                        policy=accelerator.unwrap_model(policy),
@@ -743,9 +750,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
    if is_main_process:
        progbar.close()

-    if eval_env:
-        close_envs(eval_env)
-
    is_fsdp = accelerator.distributed_type == DistributedType.FSDP
    model_state_dict = accelerator.get_state_dict(policy) if is_fsdp else None
    if is_main_process:
@@ -171,7 +171,13 @@ class IOSPhone(BasePhone, Teleoperator):
        # HEBI provides orientation in w, x, y, z format.
        # Scipy's Rotation expects x, y, z, w.
        quat_xyzw = np.concatenate((ar_quat[1:], [ar_quat[0]]))  # wxyz to xyzw
-        rot = Rotation.from_quat(quat_xyzw)
+        # ARKit can emit zero/NaN quaternions before tracking is ready or on a
+        # dropped packet. Rotation.from_quat now rejects those; degrade the same
+        # way as a missing pose so teleop stays alive mid-session.
+        try:
+            rot = Rotation.from_quat(quat_xyzw)
+        except ValueError:
+            return False, None, None, None
        pos = ar_pos - rot.apply(self.config.camera_offset)
        return True, pos, rot, pose

@@ -29,6 +29,12 @@ class SOLeaderConfig:
    # Whether to use degrees for angles
    use_degrees: bool = True

+    # Number of extra attempts when a `sync_read` of the motors fails. Feetech buses can occasionally
+    # return a corrupted status packet ("Incorrect status packet!"), especially when several joints move
+    # at once, which otherwise aborts the teleoperation loop. Retries are immediate (no sleep) and only
+    # happen on failure, so the steady-state read cost is unchanged.
+    num_read_retries: int = 2
+

@TeleoperatorConfig.register_subclass("so101_leader")
@TeleoperatorConfig.register_subclass("so100_leader")
@@ -145,7 +145,7 @@ class SOLeader(Teleoperator):
    @check_if_not_connected
    def get_action(self) -> dict[str, float]:
        start = time.perf_counter()
-        action = self.bus.sync_read("Present_Position")
+        action = self.bus.sync_read("Present_Position", num_retry=self.config.num_read_retries)
        action = {f"{motor}.pos": val for motor, val in action.items()}
        dt_ms = (time.perf_counter() - start) * 1e3
        logger.debug(f"{self} read action: {dt_ms:.1f}ms")
@@ -41,7 +41,7 @@ class RandomSubsetApply(Transform):

    def __init__(
        self,
-        transforms: Sequence[Callable],
+        transforms: Sequence[Callable[..., Any]],
        p: list[float] | None = None,
        n_subset: int | None = None,
        random_order: bool = False,
@@ -50,7 +50,7 @@ class RandomSubsetApply(Transform):
        if not isinstance(transforms, Sequence):
            raise TypeError("Argument transforms should be a sequence of callables")
        if p is None:
-            p = [1] * len(transforms)
+            p = [1.0] * len(transforms)
        elif len(p) != len(transforms):
            raise ValueError(
                f"Length of p doesn't match the number of transforms: {len(p)} != {len(transforms)}"
@@ -69,7 +69,7 @@ class RandomSubsetApply(Transform):
        self.n_subset = n_subset
        self.random_order = random_order

-        self.selected_transforms = None
+        self.selected_transforms: list[Callable[..., Any]] = []

    def forward(self, *inputs: Any) -> Any:
        needs_unpacking = len(inputs) > 1
@@ -119,7 +119,7 @@ class SharpnessJitter(Transform):
        super().__init__()
        self.sharpness = self._check_input(sharpness)

-    def _check_input(self, sharpness):
+    def _check_input(self, sharpness: float | Sequence[float]) -> tuple[float, float]:
        if isinstance(sharpness, (int | float)):
            if sharpness < 0:
                raise ValueError("If sharpness is a single number, it must be non negative.")
@@ -215,7 +215,7 @@ class ImageTransformsConfig:
    )


-def make_transform_from_config(cfg: ImageTransformConfig):
+def make_transform_from_config(cfg: ImageTransformConfig) -> Transform:
    if cfg.type == "SharpnessJitter":
        return SharpnessJitter(**cfg.kwargs)

@@ -236,8 +236,8 @@ class ImageTransforms(Transform):
        super().__init__()
        self._cfg = cfg

-        self.weights = []
-        self.transforms = {}
+        self.weights: list[float] = []
+        self.transforms: dict[str, Transform] = {}
        for tf_name, tf_cfg in cfg.tfs.items():
            if tf_cfg.weight <= 0.0:
                continue
@@ -37,16 +37,25 @@ def auto_select_torch_device() -> torch.device:

 # TODO(Steven): Remove log. log shouldn't be an argument, this should be handled by the logger level
 def get_safe_torch_device(try_device: str, log: bool = False) -> torch.device:
-    """Given a string, return a torch.device with checks on whether the device is available."""
+    """Given a string, return a torch.device with checks on whether the device is available.
+
+    Raises:
+        ValueError: If the requested device family is known but not available on
+            this machine (``AssertionError`` was previously used and is easy to
+            mistake for a programmer bug under ``python -O`` where asserts vanish).
+    """
    try_device = str(try_device)
    if try_device.startswith("cuda"):
-        assert torch.cuda.is_available()
+        if not torch.cuda.is_available():
+            raise ValueError(f"Requested device {try_device!r} but CUDA is not available.")
        device = torch.device(try_device)
    elif try_device == "mps":
-        assert torch.backends.mps.is_available()
+        if not torch.backends.mps.is_available():
+            raise ValueError("Requested device 'mps' but MPS is not available.")
        device = torch.device("mps")
    elif try_device == "xpu":
-        assert torch.xpu.is_available()
+        if not torch.xpu.is_available():
+            raise ValueError("Requested device 'xpu' but XPU is not available.")
        device = torch.device("xpu")
    elif try_device == "cpu":
        device = torch.device("cpu")
@@ -32,21 +32,21 @@ def load_json(fpath: Path) -> Any:
    Returns:
        Any: The data loaded from the JSON file.
    """
-    with open(fpath) as f:
+    with open(fpath, encoding="utf-8") as f:
        return json.load(f)


-def write_json(data: dict, fpath: Path) -> None:
-    """Write data to a JSON file.
+def write_json(data: JsonLike, fpath: Path) -> None:
+    """Write JSON-serializable data to a file.

    Creates parent directories if they don't exist.

    Args:
-        data (dict): The dictionary to write.
+        data: JSON-serializable data to write.
        fpath (Path): The path to the output JSON file.
    """
    fpath.parent.mkdir(exist_ok=True, parents=True)
-    with open(fpath, "w") as f:
+    with open(fpath, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=4, ensure_ascii=False)


@@ -30,6 +30,10 @@ def precise_sleep(seconds: float, spin_threshold: float = 0.010, sleep_margin: f
    """
    if seconds <= 0:
        return
+    if spin_threshold < 0:
+        raise ValueError(f"spin_threshold must be >= 0, got {spin_threshold}")
+    if sleep_margin < 0:
+        raise ValueError(f"sleep_margin must be >= 0, got {sleep_margin}")

    system = platform.system()
    # On macOS and Windows the scheduler / sleep granularity can make
@@ -29,10 +29,13 @@ class Rotation:
    def __init__(self, quat: np.ndarray) -> None:
        """Initialize rotation from quaternion [x, y, z, w]."""
        self._quat = np.asarray(quat, dtype=float)
-        # Normalize quaternion
+        if self._quat.shape != (4,):
+            raise ValueError(f"Quaternion must have shape (4,), got {self._quat.shape}")
+        # Normalize quaternion. Reject the zero vector — it has no orientation.
        norm = np.linalg.norm(self._quat)
-        if norm > 0:
-            self._quat = self._quat / norm
+        if norm <= 0.0 or not np.isfinite(norm):
+            raise ValueError(f"Quaternion must be a non-zero finite vector; got {self._quat} (norm={norm})")
+        self._quat = self._quat / norm

    @classmethod
    def from_rotvec(cls, rotvec: np.ndarray) -> "Rotation":
@@ -14,7 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.

-from typing import TypedDict
+from typing import NotRequired, TypedDict

 import torch

@@ -28,7 +28,7 @@ class Transition(TypedDict):
    next_state: dict[str, torch.Tensor]
    done: bool
    truncated: bool
-    complementary_info: dict[str, torch.Tensor | float | int] | None = None
+    complementary_info: NotRequired[dict[str, torch.Tensor | float | int] | None]


 def move_transition_to_device(transition: Transition, device: str = "cpu") -> Transition:
@@ -24,7 +24,6 @@ import sys
 import time
 from collections.abc import Iterator
 from copy import copy, deepcopy
-from datetime import datetime
 from pathlib import Path
 from statistics import mean
 from typing import TYPE_CHECKING, Any
@@ -61,14 +60,16 @@ def init_logging(
        accelerator: Optional Accelerator instance (for multi-GPU detection)
    """

-    def custom_format(record: logging.LogRecord) -> str:
-        dt = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-        fnameline = f"{record.pathname}:{record.lineno}"
-        pid_str = f"[PID: {os.getpid()}] " if display_pid else ""
-        return f"{record.levelname} {pid_str}{dt} {fnameline[-15:]:>15} {record.getMessage()}"
+    class LeRobotFormatter(logging.Formatter):
+        def format(self, record: logging.LogRecord) -> str:
+            record.lerobot_location = f"{record.pathname}:{record.lineno}"[-15:]
+            record.lerobot_pid = f"[PID: {os.getpid()}] " if display_pid else ""
+            return super().format(record)

-    formatter = logging.Formatter()
-    formatter.format = custom_format
+    formatter = LeRobotFormatter(
+        "%(levelname)s %(lerobot_pid)s%(asctime)s %(lerobot_location)15s %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )

    logger = logging.getLogger()
    logger.setLevel(logging.NOTSET)
@@ -133,10 +134,13 @@ def say(text: str, blocking: bool = False):
    else:
        raise RuntimeError("Unsupported operating system for text-to-speech.")

-    if blocking:
-        subprocess.run(cmd, check=True)
-    else:
-        subprocess.Popen(cmd, creationflags=subprocess.CREATE_NO_WINDOW if system == "Windows" else 0)
+    try:
+        if blocking:
+            subprocess.run(cmd, check=True, timeout=5)
+        else:
+            subprocess.Popen(cmd, creationflags=subprocess.CREATE_NO_WINDOW if system == "Windows" else 0)
+    except (FileNotFoundError, subprocess.TimeoutExpired) as e:
+        logging.warning("Text-to-speech command failed: %s | Error: %s", cmd, e)


 def log_say(text: str, play_sounds: bool = True, blocking: bool = False):
@@ -0,0 +1,104 @@
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import pytest
+
+pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
+
+from lerobot.scripts.augment_dataset_quantile_stats import (
+    compute_quantile_stats_for_dataset,
+    has_quantile_stats,
+)
+
+
+def _numeric_keys(dataset):
+    return [k for k, v in dataset.features.items() if v["dtype"] not in ("image", "video", "string")]
+
+
+def _image_keys(dataset):
+    return [k for k, v in dataset.features.items() if v["dtype"] in ("image", "video")]
+
+
+def test_numeric_stats_are_unaffected_by_sampling(tmp_path, lerobot_dataset_factory):
+    """Sampling only touches image/video frames; numeric features are read in
+    full either way, so their stats must be identical with and without sampling."""
+    dataset = lerobot_dataset_factory(
+        root=tmp_path / "ds", total_episodes=2, total_frames=400, use_videos=False
+    )
+
+    exact = compute_quantile_stats_for_dataset(dataset, use_sampling=False)
+    sampled = compute_quantile_stats_for_dataset(dataset, use_sampling=True)
+
+    numeric_keys = _numeric_keys(dataset)
+    assert numeric_keys, "fixture should expose numeric features"
+    for key in numeric_keys:
+        if key not in exact:
+            continue
+        for stat in ("mean", "std", "q01", "q50", "q99"):
+            if stat in exact[key]:
+                np.testing.assert_allclose(
+                    sampled[key][stat],
+                    exact[key][stat],
+                    rtol=1e-6,
+                    atol=1e-6,
+                    err_msg=f"numeric feature '{key}' stat '{stat}' changed under sampling",
+                )
+
+
+def test_image_sampling_reduces_data_but_keeps_stats_close(tmp_path, lerobot_dataset_factory):
+    """For images, sampling should reduce the number of samples considered while
+    keeping the resulting statistics close to the exact ones."""
+    dataset = lerobot_dataset_factory(
+        root=tmp_path / "ds", total_episodes=2, total_frames=400, use_videos=False
+    )
+
+    exact = compute_quantile_stats_for_dataset(dataset, use_sampling=False)
+    sampled = compute_quantile_stats_for_dataset(dataset, use_sampling=True)
+
+    image_keys = _image_keys(dataset)
+    assert image_keys, "fixture should expose at least one image feature"
+    for key in image_keys:
+        # sampling actually looked at fewer pixels
+        assert sampled[key]["count"][0] < exact[key]["count"][0]
+        # but per-channel mean stays close
+        np.testing.assert_allclose(
+            sampled[key]["mean"],
+            exact[key]["mean"],
+            rtol=0.15,
+            err_msg=f"image feature '{key}' mean drifted too far under sampling",
+        )
+
+
+def test_short_episodes_use_all_frames(tmp_path, lerobot_dataset_factory):
+    """With episodes shorter than the sampling floor, sampling is a no-op and
+    must produce exactly the same stats as the exact path."""
+    dataset = lerobot_dataset_factory(
+        root=tmp_path / "ds", total_episodes=2, total_frames=40, use_videos=False
+    )
+
+    exact = compute_quantile_stats_for_dataset(dataset, use_sampling=False)
+    sampled = compute_quantile_stats_for_dataset(dataset, use_sampling=True)
+
+    for key in _image_keys(dataset):
+        assert sampled[key]["count"][0] == exact[key]["count"][0]
+
+
+def test_quantile_stats_present_after_compute(tmp_path, lerobot_dataset_factory):
+    """The computed stats should contain quantile keys for the dataset."""
+    dataset = lerobot_dataset_factory(
+        root=tmp_path / "ds", total_episodes=2, total_frames=200, use_videos=False
+    )
+    stats = compute_quantile_stats_for_dataset(dataset, use_sampling=True)
+    assert has_quantile_stats(stats)
@@ -482,6 +482,20 @@ def test_add_frame_works_in_write_mode(tmp_path):
 # ── Resume mode ──────────────────────────────────────────────────────


+def test_resume_freshly_created_empty_dataset(tmp_path):
+    """resume() accepts a local dataset created before any episode was recorded."""
+    root = tmp_path / "resume_empty_ds"
+    LeRobotDataset.create(repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root)
+
+    resumed = LeRobotDataset.resume(repo_id=DUMMY_REPO_ID, root=root)
+
+    assert isinstance(resumed.writer, DatasetWriter)
+    assert resumed.meta.total_episodes == 0
+    assert resumed.meta.total_frames == 0
+    assert resumed.meta.tasks is None
+    assert resumed.meta.episodes is None
+
+
 def test_resume_creates_writer(tmp_path):
    """After resume(), writer is a DatasetWriter."""
    root = tmp_path / "resume_ds"
@@ -294,6 +294,19 @@ def test__sync_read(addr, length, ids_values, mock_motors, dummy_motors):
    assert read_values == ids_values


+def test__sync_read_retries_after_transient_failure(mock_motors, dummy_motors):
+    addr, length, ids_values = (10, 4, {1: 1337})
+    stub = mock_motors.build_sync_read_stub(addr, length, ids_values, num_invalid_try=1)
+    bus = FeetechMotorsBus(port=mock_motors.port, motors=dummy_motors)
+    bus.connect(handshake=False)
+
+    read_values, read_comm = bus._sync_read(addr, length, list(ids_values), num_retry=1)
+
+    assert read_comm == scs.COMM_SUCCESS
+    assert read_values == ids_values
+    assert mock_motors.stubs[stub].calls == 2
+
+
@pytest.mark.parametrize("raise_on_error", (True, False))
 def test__sync_read_comm(raise_on_error, mock_motors, dummy_motors):
    addr, length, ids_values = (10, 4, {1: 1337})
@@ -496,6 +496,60 @@ def test_evo1_processor_save_load_round_trip_applies_config_overrides(tmp_path):
    assert "embodiment_id" in processed


+def test_reconcile_evo1_processors_repads_overridden_stats(tmp_path):
+    """Loading a checkpoint and injecting raw (unpadded) dataset stats must be re-padded.
+
+    Regression test: lerobot-train passes the raw dataset stats as normalizer/unnormalizer
+    overrides when resuming from a checkpoint (e.g. stage2 from a stage1 checkpoint). Those stats
+    are at the dataset dims (e.g. LIBERO state=8/action=7), but EVO1 pads state/action to
+    max_state_dim/max_action_dim before normalization, so reconcile_evo1_processors must re-pad the
+    stats or normalization crashes with a shape mismatch.
+    """
+    config = make_config()
+    preprocessor, postprocessor = make_evo1_pre_post_processors(config, dataset_stats=make_stats())
+    preprocessor.save_pretrained(tmp_path)
+    postprocessor.save_pretrained(tmp_path)
+
+    # Reload with the generic override path injecting raw, unpadded dataset stats.
+    raw_stats = make_stats()
+    loaded_pre = PolicyProcessorPipeline.from_pretrained(
+        tmp_path,
+        config_filename=f"{POLICY_PREPROCESSOR_DEFAULT_NAME}.json",
+        overrides={"normalizer_processor": {"stats": raw_stats}},
+        to_transition=batch_to_transition,
+        to_output=transition_to_batch,
+    )
+    loaded_post = PolicyProcessorPipeline.from_pretrained(
+        tmp_path,
+        config_filename=f"{POLICY_POSTPROCESSOR_DEFAULT_NAME}.json",
+        overrides={"unnormalizer_processor": {"stats": raw_stats}},
+        to_transition=policy_action_to_transition,
+        to_output=transition_to_policy_action,
+    )
+
+    # Sanity: the override really injected unpadded stats before reconciliation.
+    normalizer = next(step for step in loaded_pre.steps if isinstance(step, NormalizerProcessorStep))
+    assert normalizer._tensor_stats[OBS_STATE]["min"].shape == (STATE_DIM,)
+
+    loaded_pre, loaded_post = reconcile_evo1_processors(config, loaded_pre, loaded_post)
+
+    normalizer = next(step for step in loaded_pre.steps if isinstance(step, NormalizerProcessorStep))
+    unnormalizer = next(step for step in loaded_post.steps if isinstance(step, UnnormalizerProcessorStep))
+    assert normalizer._tensor_stats[OBS_STATE]["min"].shape == (MAX_STATE_DIM,)
+    assert normalizer._tensor_stats[ACTION]["min"].shape == (MAX_ACTION_DIM,)
+    assert unnormalizer._tensor_stats[ACTION]["min"].shape == (MAX_ACTION_DIM,)
+
+    # Normalizing a padded state must not raise (this is the exact runtime path that crashed).
+    processed = loaded_pre(
+        {
+            "task": "pick the block",
+            OBS_STATE: torch.zeros(STATE_DIM),
+            f"{OBS_IMAGES}.front": torch.rand(3, 16, 16),
+        }
+    )
+    assert processed[OBS_STATE].shape == (1, MAX_STATE_DIM)
+
+
 def test_evo1_policy_forward_and_inference_use_batched_embedding(monkeypatch):
    monkeypatch.setattr(modeling_evo1, "Evo1Model", DummyEvo1Model)
    policy = modeling_evo1.Evo1Policy(make_config())
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.

-import sys
 from types import SimpleNamespace
 from unittest.mock import MagicMock

@@ -47,18 +46,23 @@ def test_make_policy_keeps_peft_adapter_and_base_revisions_separate(monkeypatch)
    peft_config_from_pretrained = MagicMock(return_value=peft_config)
    adapted_policy = torch.nn.Linear(1, 1)
    peft_model_from_pretrained = MagicMock(return_value=adapted_policy)
-    monkeypatch.setitem(
-        sys.modules,
-        "peft",
-        SimpleNamespace(
-            PeftConfig=SimpleNamespace(from_pretrained=peft_config_from_pretrained),
-            PeftModel=SimpleNamespace(from_pretrained=peft_model_from_pretrained),
-        ),
+    require_package = MagicMock()
+    monkeypatch.setattr(policy_factory, "require_package", require_package)
+    monkeypatch.setattr(
+        policy_factory,
+        "PeftConfig",
+        SimpleNamespace(from_pretrained=peft_config_from_pretrained),
+    )
+    monkeypatch.setattr(
+        policy_factory,
+        "PeftModel",
+        SimpleNamespace(from_pretrained=peft_model_from_pretrained),
    )

    policy = policy_factory.make_policy(cfg, ds_meta=dataset_meta)

    assert policy is adapted_policy
+    require_package.assert_called_once_with("peft", extra="peft")
    peft_config_from_pretrained.assert_called_once_with(
        "user/adapter",
        revision="adapter-sha",
@@ -0,0 +1,45 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+
+from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
+from lerobot.robots.so_follower.robot_kinematic_processor import (
+    ForwardKinematicsJointsToEEAction,
+    ForwardKinematicsJointsToEEObservation,
+)
+
+MOTOR_NAMES = ["shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll", "gripper"]
+EE_KEYS = {f"ee.{k}" for k in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]}
+
+
+def _joint_bucket(feature_type: FeatureType) -> dict[str, PolicyFeature]:
+    return {f"{n}.pos": PolicyFeature(type=feature_type, shape=(1,)) for n in MOTOR_NAMES}
+
+
+@pytest.mark.parametrize(
+    ("step_cls", "bucket", "feature_type"),
+    [
+        (ForwardKinematicsJointsToEEAction, PipelineFeatureType.ACTION, FeatureType.ACTION),
+        (ForwardKinematicsJointsToEEObservation, PipelineFeatureType.OBSERVATION, FeatureType.STATE),
+    ],
+)
+def test_fk_feature_schema(step_cls, bucket, feature_type):
+    features = {PipelineFeatureType.ACTION: {}, PipelineFeatureType.OBSERVATION: {}}
+    features[bucket] = _joint_bucket(feature_type)
+    out = step_cls(kinematics=None, motor_names=MOTOR_NAMES).transform_features(features)[bucket]
+    assert set(out) == EE_KEYS
+    assert {feature.type for feature in out.values()} == {feature_type}
@@ -49,7 +49,7 @@ def _make_bus_mock() -> MagicMock:


@pytest.fixture
-def follower():
+def follower(tmp_path):
    bus_mock = _make_bus_mock()

    def _bus_side_effect(*_args, **kwargs):
@@ -71,7 +71,7 @@ def follower():
        ),
        patch.object(SO100Follower, "configure", lambda self: None),
    ):
-        cfg = SO100FollowerConfig(port="/dev/null")
+        cfg = SO100FollowerConfig(port="/dev/null", calibration_dir=tmp_path)
        robot = SO100Follower(cfg)
        yield robot
        if robot.is_connected:
@@ -99,6 +99,27 @@ def test_get_observation(follower):
        assert obs[f"{motor}.pos"] == idx


+def test_get_observation_uses_read_retries(follower):
+    # Feetech buses can intermittently fail a sync_read; the follower should forward the configured
+    # retry count so transient failures don't abort the control loop (see #3131).
+    follower.config.num_read_retries = 7
+    follower.connect()
+    follower.get_observation()
+
+    follower.bus.sync_read.assert_called_once_with("Present_Position", num_retry=7)
+
+
+def test_send_action_uses_read_retries(follower):
+    follower.config.max_relative_target = 10.0
+    follower.config.num_read_retries = 7
+    follower.connect()
+
+    action = {f"{motor}.pos": value * 10 for value, motor in enumerate(follower.bus.motors, 1)}
+    follower.send_action(action)
+
+    follower.bus.sync_read.assert_called_once_with("Present_Position", num_retry=7)
+
+
 def test_send_action(follower):
    follower.connect()

@@ -0,0 +1,36 @@
+#!/usr/bin/env python
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from lerobot.utils.device_utils import get_safe_torch_device, is_torch_device_available
+
+
+def test_cpu_always_available():
+    assert get_safe_torch_device("cpu") == torch.device("cpu")
+    assert is_torch_device_available("cpu")
+
+
+def test_missing_cuda_raises_valueerror():
+    with patch("torch.cuda.is_available", return_value=False), pytest.raises(ValueError, match="CUDA"):
+        get_safe_torch_device("cuda")
+
+
+def test_missing_mps_raises_valueerror():
+    with patch("torch.backends.mps.is_available", return_value=False), pytest.raises(ValueError, match="MPS"):
+        get_safe_torch_device("mps")
@@ -0,0 +1,46 @@
+#!/usr/bin/env python
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import pytest
+
+from lerobot.utils.rotation import Rotation
+
+
+def test_zero_quaternion_rejected():
+    with pytest.raises(ValueError, match="non-zero"):
+        Rotation(np.zeros(4))
+
+
+def test_non_finite_quaternion_rejected():
+    with pytest.raises(ValueError, match="non-zero|finite"):
+        Rotation(np.array([np.nan, 0.0, 0.0, 1.0]))
+
+
+def test_wrong_shape_rejected():
+    with pytest.raises(ValueError, match="shape"):
+        Rotation(np.array([1.0, 0.0, 0.0]))
+
+
+def test_identity_roundtrip():
+    r = Rotation.from_rotvec(np.zeros(3))
+    assert np.allclose(r.as_rotvec(), 0.0)
+    assert np.allclose(r.as_matrix(), np.eye(3))
+
+
+def test_rotvec_roundtrip():
+    rotvec = np.array([0.1, -0.2, 0.3])
+    r = Rotation.from_rotvec(rotvec)
+    assert np.allclose(r.as_rotvec(), rotvec, atol=1e-6)