mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-30 13:09:40 +00:00
Compare commits
66 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 4cd9fc73f8 | |||
| 57d33c295e | |||
| 6e23429105 | |||
| fffa42cc5e | |||
| f33089b027 | |||
| 9f32f57b59 | |||
| 0a9b9d9a93 | |||
| 16b915ede5 | |||
| 95d9029039 | |||
| bbfc4ff443 | |||
| 0c57cd03f2 | |||
| a6b06eac38 | |||
| 962ed236af | |||
| 3ae036ee40 | |||
| cdf5141688 | |||
| 9f0663e9e3 | |||
| 503f3e57ae | |||
| 6d24f20eb4 | |||
| af163fd032 | |||
| 77259f436e | |||
| 85f5c3606d | |||
| b587e81587 | |||
| 4658dada9b | |||
| 57ea6f4106 | |||
| 4209639f33 | |||
| fc7a0bc2fd | |||
| 5f6513551c | |||
| 70e157e00f | |||
| 1837be51bf | |||
| bedd56eed9 | |||
| c165e4df68 | |||
| 5e24da483a | |||
| 9c54665a76 | |||
| f6a845c30c | |||
| 45e8336854 | |||
| 5046e2df32 | |||
| 1c88e26c6d | |||
| 69a3edfa33 | |||
| 2492ce2c29 | |||
| c8e75da55f | |||
| 2eae31ea2b | |||
| c997abe739 | |||
| c73579055e | |||
| 4be438161b | |||
| 806d28a883 | |||
| 573b65ff6b | |||
| bc55713e7c | |||
| 4f53c42583 | |||
| bfced3d149 | |||
| 4969813d4e | |||
| 1c87ca31a3 | |||
| 4bcde762cc | |||
| 943ae78cfe | |||
| 3363688f1e | |||
| 0876629e72 | |||
| 305614b8c6 | |||
| 02d3202c4f | |||
| 3b6de2fdf8 | |||
| 744f3667c0 | |||
| fdde436776 | |||
| 5c683c65c6 | |||
| dfbc25c58f | |||
| 804c76bcc2 | |||
| e6afa69be9 | |||
| 31d1439e29 | |||
| 1c118c6359 |
@@ -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:
|
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
|
```python
|
||||||
# Use SmolVLA policy with LIBERO environment
|
# Use SmolVLA policy with LIBERO environment
|
||||||
libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
|
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,
|
policy_cfg=act_cfg,
|
||||||
)
|
)
|
||||||
act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg)
|
act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg)
|
||||||
|
```
|
||||||
|
|
||||||
### 3. **Easier Experimentation**
|
### 3. **Easier Experimentation**
|
||||||
|
|
||||||
@@ -145,7 +132,7 @@ class LiberoVelocityProcessorStep(ObservationProcessorStep):
|
|||||||
state = torch.cat([eef_pos, eef_axisangle, eef_vel,
|
state = torch.cat([eef_pos, eef_axisangle, eef_vel,
|
||||||
gripper_pos, gripper_vel], dim=-1) # 14D
|
gripper_pos, gripper_vel], dim=-1) # 14D
|
||||||
return state
|
return state
|
||||||
````
|
```
|
||||||
|
|
||||||
### 4. **Cleaner Environment Code**
|
### 4. **Cleaner Environment Code**
|
||||||
|
|
||||||
|
|||||||
@@ -211,7 +211,7 @@ Record, Replay and Train with Hope-JR is still experimental.
|
|||||||
|
|
||||||
### Record
|
### 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
|
```bash
|
||||||
lerobot-record \
|
lerobot-record \
|
||||||
|
|||||||
@@ -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
|
- [`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.
|
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):
|
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]"
|
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.
|
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
|
# 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
|
## 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:
|
The finetuned model can be found here:
|
||||||
|
|
||||||
|
|||||||
@@ -93,7 +93,7 @@ lerobot-train --help
|
|||||||
|
|
||||||
## Evaluate the finetuned model and run it in real-time
|
## 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:
|
Once you are logged in, you can run inference in your setup by doing:
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
|
|||||||
@@ -68,6 +68,10 @@ class UnitreeG1Config(RobotConfig):
|
|||||||
# Compensates for gravity on the unitree's arms using the arm ik solver
|
# Compensates for gravity on the unitree's arms using the arm ik solver
|
||||||
gravity_compensation: bool = False
|
gravity_compensation: bool = False
|
||||||
|
|
||||||
# Lower-body controller class name, e.g. "GrootLocomotionController" or
|
# Locomotion controller class name, e.g. "GrootLocomotionController",
|
||||||
# "HolosomaLocomotionController". None disables it.
|
# "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
|
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",
|
||||||
|
]
|
||||||
+1
-1
@@ -21,7 +21,7 @@ import numpy as np
|
|||||||
import onnxruntime as ort
|
import onnxruntime as ort
|
||||||
from huggingface_hub import hf_hub_download
|
from huggingface_hub import hf_hub_download
|
||||||
|
|
||||||
from .g1_utils import (
|
from ..g1_utils import (
|
||||||
REMOTE_AXES,
|
REMOTE_AXES,
|
||||||
REMOTE_BUTTONS,
|
REMOTE_BUTTONS,
|
||||||
G1_29_JointIndex,
|
G1_29_JointIndex,
|
||||||
+1
-1
@@ -22,7 +22,7 @@ import onnx
|
|||||||
import onnxruntime as ort
|
import onnxruntime as ort
|
||||||
from huggingface_hub import hf_hub_download
|
from huggingface_hub import hf_hub_download
|
||||||
|
|
||||||
from .g1_utils import (
|
from ..g1_utils import (
|
||||||
REMOTE_AXES,
|
REMOTE_AXES,
|
||||||
G1_29_JointArmIndex,
|
G1_29_JointArmIndex,
|
||||||
G1_29_JointIndex,
|
G1_29_JointIndex,
|
||||||
@@ -0,0 +1,360 @@
|
|||||||
|
#!/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 the standing pose (``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 json
|
||||||
|
import logging
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import onnx
|
||||||
|
import onnxruntime as ort
|
||||||
|
from huggingface_hub import hf_hub_download
|
||||||
|
|
||||||
|
from ..g1_utils import (
|
||||||
|
ISAACLAB_TO_MUJOCO,
|
||||||
|
MUJOCO_TO_ISAACLAB,
|
||||||
|
G1_29_JointIndex,
|
||||||
|
get_gravity_orientation,
|
||||||
|
)
|
||||||
|
|
||||||
|
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
|
||||||
|
|
||||||
|
# SONIC decoder checkpoint: NVIDIA's decoder ONNX re-packaged with its deploy constants
|
||||||
|
# (kp/kd PD gains, the standing pose default_angles, and the residual action_scale) embedded
|
||||||
|
# in the ONNX metadata; see upload_sonic_decoder.py for provisioning. The runtime loads the
|
||||||
|
# model *and* all of these straight from the checkpoint (the Holosoma convention), so no
|
||||||
|
# motor-physics math happens at deploy time.
|
||||||
|
DEFAULT_SONIC_REPO_ID = "lerobot/sonic_decoder"
|
||||||
|
DECODER_FILENAME = "model_decoder.onnx"
|
||||||
|
DECODER_INPUT_DIM = 994 # token(64) + 10-frame proprio history + gravity
|
||||||
|
|
||||||
|
|
||||||
|
def load_sonic_decoder(repo_id: str = DEFAULT_SONIC_REPO_ID):
|
||||||
|
"""Load the SONIC decoder ONNX and its baked-in deploy constants from the checkpoint.
|
||||||
|
|
||||||
|
Returns ``(decoder_session, kp, kd, default_angles, action_scale, neutral_token)``. The
|
||||||
|
gains/pose/scale are (29,) float32 in IsaacLab joint order and ``neutral_token`` is the
|
||||||
|
(64,) float32 idle latent -- all read from the ONNX ``metadata_props`` rather than
|
||||||
|
recomputed/hardcoded at deploy time (mirrors ``holosoma_locomotion.load_policy``).
|
||||||
|
"""
|
||||||
|
decoder_path = hf_hub_download(repo_id=repo_id, filename=DECODER_FILENAME)
|
||||||
|
so = ort.SessionOptions()
|
||||||
|
so.log_severity_level = 3 # quiet ORT logs
|
||||||
|
session = ort.InferenceSession(decoder_path, sess_options=so)
|
||||||
|
dec_dim = int(session.get_inputs()[0].shape[1])
|
||||||
|
if dec_dim != DECODER_INPUT_DIM:
|
||||||
|
raise RuntimeError(f"Unexpected decoder input dim {dec_dim} (expected {DECODER_INPUT_DIM})")
|
||||||
|
|
||||||
|
meta = {p.key: p.value for p in onnx.load(decoder_path, load_external_data=False).metadata_props}
|
||||||
|
required = ("kp", "kd", "default_angles", "action_scale", "neutral_token")
|
||||||
|
missing = [k for k in required if k not in meta]
|
||||||
|
if missing:
|
||||||
|
raise ValueError(
|
||||||
|
f"SONIC decoder ONNX at {repo_id} is missing metadata {missing}; "
|
||||||
|
"re-run upload_sonic_decoder.py to (re)provision the checkpoint."
|
||||||
|
)
|
||||||
|
arr = {k: np.array(json.loads(meta[k]), dtype=np.float32) for k in required}
|
||||||
|
logger.info("Loaded SONIC deploy constants from %s (%d joints)", repo_id, len(arr["kp"]))
|
||||||
|
return session, arr["kp"], arr["kd"], arr["default_angles"], arr["action_scale"], arr["neutral_token"]
|
||||||
|
|
||||||
|
|
||||||
|
# Action-feature prefix for the latent-token interface (see _extract_token_from_action).
|
||||||
|
TOKEN_ACTION_PREFIX = "motion_token" # nosec B105 - feature-key prefix, not a secret
|
||||||
|
# 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" # nosec B105 - feature-key prefix, not a secret
|
||||||
|
|
||||||
|
|
||||||
|
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
|
||||||
|
|
||||||
|
|
||||||
|
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). ``default_angles`` and
|
||||||
|
``action_scale`` are (29,) float32 in IsaacLab order, loaded from the checkpoint.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, decoder, default_angles, action_scale):
|
||||||
|
self.decoder = decoder
|
||||||
|
self.decoder_input = decoder.get_inputs()[0].name
|
||||||
|
self.default_angles = np.asarray(default_angles, np.float32)
|
||||||
|
self.action_scale = np.asarray(action_scale, np.float32)
|
||||||
|
self.default_angles_mj = self.default_angles[MUJOCO_TO_ISAACLAB]
|
||||||
|
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)
|
||||||
|
# Reorder IsaacLab-order state into the MuJoCo order the decoder consumes. This
|
||||||
|
# permutation direction is validated against the deployed SONIC ONNX; don't flip it.
|
||||||
|
q_mj = q[MUJOCO_TO_ISAACLAB]
|
||||||
|
dq_mj = dq[MUJOCO_TO_ISAACLAB]
|
||||||
|
self.h_q_mj = [q_mj - self.default_angles_mj] + 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", self.default_angles[m.value])
|
||||||
|
for m, n in zip(G1_29_JointIndex, jnames, strict=False)
|
||||||
|
],
|
||||||
|
np.float32,
|
||||||
|
)
|
||||||
|
dq = np.array([robot_obs.get(f"{n}.dq", 0.0) for n in jnames], np.float32)
|
||||||
|
quat = np.array(
|
||||||
|
[
|
||||||
|
robot_obs.get("imu.quat.w", 1),
|
||||||
|
robot_obs.get("imu.quat.x", 0),
|
||||||
|
robot_obs.get("imu.quat.y", 0),
|
||||||
|
robot_obs.get("imu.quat.z", 0),
|
||||||
|
],
|
||||||
|
np.float32,
|
||||||
|
)
|
||||||
|
ang = np.array([robot_obs.get(f"imu.gyro.{a}", 0) for a in "xyz"], np.float32)
|
||||||
|
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 = self.default_angles + action_mj[ISAACLAB_TO_MUJOCO] * self.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_sess, self.kp, self.kd, default_angles, action_scale, neutral_token = load_sonic_decoder()
|
||||||
|
self.default_angles = default_angles
|
||||||
|
self.neutral_token = neutral_token
|
||||||
|
self.controller = SonicDecoder(decoder_sess, default_angles, action_scale)
|
||||||
|
|
||||||
|
@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
|
||||||
|
self._default_angles = self._runtime.default_angles
|
||||||
|
self._neutral_token = self._runtime.neutral_token
|
||||||
|
|
||||||
|
# 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. The controller 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).
|
||||||
|
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", self._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, obs: dict) -> dict:
|
||||||
|
if not obs:
|
||||||
|
return {}
|
||||||
|
|
||||||
|
# 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:
|
||||||
|
# No token has arrived yet: hold the checkpoint's neutral token, which the decoder
|
||||||
|
# maps to a stable, natural standing pose.
|
||||||
|
self._last_token = self._neutral_token.copy()
|
||||||
|
# Either a fresh token this tick or the last one received (held between the ~30 Hz
|
||||||
|
# token stream and the ~50 Hz control loop).
|
||||||
|
return self._startup_blend(obs, self.controller.step(obs, self._last_token))
|
||||||
|
|
||||||
|
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 the neutral token is re-seeded after a reset.
|
||||||
|
self._last_token = None
|
||||||
|
|
||||||
|
def shutdown(self):
|
||||||
|
self._runtime.shutdown()
|
||||||
@@ -23,6 +23,47 @@ import numpy as np
|
|||||||
|
|
||||||
NUM_MOTORS = 29
|
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,
|
||||||
|
)
|
||||||
|
# The two orderings are inverses of each other, so derive one from the other (argsort) to
|
||||||
|
# guarantee they can never drift out of sync.
|
||||||
|
MUJOCO_TO_ISAACLAB = np.argsort(ISAACLAB_TO_MUJOCO).astype(np.int32)
|
||||||
|
|
||||||
REMOTE_AXES = ("remote.lx", "remote.ly", "remote.rx", "remote.ry")
|
REMOTE_AXES = ("remote.lx", "remote.ly", "remote.rx", "remote.ry")
|
||||||
REMOTE_BUTTONS = tuple(f"remote.button.{i}" for i in range(16))
|
REMOTE_BUTTONS = tuple(f"remote.button.{i}" for i in range(16))
|
||||||
REMOTE_KEYS = REMOTE_AXES + REMOTE_BUTTONS
|
REMOTE_KEYS = REMOTE_AXES + REMOTE_BUTTONS
|
||||||
@@ -68,8 +109,9 @@ def make_locomotion_controller(name: str | None):
|
|||||||
if name is None:
|
if name is None:
|
||||||
return None
|
return None
|
||||||
controllers = {
|
controllers = {
|
||||||
"GrootLocomotionController": "lerobot.robots.unitree_g1.gr00t_locomotion",
|
"GrootLocomotionController": "lerobot.robots.unitree_g1.controllers.gr00t_locomotion",
|
||||||
"HolosomaLocomotionController": "lerobot.robots.unitree_g1.holosoma_locomotion",
|
"HolosomaLocomotionController": "lerobot.robots.unitree_g1.controllers.holosoma_locomotion",
|
||||||
|
"SonicWholeBodyController": "lerobot.robots.unitree_g1.controllers.sonic_whole_body",
|
||||||
}
|
}
|
||||||
module_path = controllers.get(name)
|
module_path = controllers.get(name)
|
||||||
if module_path is None:
|
if module_path is None:
|
||||||
|
|||||||
@@ -34,7 +34,6 @@ from .config_unitree_g1 import UnitreeG1Config
|
|||||||
from .g1_kinematics import G1_29_ArmIK
|
from .g1_kinematics import G1_29_ArmIK
|
||||||
from .g1_utils import (
|
from .g1_utils import (
|
||||||
REMOTE_AXES,
|
REMOTE_AXES,
|
||||||
REMOTE_KEYS,
|
|
||||||
G1_29_JointArmIndex,
|
G1_29_JointArmIndex,
|
||||||
G1_29_JointIndex,
|
G1_29_JointIndex,
|
||||||
default_remote_input,
|
default_remote_input,
|
||||||
@@ -148,22 +147,54 @@ class UnitreeG1(Robot):
|
|||||||
|
|
||||||
self.arm_ik = G1_29_ArmIK() if config.gravity_compensation else None
|
self.arm_ik = G1_29_ArmIK() if config.gravity_compensation else None
|
||||||
|
|
||||||
# Lower-body controller loaded dynamically
|
# Lower-body / whole-body controller loaded dynamically
|
||||||
self.controller: LocomotionController | None = make_locomotion_controller(config.controller)
|
self.controller: LocomotionController | None = make_locomotion_controller(config.controller)
|
||||||
|
|
||||||
# Controller thread state
|
# Controller thread state
|
||||||
self._controller_thread = None
|
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_action_lock = threading.Lock()
|
||||||
self.controller_input = default_remote_input()
|
self.controller_input = default_remote_input()
|
||||||
self.controller_output = {}
|
self.controller_output = {}
|
||||||
|
|
||||||
|
# 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
|
def _subscribe_lowstate(self): # polls robot state @ 250Hz
|
||||||
while not self._shutdown_event.is_set():
|
while not self._shutdown_event.is_set():
|
||||||
start_time = time.time()
|
start_time = time.time()
|
||||||
|
|
||||||
# Step simulation if in simulation mode
|
# Step simulation if in simulation mode
|
||||||
if self.config.is_simulation and self.sim_env is not None:
|
if self.config.is_simulation and self.sim_env is not None:
|
||||||
|
try:
|
||||||
self.sim_env.step()
|
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()
|
msg = self.lowstate_subscriber.Read()
|
||||||
if msg is not None:
|
if msg is not None:
|
||||||
@@ -231,15 +262,38 @@ class UnitreeG1(Robot):
|
|||||||
features[f"{cam}_depth"] = (cfg.height, cfg.width, 1)
|
features[f"{cam}_depth"] = (cfg.height, cfg.width, 1)
|
||||||
return features
|
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
|
@cached_property
|
||||||
def observation_features(self) -> dict[str, type | tuple]:
|
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
|
@cached_property
|
||||||
def action_features(self) -> dict[str, type]:
|
def action_features(self) -> dict[str, type]:
|
||||||
|
# No controller configured at all: raw 29-DoF joint teleop.
|
||||||
if self.controller is None:
|
if self.controller is None:
|
||||||
return {f"{G1_29_JointIndex(motor).name}.q": float for motor in G1_29_JointIndex}
|
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}
|
arm_features = {f"{G1_29_JointArmIndex(motor).name}.q": float for motor in G1_29_JointArmIndex}
|
||||||
remote_features = dict.fromkeys(REMOTE_AXES, float)
|
remote_features = dict.fromkeys(REMOTE_AXES, float)
|
||||||
return {**arm_features, **remote_features}
|
return {**arm_features, **remote_features}
|
||||||
@@ -255,6 +309,11 @@ class UnitreeG1(Robot):
|
|||||||
while not self._shutdown_event.is_set():
|
while not self._shutdown_event.is_set():
|
||||||
start_time = time.time()
|
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:
|
with self._lowstate_lock:
|
||||||
lowstate = self._lowstate
|
lowstate = self._lowstate
|
||||||
|
|
||||||
@@ -271,8 +330,12 @@ class UnitreeG1(Robot):
|
|||||||
with self._controller_action_lock:
|
with self._controller_action_lock:
|
||||||
controller_input = dict(self.controller_input)
|
controller_input = dict(self.controller_input)
|
||||||
|
|
||||||
# Run controller step
|
# Full-body controllers (SONIC) consume the full observation dict; others
|
||||||
controller_action = self.controller.run_step(controller_input, lowstate)
|
# take the raw lowstate. get_observation() is the single lowstate -> obs builder.
|
||||||
|
controller_state = (
|
||||||
|
self.get_observation() if getattr(self.controller, "full_body", False) else lowstate
|
||||||
|
)
|
||||||
|
controller_action = self.controller.run_step(controller_input, controller_state)
|
||||||
|
|
||||||
# Write controller output snapshot
|
# Write controller output snapshot
|
||||||
with self._controller_action_lock:
|
with self._controller_action_lock:
|
||||||
@@ -343,6 +406,9 @@ class UnitreeG1(Robot):
|
|||||||
|
|
||||||
self.kp = np.array(self.config.kp, dtype=np.float32)
|
self.kp = np.array(self.config.kp, dtype=np.float32)
|
||||||
self.kd = np.array(self.config.kd, 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:
|
for joint in G1_29_JointIndex:
|
||||||
self.msg.motor_cmd[joint].mode = 1
|
self.msg.motor_cmd[joint].mode = 1
|
||||||
@@ -391,6 +457,10 @@ class UnitreeG1(Robot):
|
|||||||
if self._controller_thread.is_alive():
|
if self._controller_thread.is_alive():
|
||||||
logger.warning("Controller thread did not stop cleanly")
|
logger.warning("Controller thread did not stop cleanly")
|
||||||
|
|
||||||
|
# Release controller resources (e.g. SONIC decoder sessions).
|
||||||
|
if self.controller is not None and hasattr(self.controller, "shutdown"):
|
||||||
|
self.controller.shutdown()
|
||||||
|
|
||||||
# Close simulation environment
|
# Close simulation environment
|
||||||
if self.config.is_simulation and self.sim_env is not None:
|
if self.config.is_simulation and self.sim_env is not None:
|
||||||
try:
|
try:
|
||||||
@@ -461,6 +531,15 @@ class UnitreeG1(Robot):
|
|||||||
if lowstate.wireless_remote:
|
if lowstate.wireless_remote:
|
||||||
obs["wireless_remote"] = lowstate.wireless_remote
|
obs["wireless_remote"] = lowstate.wireless_remote
|
||||||
|
|
||||||
|
# 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
|
||||||
|
|
||||||
|
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
|
# Cameras - read images from ZMQ cameras
|
||||||
for cam_name, cam in self._cameras.items():
|
for cam_name, cam in self._cameras.items():
|
||||||
if getattr(cam, "use_rgb", True):
|
if getattr(cam, "use_rgb", True):
|
||||||
@@ -473,9 +552,22 @@ class UnitreeG1(Robot):
|
|||||||
def send_action(self, action: RobotAction) -> RobotAction:
|
def send_action(self, action: RobotAction) -> RobotAction:
|
||||||
action_to_publish = action
|
action_to_publish = action
|
||||||
if self.controller is not None:
|
if self.controller is not None:
|
||||||
|
# SONIC decoder: pull the 64-D latent token out of the action and remember it
|
||||||
|
# for the observation.state echo. The controller thread reads it back from
|
||||||
|
# controller_input (populated below) and decodes it into a 29-DoF command.
|
||||||
|
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)
|
||||||
|
# Full-body controllers (SONIC) own the whole 29-DoF command; nothing to
|
||||||
|
# publish here (the controller thread is the sole publisher).
|
||||||
|
if getattr(self.controller, "full_body", False):
|
||||||
|
return action
|
||||||
# Controller thread owns legs/waist. Here we only update joystick inputs
|
# Controller thread owns legs/waist. Here we only update joystick inputs
|
||||||
# and publish arm targets from the teleoperator.
|
# and publish arm targets from the teleoperator.
|
||||||
self._update_controller_action(action)
|
|
||||||
arm_prefixes = tuple(j.name for j in G1_29_JointArmIndex)
|
arm_prefixes = tuple(j.name for j in G1_29_JointArmIndex)
|
||||||
action_to_publish = {
|
action_to_publish = {
|
||||||
key: value
|
key: value
|
||||||
@@ -503,11 +595,17 @@ class UnitreeG1(Robot):
|
|||||||
return action
|
return action
|
||||||
|
|
||||||
def _update_controller_action(self, action: RobotAction) -> None:
|
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 (locomotion ``remote.*`` axes or
|
||||||
|
SONIC ``motion_token.*`` values) 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:
|
with self._controller_action_lock:
|
||||||
for key in REMOTE_KEYS:
|
for key, value in action.items():
|
||||||
if key in action:
|
if isinstance(key, str) and value is not None:
|
||||||
self.controller_input[key] = action[key]
|
self.controller_input[key] = value
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def is_calibrated(self) -> bool:
|
def is_calibrated(self) -> bool:
|
||||||
@@ -537,6 +635,18 @@ class UnitreeG1(Robot):
|
|||||||
if default_positions is None:
|
if default_positions is None:
|
||||||
default_positions = np.array(self.config.default_positions, dtype=np.float32)
|
default_positions = np.array(self.config.default_positions, dtype=np.float32)
|
||||||
|
|
||||||
|
# Full-body controllers (SONIC) own the whole 29-DoF command and ignore
|
||||||
|
# ``<joint>.q`` in send_action(), so reset() must publish the default pose
|
||||||
|
# directly. Pause the background controller first so the two aren't both writing
|
||||||
|
# low commands while the robot moves to the default pose.
|
||||||
|
full_body = getattr(self.controller, "full_body", False)
|
||||||
|
paused = False
|
||||||
|
if full_body and self._controller_thread is not None:
|
||||||
|
self._controller_paused.set()
|
||||||
|
paused = True
|
||||||
|
time.sleep(control_dt) # let any in-flight controller tick settle
|
||||||
|
|
||||||
|
try:
|
||||||
if self.config.is_simulation and self.sim_env is not None:
|
if self.config.is_simulation and self.sim_env is not None:
|
||||||
self.sim_env.reset()
|
self.sim_env.reset()
|
||||||
self.publish_lowcmd(
|
self.publish_lowcmd(
|
||||||
@@ -565,6 +675,11 @@ class UnitreeG1(Robot):
|
|||||||
interp_pos = init_dof_pos[motor.value] * (1 - alpha) + target_pos * alpha
|
interp_pos = init_dof_pos[motor.value] * (1 - alpha) + target_pos * alpha
|
||||||
action_dict[f"{motor.name}.q"] = float(interp_pos)
|
action_dict[f"{motor.name}.q"] = float(interp_pos)
|
||||||
|
|
||||||
|
# Full-body controllers no-op in send_action(); publish the pose
|
||||||
|
# directly (arm-only controllers keep the send_action() path).
|
||||||
|
if full_body:
|
||||||
|
self.publish_lowcmd(action_dict)
|
||||||
|
else:
|
||||||
self.send_action(action_dict)
|
self.send_action(action_dict)
|
||||||
|
|
||||||
# Maintain constant control rate
|
# Maintain constant control rate
|
||||||
@@ -572,8 +687,12 @@ class UnitreeG1(Robot):
|
|||||||
sleep_time = max(0, control_dt - elapsed)
|
sleep_time = max(0, control_dt - elapsed)
|
||||||
time.sleep(sleep_time)
|
time.sleep(sleep_time)
|
||||||
|
|
||||||
# Reset controller internal state (gait phase, obs history, etc.)
|
# Reset controller internal state (gait phase, obs history, etc.) before
|
||||||
|
# resuming so its buffers reflect the post-reset pose.
|
||||||
if self.controller is not None and hasattr(self.controller, "reset"):
|
if self.controller is not None and hasattr(self.controller, "reset"):
|
||||||
self.controller.reset()
|
self.controller.reset()
|
||||||
|
finally:
|
||||||
|
if paused:
|
||||||
|
self._controller_paused.clear()
|
||||||
|
|
||||||
logger.info("Reached default position")
|
logger.info("Reached default position")
|
||||||
|
|||||||
Reference in New Issue
Block a user