mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-30 13:09:40 +00:00
Compare commits
74 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 4cd9fc73f8 | |||
| 57d33c295e | |||
| 6e23429105 | |||
| fffa42cc5e | |||
| f33089b027 | |||
| 9f32f57b59 | |||
| 0a9b9d9a93 | |||
| 16b915ede5 | |||
| 95d9029039 | |||
| bbfc4ff443 | |||
| 0c57cd03f2 | |||
| a6b06eac38 | |||
| 36b8face98 | |||
| cd8984cc0a | |||
| b9ded9e761 | |||
| 185f3e1708 | |||
| e36783253a | |||
| 289e577fc7 | |||
| 9c32722eb9 | |||
| b49cb50e01 | |||
| 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 |
+14
-11
@@ -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.
|
||||
|
||||
@@ -321,11 +325,10 @@ SmolVLA ships with `freeze_vision_encoder=True`. Unfreezing usually **improves p
|
||||
|
||||
```bash
|
||||
lerobot-train ... --policy.type=smolvla \
|
||||
--policy.fine_tune_vision_encoder=true
|
||||
--policy.freeze_vision_encoder=false \
|
||||
--policy.train_expert_only=false
|
||||
```
|
||||
|
||||
This selectively trains the vision encoder and connector while leaving the language model frozen. Their learning rate defaults to `0.1 × optimizer_lr`; adjust it with `--policy.vision_encoder_lr_multiplier` if needed.
|
||||
|
||||
### 7.7 Signals to stop / keep going
|
||||
|
||||
- Train loss plateaus → stop, save a Hub checkpoint.
|
||||
|
||||
@@ -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**
|
||||
|
||||
|
||||
@@ -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 \
|
||||
|
||||
@@ -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:
|
||||
|
||||
|
||||
+1
-14
@@ -70,19 +70,6 @@ cd lerobot && lerobot-train \
|
||||
GPU allows it, as long as loading times remain short.
|
||||
</Tip>
|
||||
|
||||
For tasks that require adapting visual features, such as distinguishing new colors or shapes, selectively
|
||||
fine-tune the vision encoder and its connector:
|
||||
|
||||
```bash
|
||||
lerobot-train ... \
|
||||
--policy.path=lerobot/smolvla_base \
|
||||
--policy.fine_tune_vision_encoder=true
|
||||
```
|
||||
|
||||
This keeps the language model frozen with the default `train_expert_only=true` setting and trains the vision
|
||||
path at `0.1` times the main learning rate by default. Fine-tuning the vision encoder increases memory use and
|
||||
can reduce the model's general visual knowledge, so enable it only when the frozen encoder is insufficient.
|
||||
|
||||
Fine-tuning is an art. For a complete overview of the options for finetuning, run
|
||||
|
||||
```bash
|
||||
@@ -106,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
|
||||
|
||||
@@ -67,8 +67,6 @@ class SmolVLAConfig(PreTrainedConfig):
|
||||
|
||||
# Finetuning settings
|
||||
freeze_vision_encoder: bool = True
|
||||
fine_tune_vision_encoder: bool = False # Fine-tune vision + connector; takes priority over freezing.
|
||||
vision_encoder_lr_multiplier: float = 0.1
|
||||
train_expert_only: bool = True
|
||||
train_state_proj: bool = True
|
||||
|
||||
@@ -112,12 +110,6 @@ class SmolVLAConfig(PreTrainedConfig):
|
||||
super().__post_init__()
|
||||
|
||||
"""Input validation (not exhaustive)."""
|
||||
if self.fine_tune_vision_encoder:
|
||||
self.freeze_vision_encoder = False
|
||||
if self.vision_encoder_lr_multiplier <= 0:
|
||||
raise ValueError(
|
||||
f"`vision_encoder_lr_multiplier` must be positive, got {self.vision_encoder_lr_multiplier}."
|
||||
)
|
||||
if self.n_action_steps > self.chunk_size:
|
||||
raise ValueError(
|
||||
f"The chunk size is the upper bound for the number of action steps per model invocation. Got "
|
||||
|
||||
@@ -186,27 +186,8 @@ class SmolVLAPolicy(PreTrainedPolicy):
|
||||
if model_value is not None:
|
||||
model_value.rtc_processor = self.rtc_processor
|
||||
|
||||
def get_optim_params(self):
|
||||
if not self.config.fine_tune_vision_encoder:
|
||||
return self.parameters()
|
||||
|
||||
vision_params = []
|
||||
other_params = []
|
||||
for name, param in self.named_parameters():
|
||||
if not param.requires_grad:
|
||||
continue
|
||||
if ".vision_model." in name or ".connector." in name:
|
||||
vision_params.append(param)
|
||||
else:
|
||||
other_params.append(param)
|
||||
|
||||
return [
|
||||
{"params": other_params},
|
||||
{
|
||||
"params": vision_params,
|
||||
"lr": self.config.optimizer_lr * self.config.vision_encoder_lr_multiplier,
|
||||
},
|
||||
]
|
||||
def get_optim_params(self) -> dict:
|
||||
return self.parameters()
|
||||
|
||||
def _get_action_chunk(
|
||||
self, batch: dict[str, Tensor], noise: Tensor | None = None, **kwargs: Unpack[ActionSelectKwargs]
|
||||
@@ -512,7 +493,6 @@ class VLAFlowMatching(nn.Module):
|
||||
self.vlm_with_expert = SmolVLMWithExpertModel(
|
||||
model_id=self.config.vlm_model_name,
|
||||
freeze_vision_encoder=self.config.freeze_vision_encoder,
|
||||
fine_tune_vision_encoder=self.config.fine_tune_vision_encoder,
|
||||
train_expert_only=self.config.train_expert_only,
|
||||
load_vlm_weights=self.config.load_vlm_weights,
|
||||
attention_mode=self.config.attention_mode,
|
||||
|
||||
@@ -78,7 +78,6 @@ class SmolVLMWithExpertModel(nn.Module):
|
||||
load_vlm_weights: bool = True,
|
||||
train_expert_only: bool = True,
|
||||
freeze_vision_encoder: bool = False,
|
||||
fine_tune_vision_encoder: bool = False,
|
||||
attention_mode: str = "self_attn",
|
||||
num_expert_layers: int = -1,
|
||||
num_vlm_layers: int = -1,
|
||||
@@ -142,7 +141,6 @@ class SmolVLMWithExpertModel(nn.Module):
|
||||
self.num_key_value_heads = self.config.text_config.num_key_value_heads
|
||||
|
||||
self.freeze_vision_encoder = freeze_vision_encoder
|
||||
self.fine_tune_vision_encoder = fine_tune_vision_encoder
|
||||
self.train_expert_only = train_expert_only
|
||||
self.attention_mode = attention_mode
|
||||
self.expert_hidden_size = lm_expert_config.hidden_size
|
||||
@@ -152,6 +150,10 @@ class SmolVLMWithExpertModel(nn.Module):
|
||||
return self.vlm.model
|
||||
|
||||
def set_requires_grad(self):
|
||||
if self.freeze_vision_encoder:
|
||||
self.get_vlm_model().vision_model.eval()
|
||||
for params in self.get_vlm_model().vision_model.parameters():
|
||||
params.requires_grad = False
|
||||
if self.train_expert_only:
|
||||
self.vlm.eval()
|
||||
for params in self.vlm.parameters():
|
||||
@@ -174,18 +176,6 @@ class SmolVLMWithExpertModel(nn.Module):
|
||||
for name, params in self.vlm.named_parameters():
|
||||
if any(k in name for k in frozen_layers):
|
||||
params.requires_grad = False
|
||||
|
||||
if self.freeze_vision_encoder:
|
||||
self.get_vlm_model().vision_model.eval()
|
||||
for params in self.get_vlm_model().vision_model.parameters():
|
||||
params.requires_grad = False
|
||||
|
||||
if self.fine_tune_vision_encoder:
|
||||
for params in self.get_vlm_model().vision_model.parameters():
|
||||
params.requires_grad = True
|
||||
for params in self.get_vlm_model().connector.parameters():
|
||||
params.requires_grad = True
|
||||
|
||||
# To avoid unused params issue with distributed training
|
||||
for name, params in self.lm_expert.named_parameters():
|
||||
if "lm_head" in name:
|
||||
@@ -194,15 +184,11 @@ class SmolVLMWithExpertModel(nn.Module):
|
||||
def train(self, mode: bool = True):
|
||||
super().train(mode)
|
||||
|
||||
if self.train_expert_only:
|
||||
self.vlm.eval()
|
||||
|
||||
if self.freeze_vision_encoder:
|
||||
self.get_vlm_model().vision_model.eval()
|
||||
|
||||
if self.fine_tune_vision_encoder:
|
||||
self.get_vlm_model().vision_model.train(mode)
|
||||
self.get_vlm_model().connector.train(mode)
|
||||
if self.train_expert_only:
|
||||
self.vlm.eval()
|
||||
|
||||
def embed_image(self, image: torch.Tensor):
|
||||
patch_attention_mask = None
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -68,6 +68,10 @@ class UnitreeG1Config(RobotConfig):
|
||||
# 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",
|
||||
]
|
||||
+1
-1
@@ -21,7 +21,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,
|
||||
+1
-1
@@ -22,7 +22,7 @@ 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,
|
||||
@@ -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
|
||||
|
||||
# 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_BUTTONS = tuple(f"remote.button.{i}" for i in range(16))
|
||||
REMOTE_KEYS = REMOTE_AXES + REMOTE_BUTTONS
|
||||
@@ -68,8 +109,9 @@ def make_locomotion_controller(name: str | 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:
|
||||
|
||||
@@ -34,7 +34,6 @@ 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,
|
||||
@@ -148,22 +147,54 @@ class UnitreeG1(Robot):
|
||||
|
||||
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)
|
||||
|
||||
# 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 = {}
|
||||
|
||||
# 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 +262,38 @@ 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]:
|
||||
# No controller configured at all: raw 29-DoF joint teleop.
|
||||
if self.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 +309,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,8 +330,12 @@ class UnitreeG1(Robot):
|
||||
with self._controller_action_lock:
|
||||
controller_input = dict(self.controller_input)
|
||||
|
||||
# Run controller step
|
||||
controller_action = self.controller.run_step(controller_input, lowstate)
|
||||
# Full-body controllers (SONIC) consume the full observation dict; others
|
||||
# take the raw lowstate. get_observation() is the single lowstate -> obs builder.
|
||||
controller_state = (
|
||||
self.get_observation() if getattr(self.controller, "full_body", False) else lowstate
|
||||
)
|
||||
controller_action = self.controller.run_step(controller_input, controller_state)
|
||||
|
||||
# Write controller output snapshot
|
||||
with self._controller_action_lock:
|
||||
@@ -343,6 +406,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
|
||||
@@ -391,6 +457,10 @@ class UnitreeG1(Robot):
|
||||
if self._controller_thread.is_alive():
|
||||
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
|
||||
if self.config.is_simulation and self.sim_env is not None:
|
||||
try:
|
||||
@@ -461,6 +531,15 @@ class UnitreeG1(Robot):
|
||||
if 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
|
||||
for cam_name, cam in self._cameras.items():
|
||||
if getattr(cam, "use_rgb", True):
|
||||
@@ -473,9 +552,22 @@ class UnitreeG1(Robot):
|
||||
def send_action(self, action: RobotAction) -> RobotAction:
|
||||
action_to_publish = action
|
||||
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
|
||||
# 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 +595,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 (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:
|
||||
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:
|
||||
@@ -537,43 +635,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) 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")
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user