admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-29 12:39:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(pi052): add language-supervised policy

Browse Source
This commit is contained in:
Pepijn
2026-07-28 10:08:20 +02:00
parent f2b90e3ad6
commit a6f533a6dd
41 changed files with 7406 additions and 289 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+7 -7
View File
@@ -101,13 +101,13 @@ lerobot-train \
--dataset.repo_id=lerobot/aloha_mobile_cabinet
```
| Category | Models |
| -------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
| **Imitation Learning** | [ACT](./docs/source/policy_act_README.md), [Diffusion](./docs/source/policy_diffusion_README.md), [VQ-BeT](./docs/source/policy_vqbet_README.md), [Multitask DiT Policy](./docs/source/policy_multi_task_dit_README.md) |
| **Reinforcement Learning** | [HIL-SERL](./docs/source/hilserl.mdx), [TDMPC](./docs/source/policy_tdmpc_README.md) & QC-FQL (coming soon) |
| **VLAs Models** | [Pi0](./docs/source/pi0.mdx), [Pi0Fast](./docs/source/pi0fast.mdx), [Pi0.5](./docs/source/pi05.mdx), [GR00T N1.7](./docs/source/policy_groot_README.md), [SmolVLA](./docs/source/policy_smolvla_README.md), [XVLA](./docs/source/xvla.mdx), [EO-1](./docs/source/eo1.mdx), [MolmoAct2](./docs/source/molmoact2.mdx), [WALL-OSS](./docs/source/walloss.mdx), [EVO1](./docs/source/evo1.mdx) |
| **World Models** | [VLA-JEPA](./docs/source/vla_jepa.mdx), [LingBot-VA](./docs/source/lingbot_va.mdx), [FastWAM](./docs/source/fastwam.mdx) |
| **Reward Models** | [SARM](./docs/source/sarm.mdx), [TOPReward](./docs/source/topreward.mdx), [Robometer](./docs/source/robometer.mdx) |
| Category | Models |
| -------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
| **Imitation Learning** | [ACT](./docs/source/policy_act_README.md), [Diffusion](./docs/source/policy_diffusion_README.md), [VQ-BeT](./docs/source/policy_vqbet_README.md), [Multitask DiT Policy](./docs/source/policy_multi_task_dit_README.md) |
| **Reinforcement Learning** | [HIL-SERL](./docs/source/hilserl.mdx), [TDMPC](./docs/source/policy_tdmpc_README.md) & QC-FQL (coming soon) |
| **VLAs Models** | [Pi0](./docs/source/pi0.mdx), [Pi0Fast](./docs/source/pi0fast.mdx), [Pi0.5](./docs/source/pi05.mdx), [Pi052](./docs/source/pi052.mdx), [GR00T N1.7](./docs/source/policy_groot_README.md), [SmolVLA](./docs/source/policy_smolvla_README.md), [XVLA](./docs/source/xvla.mdx), [EO-1](./docs/source/eo1.mdx), [MolmoAct2](./docs/source/molmoact2.mdx), [WALL-OSS](./docs/source/walloss.mdx), [EVO1](./docs/source/evo1.mdx) |
| **World Models** | [VLA-JEPA](./docs/source/vla_jepa.mdx), [LingBot-VA](./docs/source/lingbot_va.mdx), [FastWAM](./docs/source/fastwam.mdx) |
| **Reward Models** | [SARM](./docs/source/sarm.mdx), [TOPReward](./docs/source/topreward.mdx), [Robometer](./docs/source/robometer.mdx) |
Similarly to the hardware, you can easily implement your own policy & leverage LeRobot's data collection, training, and visualization tools, and share your model to the HF Hub.
docs/source/_toctree.yml
+2
View File
@@ -63,6 +63,8 @@
title: π₀-FAST (Pi0Fast)
- local: pi05
title: π₀.₅ (Pi05)
- local: pi052
title: π₀.₅ with language supervision (Pi052)
- local: molmoact2
title: MolmoAct2
- local: vla_jepa
docs/source/pi052.mdx
+274
View File
@@ -0,0 +1,274 @@
# π₀.₅ with language supervision (Pi052)
Pi052 extends [Pi05](./pi05) with a trainable PaliGemma language head and a
runtime that alternates language generation with action generation. A single
checkpoint can predict a low-level subtask, optionally update memory or answer
visual questions, and condition its flow-matching action expert on that text.
Use Pi05 when you only need task-conditioned actions. Use Pi052 when the policy
must generate or consume intermediate language during a rollout.
## How Pi052 differs from Pi05
| Capability | Pi05 | Pi052 |
| ------------------- | ------------------------------------------------------ | --------------------------------------------------------------------------------- |
| Action model | PaliGemma vision-language prefix + Gemma action expert | Same base architecture |
| Language head | Not trained for runtime generation | Re-enabled and trained with text cross-entropy |
| Action conditioning | Episode task | Active low-level subtask plus normalized robot state |
| Training targets | Flow-matching actions | Flow actions, recipe-selected text, and optional FAST action tokens |
| Dataset requirement | Standard images, state, actions, and task | The same fields plus language annotations for every language capability you train |
| Rollout | Direct task-to-action policy | Hierarchical task → subtask → action loop, with optional memory and VQA |
Pi052 can initialize from a Pi05 checkpoint. The policy architecture remains
compatible, while Pi052 builds its own processors so recipe labels and FAST
labels are not silently replaced by the Pi05 processor stack.
## Install
Install LeRobot with the PI dependencies:
```bash
git clone https://github.com/huggingface/lerobot.git
cd lerobot
python -m venv .venv
source .venv/bin/activate
pip install -e ".[pi]"
```
The `pi` extra includes the PaliGemma/FAST dependencies. Install
`liger-kernel` for the supported fused training kernels; optional FlashRT
backends also require the Hugging Face `kernels` package and a supported CUDA
GPU.
## Prepare language-annotated data
Pi052 does not infer supervised subtasks from a normal LeRobot dataset during
training. The dataset must contain the language targets used by the selected
recipe in the optional `language_persistent` and `language_events` columns.
At minimum, annotate a continuous `subtask` timeline so each training frame has
an active low-level instruction. Add `memory`, VQA, interjections, and speech
annotations only if the recipe trains those capabilities.
The provided recipes are:
| Recipe | Required annotations | Trains |
| ------------------------------------- | ----------------------------------------------------------------------- | ------------------------------------------------------------------ |
| `recipes/subtask.yaml` | `subtask` | Subtask prediction and subtask-conditioned actions |
| `recipes/subtask_joint.yaml` | `subtask` | Paper-style joint sequence: subtask text and actions in one sample |
| `recipes/subtask_mem.yaml` | `subtask`, `memory` | Subtasks, actions, and memory updates |
| `recipes/subtask_mem_vqa_speech.yaml` | `subtask`, `memory`, `vqa`; interjection/speech rows for those branches | Subtasks, actions, memory, VQA, and spoken replies |
The blend recipes factorize training into separate high-level (task → subtask)
and low-level (subtask → actions) samples, matching how inference decomposes
π(a|o, subtask)·π(subtask|o, task). `recipes/subtask_joint.yaml` instead uses
the π0.5 paper's single-sequence layout — the supervised subtask span is
attended causally and conditions the FAST and flow losses in the same forward.
Checkpoints trained with the joint recipe must set
`--policy.joint_subtask_conditioning=true` at inference so the flow prefix
rebuilds the same layout (task turn with state, then the generated subtask as a
causal assistant turn); leave it `false` for the blend recipes.
Use `lerobot-annotate` to generate these columns. The repository includes a
Hugging Face Jobs launcher that you can edit for your source and destination
datasets. For a local annotation run, first install
`pip install -e ".[annotations]"`:
```bash
HF_TOKEN=hf_... uv run python examples/annotations/run_hf_job.py
```
Before a long training run, inspect several episodes and verify that subtasks
are temporally correct and cover the full demonstration. See
[Annotation Pipeline](./annotation_pipeline) for generation and validation, and
[Language Columns and Recipes](./language_and_recipes) for the schema and
recipe resolver.
<Tip>
If a dataset has no language columns, recipe rendering becomes a no-op and
Pi052 falls back to the plain Pi05 prompt path. This is useful for
compatibility but does not train the language planner.
</Tip>
## Train Pi052
This example initializes Pi052 from the public Pi05 base checkpoint and trains
the default subtask-and-memory recipe:
```bash
lerobot-train \
--dataset.repo_id=${HF_USER}/my_language_annotated_dataset \
--policy.type=pi052 \
--policy.pretrained_path=lerobot/pi05_base \
--policy.recipe_path=recipes/subtask_mem.yaml \
--policy.dtype=bfloat16 \
--policy.device=cuda \
--policy.freeze_vision_encoder=false \
--policy.gradient_checkpointing=true \
--batch_size=8 \
--steps=30000 \
--output_dir=outputs/pi052 \
--job_name=pi052 \
--wandb.enable=true
```
For subtask-only data, change the recipe to `recipes/subtask.yaml` and disable
memory during rollout. Start with a small run and confirm that W&B examples show
the expected prompt, text target, and action endpoints before scaling up.
### Main training controls
| Option | Default | Purpose |
| ----------------------------------- | -------------------------: | ------------------------------------------------------------------- |
| `policy.recipe_path` | `recipes/subtask_mem.yaml` | Selects the language/action objective mixture |
| `policy.text_loss_weight` | `1.0` | Language-head cross-entropy weight; `0` disables text training |
| `policy.flow_loss_weight` | `10.0` | Continuous action flow-loss weight |
| `policy.enable_fast_action_loss` | `true` | Adds discrete FAST action-token supervision |
| `policy.fast_action_loss_weight` | `1.0` | FAST cross-entropy weight |
| `policy.knowledge_insulation` | `true` | Blocks action-loss gradients through the VLM K/V path |
| `policy.flow_num_repeats` | `5` | Reuses one VLM prefix for independent denoising targets |
| `policy.lm_head_lr_scale` | `1.0` | Scales language-head learning rate; `1.0` uses the base rate |
| `policy.fast_skip_tokens` | `1152` | FAST id offset; skips `<seg>`+`<loc>` so VQA and FAST never collide |
| `policy.joint_subtask_conditioning` | `false` | Rebuilds the joint-sequence prefix at inference (see recipes) |
`fast_skip_tokens=1152` places FAST codes below PaliGemma's `<loc>` range.
openpi's pi0-FAST convention is `128` (FAST occupies the `<loc>` ids); use that
value only to stay weight-compatible with checkpoints trained that way, and
avoid combining it with the VQA recipe, whose `<loc>` targets would share
embedding rows with FAST codes.
The loss weights are starting points, not dataset-independent constants. Track
flow loss and text/FAST losses separately, and inspect generated subtasks rather
than selecting a checkpoint from total loss alone.
### Dataset-specific FAST tokenizer
The universal FAST tokenizer works out of the box. For a large or
embodiment-specific dataset, Pi052 can fit and cache a tokenizer on normalized
actions before training:
```bash
lerobot-train \
... \
--policy.auto_fit_fast_tokenizer=true \
--policy.fast_tokenizer_fit_samples=4096
```
The fit runs once per dataset/tokenizer configuration. Keep
`auto_fit_fast_tokenizer=false` when you do not want the extra preprocessing
pass.
## Training performance
Pi052 uses optimized training paths by default:
- batches repeated flow targets and suffix projections instead of replaying
small operations in Python;
- caches constant action masks and computes RoPE positions once per forward;
- selects the text/FAST cross-entropy implementation from target shape and
sparsity;
- skips the mathematically dead VLM/vision backward on knowledge-insulated,
flow-only batches;
- uses native non-reentrant SigLIP layer checkpointing when gradient
checkpointing is enabled; and
- retains the Liger RoPE/GeGLU kernels while avoiding the slower LayerNorm
patch at SigLIP shapes.
Optional training backends are disabled by default:
| Option | When to try it |
| -------------------------------------- | ------------------------------------------------------------------------------------------------- |
| `policy.use_flashrt_adarms=true` | Fused adaptive RMSNorm and gated residuals on supported CUDA GPUs |
| `policy.use_compiled_text_ce=true` | Compiled materialized-logit CE buckets |
| `policy.use_compiled_vision=true` | Compiled vision only when the vision pass has no gradients |
| `policy.use_flex_attention=true` | Profiled CUDA setups with knowledge insulation and `flow_num_repeats > 1`; otherwise SDPA is used |
| `policy.use_manual_attention=true` | Explicitly profiled shapes where materialized attention is faster |
| `policy.manual_attention_scope=action` | Restricts manual attention to action queries |
Do not enable every backend blindly. Flex and manual attention are mutually
exclusive, and attention/AdaRMS alternatives require knowledge insulation.
The benchmark-best configuration used compiled text CE and FlashRT AdaRMS,
with Flex/manual attention and compiled vision disabled.
### Reported training benchmarks
These benchmarks measure complete optimizer steps with three real camera
inputs, BF16 transformer/action execution, FP32 vision, fused AdamW, and no
video decoding or network I/O. Results vary with GPU, batch shape, annotation
mixture, and checkpointing:
| Workload | RTX PRO 6000 Blackwell | A100 80 GB |
| -------------------------- | -------------------------: | -------------------------: |
| Full flow + text, batch 1 | 4.75× vs checkpointing off | 3.33× vs checkpointing off |
| Full flow + text, batch 8 | 2.16× vs checkpointing off | 1.66× vs checkpointing off |
| Full flow + text, batch 64 | 1.24× vs checkpointing on | 1.15× vs checkpointing on |
| Flow-only, batch 1 | 3.70× vs checkpointing off | 3.58× vs checkpointing off |
| Flow-only, batch 64 | 3.76× vs checkpointing on | 3.61× vs checkpointing on |
On those 80 GB GPUs, full training was fastest without gradient checkpointing
through batch 8, then required checkpointing at batch 16 and above. Treat that
as a tuning rule to test on your hardware, not a universal threshold. Flow-only
means both text and FAST supervision are disabled; it is useful for action-only
ablation or post-training but does not learn the language runtime.
## Inference performance
Pi052 has two inference loops, and both avoid repeatedly encoding the expensive
multimodal prefix:
1. **Action denoising** encodes the image/language prefix once, reuses its KV
cache across flow steps, precomputes the timestep schedule on-device, and
crops temporary suffix K/V instead of cloning the prefix cache.
2. **Language decoding** uses autoregressive KV caching, so each new token only
processes the sampled token against cached image/language keys instead of
rerunning the full prefix.
The runtime also runs language and actions at different rates. Increase
`--subtask_chunks_per_gen` when a subtask remains valid across several action
chunks, lower `--high_level_hz`, or use `--direct_subtask` to bypass language
generation entirely. These settings reduce compute but also slow replanning.
`--fp8` enables the optional FlashRT inference MLP swap on supported CUDA GPUs.
It calibrates on the first observation and falls back to BF16 when unavailable;
because FP8 can change outputs slightly, validate task success before using it
for production rollouts.
## Run a checkpoint
RoboCasa:
```bash
MUJOCO_GL=egl lerobot-rollout \
--policy.path=lerobot/pi052_robocasa \
--sim --sim.task=CloseFridge --sim.split=pretrain \
--task="close the fridge" \
--disable_memory \
--sim.render_size=384 \
--sim.views=robot0_agentview_left,robot0_eye_in_hand,robot0_agentview_right \
--mode=action --ctrl_hz=20
```
Open `http://localhost:8010` for the live view. Without
`--sim.direct_subtask`, Pi052 generates the low-level subtask; with it, each
prompt becomes the action policy's subtask directly.
The same runtime supports real robots. See [Interactive language
control](./inference#interactive-language-control) for the real-arm command,
safety behavior, and runtime controls.
## Troubleshooting
- **No text loss or generated subtasks:** confirm the selected recipe can bind
the annotations on sampled frames and that `policy.text_loss_weight > 0`.
- **Subtasks look plausible but actions fail:** verify subtask boundaries,
normalized state/action statistics, and that low-level recipe samples are
present.
- **Text collapses to repeated or location tokens:** inspect text-target
coverage, language-head learning rate, and the balance between flow, FAST,
and text losses.
- **Out of memory:** reduce batch size first, then enable gradient
checkpointing. Do not enable compiled or alternative attention backends
without profiling their memory on your camera count.
- **Slow rollout:** separate action latency from language latency, then tune
`--subtask_chunks_per_gen`, `--high_level_hz`, and the number of flow
inference steps.
docs/source/pi0fast.mdx
+15 -9
View File
@@ -109,15 +109,21 @@ lerobot-train \
### Key Training Parameters
| Parameter | Description | Default |
| -------------------------------------- | -------------------------------------------------- | ------------------------------- |
| `--policy.gradient_checkpointing=true` | Reduces memory usage significantly during training | `false` |
| `--policy.dtype=bfloat16` | Use mixed precision training for efficiency | `float32` |
| `--policy.chunk_size` | Number of action steps to predict (action horizon) | `50` |
| `--policy.n_action_steps` | Number of action steps to execute | `50` |
| `--policy.max_action_tokens` | Maximum number of FAST tokens per action chunk | `256` |
| `--policy.action_tokenizer_name` | FAST tokenizer to use | `lerobot/fast-action-tokenizer` |
| `--policy.compile_model=true` | Enable torch.compile for faster training | `false` |
| Parameter | Description | Default |
| --------------------------------------- | -------------------------------------------------- | ------------------------------- |
| `--policy.gradient_checkpointing=true` | Reduces memory usage significantly during training | `false` |
| `--policy.dtype=bfloat16` | Use mixed precision training for efficiency | `float32` |
| `--policy.chunk_size` | Number of action steps to predict (action horizon) | `50` |
| `--policy.n_action_steps` | Number of action steps to execute | `50` |
| `--policy.max_action_tokens` | Maximum number of FAST tokens per action chunk | `256` |
| `--policy.action_tokenizer_name` | FAST tokenizer to use | `lerobot/fast-action-tokenizer` |
| `--policy.auto_fit_fast_tokenizer=true` | Fit and cache a tokenizer for the training dataset | `false` |
| `--policy.compile_model=true` | Enable torch.compile for faster training | `false` |
Set `--policy.auto_fit_fast_tokenizer=true` to sample action chunks from the
training dataset and cache a fitted tokenizer under
`~/.cache/lerobot/fast_tokenizers`. This also works when fine-tuning with
`--policy.path`; leave it disabled to retain the checkpoint's tokenizer.
## Inference
pyproject.toml
+2 -1
View File
@@ -150,6 +150,7 @@ pygame-dep = ["pygame>=2.5.1,<2.7.0"]
# There is no cmeel-urdfdom 5.x; <5 selects the 4.x ABI the placo/pin wheels are built against.
placo-dep = ["placo>=0.9.6,<0.9.16", "cmeel-urdfdom>=4,<5", "cmeel-tinyxml2<11"]
transformers-dep = ["transformers>=5.4.0,<5.6.0"]
sentencepiece-dep = ["sentencepiece>=0.2.0,<0.3.0"] # FAST action tokenizer backend (pi052, pi0_fast)
grpcio-dep = ["grpcio>=1.73.1,<2.0.0", "protobuf>=6.31.1,<8.0.0"]
accelerate-dep = ["accelerate>=1.14.0,<2.0.0"]
can-dep = ["python-can>=4.2.0,<5.0.0"]
@@ -212,7 +213,7 @@ wallx = [
"torchdiffeq>=0.2.4,<0.3.0",
"lerobot[qwen-vl-utils-dep]",
]
pi = ["lerobot[transformers-dep]", "lerobot[scipy-dep]"]
pi = ["lerobot[transformers-dep]", "lerobot[scipy-dep]", "lerobot[sentencepiece-dep]"]
molmoact2 = ["lerobot[transformers-dep]", "lerobot[peft-dep]", "lerobot[scipy-dep]"]
smolvla = ["lerobot[transformers-dep]", "num2words>=0.5.14,<0.6.0", "lerobot[accelerate-dep]"]
multi_task_dit = ["lerobot[transformers-dep]", "lerobot[diffusers-dep]"]
scripts/backfill_pi052_action_tokenizer.py
+227
View File
@@ -0,0 +1,227 @@
#!/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.
"""Audit or backfill checkpoint-local FAST artifacts for PI052 model repositories."""
from __future__ import annotations
import argparse
import hashlib
import io
import json
from pathlib import Path, PurePosixPath
from typing import Any
from huggingface_hub import CommitOperationAdd, HfApi, hf_hub_download
DEFAULT_REPOSITORIES = (
"pepijn223/pi052_atomic4_01_baseline",
"pepijn223/pi052_atomic4_02_lr_1e5",
"pepijn223/pi052_atomic4_03_recipe_50_50",
"pepijn223/pi052_atomic4_04_flow_weight_10",
"pepijn223/pi052_atomic4_05_flow_repeat_1",
"pepijn223/pi052_atomic4_06_ki_off",
)
CHECKPOINT_DIRECTORIES = (
"",
"checkpoints/003000/pretrained_model",
"checkpoints/006000/pretrained_model",
"checkpoints/009000/pretrained_model",
"checkpoints/012000/pretrained_model",
)
TOKENIZER_DIRECTORY = "action_tokenizer"
def artifact_fingerprint(files: list[tuple[str, bytes]]) -> str:
digest = hashlib.sha256()
for relative_path, content in sorted(files):
encoded_path = relative_path.encode()
digest.update(len(encoded_path).to_bytes(8, "big"))
digest.update(encoded_path)
digest.update(len(content).to_bytes(8, "big"))
digest.update(content)
return digest.hexdigest()
def tokenizer_files(tokenizer_path: Path) -> list[tuple[str, Path]]:
return [
(path.relative_to(tokenizer_path).as_posix(), path)
for path in sorted(tokenizer_path.rglob("*"))
if path.is_file()
]
def _repo_path(directory: str, filename: str) -> str:
return (PurePosixPath(directory) / filename).as_posix() if directory else filename
def _download_json(repo_id: str, path_in_repo: str, revision: str | None = None) -> dict[str, Any]:
path = hf_hub_download(repo_id, path_in_repo, repo_type="model", revision=revision)
return json.loads(Path(path).read_text())
def make_portable_preprocessor(config: dict[str, Any]) -> dict[str, Any]:
config = json.loads(json.dumps(config))
action_steps = [
step for step in config["steps"] if step.get("registry_name") == "action_tokenizer_processor"
]
if len(action_steps) != 1:
raise ValueError(f"Expected one action tokenizer step, found {len(action_steps)}")
action_step = action_steps[0]
action_step["config"]["action_tokenizer_name"] = TOKENIZER_DIRECTORY
action_step["artifacts"] = {"action_tokenizer_name": TOKENIZER_DIRECTORY}
recipe_steps = [
step for step in config["steps"] if step.get("registry_name") == "render_messages_processor"
]
if len(recipe_steps) != 1 or not recipe_steps[0].get("config", {}).get("recipe"):
raise ValueError("PI052 preprocessor does not contain an embedded training recipe")
return config
def _json_operation(path_in_repo: str, content: dict[str, Any]) -> CommitOperationAdd:
serialized = (json.dumps(content, indent=2) + "\n").encode()
return CommitOperationAdd(path_in_repo=path_in_repo, path_or_fileobj=io.BytesIO(serialized))
def prepare_operations(
repo_id: str,
tokenizer_path: Path,
revision: str | None = None,
) -> list[CommitOperationAdd]:
operations: list[CommitOperationAdd] = []
files = tokenizer_files(tokenizer_path)
for directory in CHECKPOINT_DIRECTORIES:
preprocessor_path = _repo_path(directory, "policy_preprocessor.json")
operations.append(
_json_operation(
preprocessor_path,
make_portable_preprocessor(_download_json(repo_id, preprocessor_path, revision)),
)
)
for relative_path, local_path in files:
operations.append(
CommitOperationAdd(
path_in_repo=_repo_path(
directory,
f"{TOKENIZER_DIRECTORY}/{relative_path}",
),
path_or_fileobj=str(local_path),
)
)
return operations
def audit_repository(
api: HfApi,
repo_id: str,
expected_tokenizer_fingerprint: str,
revision: str | None = None,
) -> None:
info = api.model_info(repo_id, revision=revision)
repository_files = {sibling.rfilename for sibling in info.siblings or []}
for directory in CHECKPOINT_DIRECTORIES:
preprocessor_path = _repo_path(directory, "policy_preprocessor.json")
policy_config_path = _repo_path(directory, "config.json")
postprocessor_path = _repo_path(directory, "policy_postprocessor.json")
for required_path in (preprocessor_path, policy_config_path, postprocessor_path):
if required_path not in repository_files:
raise FileNotFoundError(f"{repo_id}@{revision or 'main'} is missing {required_path}")
preprocessor = _download_json(repo_id, preprocessor_path, revision)
portable_preprocessor = make_portable_preprocessor(preprocessor)
if preprocessor != portable_preprocessor:
raise ValueError(f"{repo_id}:{preprocessor_path} is not portable")
normalizer_steps = [
step for step in preprocessor["steps"] if step.get("registry_name") == "normalizer_processor"
]
if len(normalizer_steps) != 1 or "state_file" not in normalizer_steps[0]:
raise ValueError(f"{repo_id}:{preprocessor_path} is missing normalizer state metadata")
normalizer_path = _repo_path(directory, normalizer_steps[0]["state_file"])
if normalizer_path not in repository_files:
raise FileNotFoundError(f"{repo_id} is missing {normalizer_path}")
remote_tokenizer_files: list[tuple[str, bytes]] = []
for relative_path in _tokenizer_relative_paths(repository_files, directory):
path_in_repo = _repo_path(directory, f"{TOKENIZER_DIRECTORY}/{relative_path}")
downloaded = hf_hub_download(repo_id, path_in_repo, repo_type="model", revision=revision)
remote_tokenizer_files.append((relative_path, Path(downloaded).read_bytes()))
fingerprint = artifact_fingerprint(remote_tokenizer_files)
if fingerprint != expected_tokenizer_fingerprint:
raise ValueError(
f"{repo_id}:{_repo_path(directory, TOKENIZER_DIRECTORY)} fingerprint "
f"{fingerprint} != {expected_tokenizer_fingerprint}"
)
def _tokenizer_relative_paths(repository_files: set[str], directory: str) -> list[str]:
prefix = _repo_path(directory, TOKENIZER_DIRECTORY).rstrip("/") + "/"
paths = sorted(path.removeprefix(prefix) for path in repository_files if path.startswith(prefix))
if not paths:
raise FileNotFoundError(f"Missing tokenizer artifact directory {prefix.rstrip('/')}")
return paths
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--tokenizer-path", type=Path, required=True)
parser.add_argument("--repo-id", action="append", dest="repo_ids")
parser.add_argument("--revision")
parser.add_argument("--dry-run", action="store_true")
parser.add_argument("--audit-only", action="store_true")
return parser.parse_args()
def main() -> None:
args = parse_args()
tokenizer_path = args.tokenizer_path.resolve()
if not tokenizer_path.is_dir():
raise FileNotFoundError(f"Tokenizer directory does not exist: {tokenizer_path}")
files = tokenizer_files(tokenizer_path)
fingerprint = artifact_fingerprint([(relative_path, path.read_bytes()) for relative_path, path in files])
api = HfApi()
repositories = tuple(args.repo_ids or DEFAULT_REPOSITORIES)
print(f"Tokenizer fingerprint: {fingerprint}")
for repo_id in repositories:
if args.audit_only:
audit_repository(api, repo_id, fingerprint, args.revision)
print(f"AUDIT OK {repo_id}@{args.revision or 'main'}")
continue
operations = prepare_operations(repo_id, tokenizer_path, args.revision)
if args.dry_run:
print(f"DRY RUN {repo_id}: {len(operations)} files")
for operation in operations:
print(f" {operation.path_in_repo}")
continue
commit = api.create_commit(
repo_id=repo_id,
repo_type="model",
operations=operations,
commit_message="Embed fitted FAST tokenizer for portable PI052 checkpoints",
revision=args.revision,
)
audit_repository(api, repo_id, fingerprint, commit.oid)
print(f"BACKFILLED {repo_id}@{commit.oid}")
if __name__ == "__main__":
main()
src/lerobot/optim/optimizers.py
+2
View File
@@ -104,6 +104,8 @@ class AdamWConfig(OptimizerConfig):
eps: float = 1e-8
weight_decay: float = 1e-2
grad_clip_norm: float = 10.0
foreach: bool | None = None
fused: bool | None = None
def build(self, params: OptimizerParams) -> torch.optim.Optimizer:
kwargs = asdict(self)
src/lerobot/policies/__init__.py
+2
View File
@@ -28,6 +28,7 @@ from .multi_task_dit.configuration_multi_task_dit import MultiTaskDiTConfig as M
from .pi0.configuration_pi0 import PI0Config as PI0Config
from .pi0_fast.configuration_pi0_fast import PI0FastConfig as PI0FastConfig
from .pi05.configuration_pi05 import PI05Config as PI05Config
from .pi052.configuration_pi052 import PI052Config as PI052Config
from .pretrained import PreTrainedPolicy as PreTrainedPolicy
from .smolvla.configuration_smolvla import SmolVLAConfig as SmolVLAConfig
from .tdmpc.configuration_tdmpc import TDMPCConfig as TDMPCConfig
@@ -56,6 +57,7 @@ __all__ = [
"PI0Config",
"PI0FastConfig",
"PI05Config",
"PI052Config",
"SmolVLAConfig",
"TDMPCConfig",
"VLAJEPAConfig",
src/lerobot/policies/factory.py
+37 -2
View File
@@ -137,6 +137,12 @@ class ProcessorConfigKwargs(TypedDict, total=False):
preprocessor_overrides: dict[str, Any] | None
postprocessor_overrides: dict[str, Any] | None
dataset_stats: dict[str, dict[str, torch.Tensor]] | None
# Dataset source used by policies that optionally fit processor artifacts.
dataset_repo_id: str | None
dataset_root: str | None
dataset_revision: str | None
dataset_episodes: list[int] | None
dataset_exclude_episodes: list[int] | None
dataset_meta: Any | None
@@ -171,6 +177,10 @@ def make_pre_post_processors(
ValueError: If no processor factory exists for the given policy configuration type.
"""
if pretrained_path:
# Register the PI052-only stateful tokenizer step before deserializing its pipeline.
if policy_cfg.type == "pi052":
from .pi052 import processor_pi052 as _processor_pi052 # noqa: F401
if isinstance(policy_cfg, GrootConfig):
from .groot.processor_groot import make_groot_pre_post_processors_from_pretrained
@@ -190,12 +200,29 @@ def make_pre_post_processors(
),
)
preprocessor_overrides = dict(kwargs.get("preprocessor_overrides") or {})
if policy_cfg.type == "pi0_fast" and getattr(policy_cfg, "auto_fit_fast_tokenizer", False):
from .pi052.fit_fast_tokenizer import resolve_fast_tokenizer
fitted_tokenizer = resolve_fast_tokenizer(
policy_cfg,
kwargs.get("dataset_repo_id"),
kwargs.get("dataset_root"),
kwargs.get("dataset_stats"),
kwargs.get("dataset_revision"),
kwargs.get("dataset_episodes"),
kwargs.get("dataset_exclude_episodes"),
)
tokenizer_overrides = dict(preprocessor_overrides.get("action_tokenizer_processor") or {})
tokenizer_overrides["action_tokenizer_name"] = fitted_tokenizer
preprocessor_overrides["action_tokenizer_processor"] = tokenizer_overrides
preprocessor = PolicyProcessorPipeline.from_pretrained(
pretrained_model_name_or_path=pretrained_path,
config_filename=kwargs.get(
"preprocessor_config_filename", f"{POLICY_PREPROCESSOR_DEFAULT_NAME}.json"
),
overrides=kwargs.get("preprocessor_overrides", {}),
overrides=preprocessor_overrides,
to_transition=batch_to_transition,
to_output=transition_to_batch,
revision=pretrained_revision,
@@ -227,6 +254,11 @@ def make_pre_post_processors(
config=policy_cfg,
dataset_stats=kwargs.get("dataset_stats"),
dataset_meta=kwargs.get("dataset_meta"),
dataset_repo_id=kwargs.get("dataset_repo_id"),
dataset_root=kwargs.get("dataset_root"),
dataset_revision=kwargs.get("dataset_revision"),
episodes=kwargs.get("dataset_episodes"),
exclude_episodes=kwargs.get("dataset_exclude_episodes"),
)
@@ -424,6 +456,7 @@ def _make_processors_from_policy_config(
config: PreTrainedConfig,
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
dataset_meta: Any | None = None,
**optional_kwargs: Any,
) -> tuple[Any, Any]:
"""Create pre- and post-processors from a policy configuration using dynamic imports.
@@ -459,7 +492,9 @@ def _make_processors_from_policy_config(
function = getattr(module, function_name, None)
if function is None:
raise ValueError(f"Processor for policy type '{policy_type}' is not implemented.")
parameters = inspect.signature(function).parameters
call_kwargs: dict[str, Any] = {"dataset_stats": dataset_stats}
if "dataset_meta" in inspect.signature(function).parameters:
if "dataset_meta" in parameters:
call_kwargs["dataset_meta"] = dataset_meta
call_kwargs.update({name: value for name, value in optional_kwargs.items() if name in parameters})
return function(config, **call_kwargs)
src/lerobot/policies/pi05/modeling_pi05.py
+387 -121
View File
@@ -15,21 +15,26 @@
# limitations under the License.
import builtins
import json
import logging
import math
from collections import deque
from pathlib import Path
from typing import TYPE_CHECKING, Literal, TypedDict, Unpack
import torch
import torch.nn.functional as F # noqa: N812
from safetensors.torch import load_file
from torch import Tensor, nn
from lerobot.utils.import_utils import _transformers_available, require_package
# Conditional import for type checking and lazy loading
if TYPE_CHECKING or _transformers_available:
from transformers.cache_utils import DynamicCache
from transformers.models.auto import CONFIG_MAPPING
from transformers.models.gemma import modeling_gemma
from transformers.utils import cached_file
from ..pi_gemma import (
PaliGemmaForConditionalGenerationWithPiGemma,
@@ -39,27 +44,22 @@ if TYPE_CHECKING or _transformers_available:
)
else:
CONFIG_MAPPING = None
DynamicCache = None
modeling_gemma = None
PiGemmaForCausalLM = None
_gated_residual = None
layernorm_forward = None
PaliGemmaForConditionalGenerationWithPiGemma = None
cached_file = None
from lerobot.configs import PreTrainedConfig
from lerobot.utils.constants import (
ACTION,
OBS_LANGUAGE_ATTENTION_MASK,
OBS_LANGUAGE_TOKENS,
OPENPI_ATTENTION_MASK_VALUE,
)
from ..common.flow_matching import euler_integrate, sample_noise, sample_time_beta
from ..common.vla_utils import (
clone_past_key_values,
create_sinusoidal_pos_embedding,
make_att_2d_masks,
pad_vector,
prepare_attention_masks_4d,
resize_with_pad_torch,
)
from ..common.flow_matching import sample_noise, sample_time_beta
from ..pretrained import PreTrainedPolicy, T
from ..rtc.modeling_rtc import RTCProcessor
from .configuration_pi05 import DEFAULT_IMAGE_SIZE, PI05Config
@@ -71,6 +71,251 @@ class ActionSelectKwargs(TypedDict, total=False):
execution_horizon: int | None
_SAFETENSORS_FILE = "model.safetensors"
_SAFETENSORS_INDEX = "model.safetensors.index.json"
def _resolve_weight_files(
pretrained_name_or_path: str | Path,
*,
force_download: bool,
resume_download: bool | None,
proxies: dict | None,
token: str | bool | None,
cache_dir: str | Path | None,
local_files_only: bool,
revision: str | None,
) -> list[Path]:
model_id = str(pretrained_name_or_path)
local_dir = Path(model_id)
load_kwargs = {
"revision": revision,
"cache_dir": cache_dir,
"force_download": force_download,
"resume_download": resume_download,
"proxies": proxies,
"token": token,
"local_files_only": local_files_only,
}
if local_dir.is_dir():
index_path = local_dir / _SAFETENSORS_INDEX
single_path = local_dir / _SAFETENSORS_FILE
else:
resolved_index = cached_file(
model_id,
_SAFETENSORS_INDEX,
_raise_exceptions_for_missing_entries=False,
**load_kwargs,
)
index_path = Path(resolved_index) if resolved_index is not None else None
single_path = None
if index_path is None:
resolved_file = cached_file(model_id, _SAFETENSORS_FILE, **load_kwargs)
single_path = Path(resolved_file) if resolved_file is not None else None
if index_path is None or not index_path.is_file():
if single_path is None or not single_path.is_file():
raise FileNotFoundError(f"No {_SAFETENSORS_FILE} found in {model_id!r}.")
return [single_path]
index = json.loads(index_path.read_text())
shard_names = sorted(set(index.get("weight_map", {}).values()))
if not shard_names:
raise ValueError(f"Invalid safetensors index without a weight_map: {index_path}")
if local_dir.is_dir():
files = [local_dir / name for name in shard_names]
else:
files = []
for name in shard_names:
resolved_file = cached_file(model_id, name, **load_kwargs)
if resolved_file is None:
raise FileNotFoundError(f"Checkpoint shard {name!r} not found in {model_id!r}.")
files.append(Path(resolved_file))
missing = [str(path) for path in files if not path.is_file()]
if missing:
raise FileNotFoundError(f"Missing checkpoint shards: {missing}")
return files
def _load_weight_files(files: list[Path]) -> dict[str, Tensor]:
state_dict: dict[str, Tensor] = {}
for path in files:
shard = load_file(path)
overlap = state_dict.keys() & shard.keys()
if overlap:
raise ValueError(f"Duplicate checkpoint keys in {path}: {sorted(overlap)[:5]}")
state_dict.update(shard)
return state_dict
def get_safe_dtype(target_dtype, device_type):
"""Get a safe dtype for the given device type."""
if device_type == "mps" and target_dtype == torch.float64:
return torch.float32
if device_type == "cpu":
# CPU doesn't support bfloat16, use float32 instead
if target_dtype == torch.bfloat16:
return torch.float32
if target_dtype == torch.float64:
return torch.float64
return target_dtype
def create_sinusoidal_pos_embedding( # see openpi `create_sinusoidal_pos_embedding` (exact copy)
time: torch.Tensor, dimension: int, min_period: float, max_period: float, device="cpu"
) -> Tensor:
"""Computes sine-cosine positional embedding vectors for scalar positions."""
if dimension % 2 != 0:
raise ValueError(f"dimension ({dimension}) must be divisible by 2")
if time.ndim != 1:
raise ValueError("The time tensor is expected to be of shape `(batch_size, )`.")
dtype = get_safe_dtype(torch.float64, device.type)
fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=dtype, device=device)
period = min_period * (max_period / min_period) ** fraction
# Compute the outer product
scaling_factor = 1.0 / period * 2 * math.pi
sin_input = scaling_factor[None, :] * time[:, None]
return torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1)
def sample_beta(alpha, beta, bsize, device): # see openpi `sample_beta` (exact copy)
# Beta sampling uses _sample_dirichlet which isn't implemented for MPS, so sample on CPU
alpha_t = torch.tensor(alpha, dtype=torch.float32)
beta_t = torch.tensor(beta, dtype=torch.float32)
dist = torch.distributions.Beta(alpha_t, beta_t)
return dist.sample((bsize,)).to(device)
def make_att_2d_masks(pad_masks, att_masks): # see openpi `make_att_2d_masks` (exact copy)
"""Copied from big_vision.
Tokens can attend to valid inputs tokens which have a cumulative mask_ar
smaller or equal to theirs. This way `mask_ar` int[B, N] can be used to
setup several types of attention, for example:
[[1 1 1 1 1 1]]: pure causal attention.
[[0 0 0 1 1 1]]: prefix-lm attention. The first 3 tokens can attend between
themselves and the last 3 tokens have a causal attention. The first
entry could also be a 1 without changing behaviour.
[[1 0 1 0 1 0 0 1 0 0]]: causal attention between 4 blocks. Tokens of a
block can attend all previous blocks and all tokens on the same block.
Args:
input_mask: bool[B, N] true if its part of the input, false if padding.
mask_ar: int32[B, N] mask that's 1 where previous tokens cannot depend on
it and 0 where it shares the same attention mask as the previous token.
"""
if att_masks.ndim != 2:
raise ValueError(att_masks.ndim)
if pad_masks.ndim != 2:
raise ValueError(pad_masks.ndim)
cumsum = torch.cumsum(att_masks, dim=1)
att_2d_masks = cumsum[:, None, :] <= cumsum[:, :, None]
pad_2d_masks = pad_masks[:, None, :] * pad_masks[:, :, None]
return att_2d_masks & pad_2d_masks
def clone_past_key_values(past_key_values):
"""Clone the DynamicCache returned by prefix prefill for compiled denoising."""
return DynamicCache(
tuple(
(keys.clone(), values.clone(), sliding_window) for keys, values, sliding_window in past_key_values
)
)
def pad_vector(vector, new_dim):
"""Pad the last dimension of a vector to new_dim with zeros.
Can be (batch_size x sequence_length x features_dimension)
or (batch_size x features_dimension)
"""
if vector.shape[-1] >= new_dim:
return vector
return F.pad(vector, (0, new_dim - vector.shape[-1]))
def resize_with_pad_torch( # see openpi `resize_with_pad_torch` (exact copy)
images: torch.Tensor,
height: int,
width: int,
mode: str = "bilinear",
) -> torch.Tensor:
"""PyTorch version of resize_with_pad. Resizes an image to a target height and width without distortion
by padding with black. If the image is float32, it must be in the range [-1, 1].
Args:
images: Tensor of shape [*b, h, w, c] or [*b, c, h, w]
height: Target height
width: Target width
mode: Interpolation mode ('bilinear', 'nearest', etc.)
Returns:
Resized and padded tensor with same shape format as input
"""
# Check if input is in channels-last format [*b, h, w, c] or channels-first [*b, c, h, w]
if images.shape[-1] <= 4: # Assume channels-last format
channels_last = True
if images.dim() == 3:
images = images.unsqueeze(0) # Add batch dimension
images = images.permute(0, 3, 1, 2) # [b, h, w, c] -> [b, c, h, w]
else:
channels_last = False
if images.dim() == 3:
images = images.unsqueeze(0) # Add batch dimension
batch_size, channels, cur_height, cur_width = images.shape
# Calculate resize ratio
ratio = max(cur_width / width, cur_height / height)
resized_height = int(cur_height / ratio)
resized_width = int(cur_width / ratio)
# Resize
resized_images = F.interpolate(
images,
size=(resized_height, resized_width),
mode=mode,
align_corners=False if mode == "bilinear" else None,
)
# Handle dtype-specific clipping
if images.dtype == torch.uint8:
resized_images = torch.round(resized_images).clamp(0, 255).to(torch.uint8)
elif images.dtype == torch.float32:
resized_images = resized_images.clamp(0.0, 1.0)
else:
raise ValueError(f"Unsupported image dtype: {images.dtype}")
# Calculate padding
pad_h0, remainder_h = divmod(height - resized_height, 2)
pad_h1 = pad_h0 + remainder_h
pad_w0, remainder_w = divmod(width - resized_width, 2)
pad_w1 = pad_w0 + remainder_w
# Pad
constant_value = 0 if images.dtype == torch.uint8 else 0.0
padded_images = F.pad(
resized_images,
(pad_w0, pad_w1, pad_h0, pad_h1), # left, right, top, bottom
mode="constant",
value=constant_value,
)
# Convert back to original format if needed
if channels_last:
padded_images = padded_images.permute(0, 2, 3, 1) # [b, c, h, w] -> [b, h, w, c]
return padded_images
# Define the complete layer computation function for gradient checkpointing
def compute_layer_complete(inputs_embeds, attention_mask, position_ids, adarms_cond, layers, rotary_emb):
query_states = []
@@ -401,6 +646,12 @@ class PaliGemmaWithExpertModel(
class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch`
"""Core PI05 PyTorch model."""
use_hf_vision_checkpointing_api = False
checkpoint_vision_embeddings = True
use_typed_attention_masks = False
use_on_device_suffix_mask = False
precompute_denoise_times = False
def __init__(self, config: PI05Config, rtc_processor: RTCProcessor | None = None):
super().__init__()
self.config = config
@@ -444,7 +695,11 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch`
"""Enable gradient checkpointing for memory optimization."""
self.gradient_checkpointing_enabled = True
self.paligemma_with_expert.paligemma.model.language_model.gradient_checkpointing = True
self.paligemma_with_expert.paligemma.model.vision_tower.gradient_checkpointing = True
vision_tower = self.paligemma_with_expert.paligemma.model.vision_tower
if self.use_hf_vision_checkpointing_api:
vision_tower.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
else:
vision_tower.gradient_checkpointing = True
self.paligemma_with_expert.gemma_expert.model.gradient_checkpointing = True
logging.info("Enabled gradient checkpointing for PI05Pytorch model")
@@ -452,7 +707,11 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch`
"""Disable gradient checkpointing."""
self.gradient_checkpointing_enabled = False
self.paligemma_with_expert.paligemma.model.language_model.gradient_checkpointing = False
self.paligemma_with_expert.paligemma.model.vision_tower.gradient_checkpointing = False
vision_tower = self.paligemma_with_expert.paligemma.model.vision_tower
if self.use_hf_vision_checkpointing_api:
vision_tower.gradient_checkpointing_disable()
else:
vision_tower.gradient_checkpointing = False
self.paligemma_with_expert.gemma_expert.model.gradient_checkpointing = False
logging.info("Disabled gradient checkpointing for PI05Pytorch model")
@@ -467,6 +726,14 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch`
)
return func(*args, **kwargs)
def _prepare_attention_masks_4d(self, att_2d_masks, dtype=None):
"""Helper method to prepare 4D attention masks for transformer."""
att_2d_masks_4d = att_2d_masks[:, None, :, :]
result = torch.where(att_2d_masks_4d, 0.0, OPENPI_ATTENTION_MASK_VALUE)
if dtype is not None:
result = result.to(dtype=dtype)
return result
def sample_noise(self, shape, device):
return sample_noise(shape, device)
@@ -488,13 +755,16 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch`
pad_masks = []
att_masks = []
# Process images
for img, img_mask in zip(images, img_masks, strict=True):
if self.checkpoint_vision_embeddings:
def image_embed_func(img):
return self.paligemma_with_expert.embed_image(img)
def embed_image(img):
return self._apply_checkpoint(self.paligemma_with_expert.embed_image, img)
img_emb = self._apply_checkpoint(image_embed_func, img)
img_embs = [embed_image(img) for img in images]
else:
img_embs = [self.paligemma_with_expert.embed_image(img) for img in images]
for img_emb, img_mask in zip(img_embs, img_masks, strict=True):
bsize, num_img_embs = img_emb.shape[:2]
embs.append(img_emb)
@@ -556,8 +826,15 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch`
# Set attention masks so that image, language and state inputs do not attend to action tokens
att_masks += [1] + ([0] * (self.config.chunk_size - 1))
att_masks = torch.tensor(att_masks, dtype=action_emb.dtype, device=action_emb.device)
att_masks = att_masks[None, :].expand(bsize, len(att_masks))
if self.use_on_device_suffix_mask:
n = len(att_masks)
att_masks = torch.zeros(n, dtype=action_emb.dtype, device=action_emb.device)
att_masks[0] = 1
att_masks = att_masks[None, :].expand(bsize, n)
else:
att_masks = torch.tensor(att_masks, dtype=action_emb.dtype, device=action_emb.device)
att_masks = att_masks[None, :].expand(bsize, len(att_masks))
return action_emb, pad_masks, att_masks, adarms_cond
@@ -583,7 +860,7 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch`
att_2d_masks = make_att_2d_masks(pad_masks, att_masks)
position_ids = torch.cumsum(pad_masks, dim=1) - 1
att_2d_masks_4d = prepare_attention_masks_4d(att_2d_masks)
att_2d_masks_4d = self._prepare_attention_masks_4d(att_2d_masks)
def forward_func(prefix_embs, suffix_embs, att_2d_masks_4d, position_ids, adarms_cond):
(_, suffix_out), _ = self.paligemma_with_expert.forward(
@@ -641,7 +918,8 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch`
prefix_att_2d_masks = make_att_2d_masks(prefix_pad_masks, prefix_att_masks)
prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1
prefix_att_2d_masks_4d = prepare_attention_masks_4d(prefix_att_2d_masks)
mask_dtype = prefix_embs.dtype if self.use_typed_attention_masks else None
prefix_att_2d_masks_4d = self._prepare_attention_masks_4d(prefix_att_2d_masks, dtype=mask_dtype)
self.paligemma_with_expert.paligemma.model.language_model.config._attn_implementation = "eager" # noqa: SLF001
_, past_key_values = self.paligemma_with_expert.forward(
@@ -652,21 +930,52 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch`
use_cache=True,
)
return euler_integrate(
lambda input_x_t, current_timestep: self.denoise_step(
prefix_pad_masks=prefix_pad_masks,
past_key_values=past_key_values,
x_t=input_x_t,
timestep=current_timestep,
),
noise,
num_steps,
rtc_processor=self.rtc_processor,
rtc_enabled=self._rtc_enabled(),
inference_delay=kwargs.get("inference_delay"),
prev_chunk_left_over=kwargs.get("prev_chunk_left_over"),
execution_horizon=kwargs.get("execution_horizon"),
)
dt = -1.0 / num_steps
times = None
if self.precompute_denoise_times:
times = torch.tensor(
[1.0 + step * dt for step in range(num_steps)], dtype=torch.float32, device=device
)
x_t = noise
for step in range(num_steps):
time = 1.0 + step * dt
if times is None:
time_tensor = torch.tensor(time, dtype=torch.float32, device=device).expand(bsize)
else:
time_tensor = times[step].expand(bsize)
def denoise_step_partial_call(input_x_t, current_timestep=time_tensor):
return self.denoise_step(
prefix_pad_masks=prefix_pad_masks,
past_key_values=past_key_values,
x_t=input_x_t,
timestep=current_timestep,
)
if self._rtc_enabled():
inference_delay = kwargs.get("inference_delay")
prev_chunk_left_over = kwargs.get("prev_chunk_left_over")
execution_horizon = kwargs.get("execution_horizon")
v_t = self.rtc_processor.denoise_step(
x_t=x_t,
prev_chunk_left_over=prev_chunk_left_over,
inference_delay=inference_delay,
time=time,
original_denoise_step_partial=denoise_step_partial_call,
execution_horizon=execution_horizon,
)
else:
v_t = denoise_step_partial_call(x_t)
x_t = x_t + dt * v_t
if self.rtc_processor is not None and self.rtc_processor.is_debug_enabled():
self.rtc_processor.track(time=time, x_t=x_t, v_t=v_t)
return x_t
def denoise_step(
self,
@@ -689,7 +998,7 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch`
prefix_offsets = torch.sum(prefix_pad_masks, dim=-1)[:, None]
position_ids = prefix_offsets + torch.cumsum(suffix_pad_masks, dim=1) - 1
full_att_2d_masks_4d = prepare_attention_masks_4d(full_att_2d_masks)
full_att_2d_masks_4d = self._prepare_attention_masks_4d(full_att_2d_masks)
self.paligemma_with_expert.gemma_expert.model.config._attn_implementation = "eager" # noqa: SLF001
past_key_values = clone_past_key_values(past_key_values)
@@ -713,6 +1022,9 @@ class PI05Policy(PreTrainedPolicy):
config_class = PI05Config
name = "pi05"
model_class = PI05Pytorch
eval_after_pretrained_load = False
show_openpi_disclaimer = True
def __init__(
self,
@@ -730,7 +1042,7 @@ class PI05Policy(PreTrainedPolicy):
# Initialize the core PI05 model
self.init_rtc_processor()
self.model = PI05Pytorch(config, rtc_processor=self.rtc_processor)
self.model = self.model_class(config, rtc_processor=self.rtc_processor)
# Enable gradient checkpointing if requested
if config.gradient_checkpointing:
@@ -756,16 +1068,16 @@ class PI05Policy(PreTrainedPolicy):
strict: bool = True,
**kwargs,
) -> T:
"""Override the from_pretrained method to handle key remapping and display important disclaimer."""
print(
"The PI05 model is a direct port of the OpenPI implementation. \n"
"This implementation follows the original OpenPI structure for compatibility. \n"
"Original implementation: https://github.com/Physical-Intelligence/openpi"
)
"""Load PI05-compatible single-file or sharded safetensors checkpoints."""
if cls.show_openpi_disclaimer:
print(
"The PI05 model is a direct port of the OpenPI implementation. \n"
"This implementation follows the original OpenPI structure for compatibility. \n"
"Original implementation: https://github.com/Physical-Intelligence/openpi"
)
if pretrained_name_or_path is None:
raise ValueError("pretrained_name_or_path is required")
# Use provided config if available, otherwise create default config
if config is None:
config = PreTrainedConfig.from_pretrained(
pretrained_name_or_path=pretrained_name_or_path,
@@ -779,85 +1091,35 @@ class PI05Policy(PreTrainedPolicy):
**kwargs,
)
# Initialize model without loading weights
# Check if dataset_stats were provided in kwargs
model = cls(config, **kwargs)
# Load state dict (expects keys with "model." prefix)
try:
print(f"Loading model from: {pretrained_name_or_path}")
try:
from transformers.utils import cached_file
resolved_file = cached_file(
pretrained_name_or_path,
"model.safetensors",
cache_dir=kwargs.get("cache_dir"),
force_download=kwargs.get("force_download", False),
resume_download=kwargs.get("resume_download"),
proxies=kwargs.get("proxies"),
token=kwargs.get("token"),
revision=kwargs.get("revision"),
local_files_only=kwargs.get("local_files_only", False),
)
from safetensors.torch import load_file
original_state_dict = load_file(resolved_file)
print("✓ Loaded state dict from model.safetensors")
except Exception as e:
print(f"Could not load state dict from remote files: {e}")
print("Returning model without loading pretrained weights")
return model
# First, fix any key differences (see openpi model.py, _fix_pytorch_state_dict_keys)
fixed_state_dict = model._fix_pytorch_state_dict_keys(original_state_dict, model.config)
# Then add "model." prefix for all keys that don't already have it
remapped_state_dict = {}
remap_count = 0
for key, value in fixed_state_dict.items():
if not key.startswith("model."):
new_key = f"model.{key}"
remapped_state_dict[new_key] = value
remap_count += 1
else:
remapped_state_dict[key] = value
if remap_count > 0:
print(f"Remapped {remap_count} state dict keys")
# Load the remapped state dict into the model
missing_keys, unexpected_keys = model.load_state_dict(remapped_state_dict, strict=strict)
if missing_keys:
print(f"Missing keys when loading state dict: {len(missing_keys)} keys")
if len(missing_keys) <= 5:
for key in missing_keys:
print(f" - {key}")
else:
for key in missing_keys[:5]:
print(f" - {key}")
print(f" ... and {len(missing_keys) - 5} more")
if unexpected_keys:
print(f"Unexpected keys when loading state dict: {len(unexpected_keys)} keys")
if len(unexpected_keys) <= 5:
for key in unexpected_keys:
print(f" - {key}")
else:
for key in unexpected_keys[:5]:
print(f" - {key}")
print(f" ... and {len(unexpected_keys) - 5} more")
if not missing_keys and not unexpected_keys:
print("All keys loaded successfully!")
except Exception as e:
print(f"Warning: Could not load state dict: {e}")
files = _resolve_weight_files(
pretrained_name_or_path,
force_download=force_download,
resume_download=resume_download,
proxies=proxies,
token=token,
cache_dir=cache_dir,
local_files_only=local_files_only,
revision=revision,
)
fixed_state_dict = model._fix_pytorch_state_dict_keys(_load_weight_files(files), model.config)
remapped_state_dict = {
key if key.startswith("model.") else f"model.{key}": value
for key, value in fixed_state_dict.items()
}
remapped_state_dict = model._prepare_pretrained_state_dict(remapped_state_dict)
missing_keys, unexpected_keys = model.load_state_dict(remapped_state_dict, strict=strict)
if missing_keys:
logging.warning("Missing %s checkpoint keys: %s", cls.name, missing_keys)
if unexpected_keys:
logging.warning("Unexpected %s checkpoint keys: %s", cls.name, unexpected_keys)
if model.eval_after_pretrained_load:
model.eval()
return model
def _prepare_pretrained_state_dict(self, state_dict: dict[str, Tensor]) -> dict[str, Tensor]:
return state_dict
def _fix_pytorch_state_dict_keys(
self, state_dict, model_config
): # see openpi `BaseModelConfig, _fix_pytorch_state_dict_keys`
@@ -1028,12 +1290,16 @@ class PI05Policy(PreTrainedPolicy):
# Action queue logic for n_action_steps > 1
if len(self._action_queue) == 0:
actions = self.predict_action_chunk(batch)[:, : self.config.n_action_steps]
action_batch = self._prepare_action_batch(batch)
actions = self.predict_action_chunk(action_batch)[:, : self.config.n_action_steps]
# Transpose to get shape (n_action_steps, batch_size, action_dim)
self._action_queue.extend(actions.transpose(0, 1))
return self._action_queue.popleft()
def _prepare_action_batch(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
return batch
@torch.no_grad()
def predict_action_chunk(self, batch: dict[str, Tensor], **kwargs: Unpack[ActionSelectKwargs]) -> Tensor:
"""Predict a chunk of actions given environment observations."""
src/lerobot/policies/pi052/__init__.py
+19
View File
@@ -0,0 +1,19 @@
# 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.
"""PI052 configuration; model and processors are imported lazily by their factories."""
from .configuration_pi052 import PI052Config
__all__ = ["PI052Config"]
src/lerobot/policies/pi052/configuration_pi052.py
+195
View File
@@ -0,0 +1,195 @@
# 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.
"""PI0.5 with hierarchical text generation and flow-matched actions."""
from dataclasses import dataclass
from lerobot.configs import PreTrainedConfig
from lerobot.optim.optimizers import AdamWConfig
from ..pi05.configuration_pi05 import PI05Config
@PreTrainedConfig.register_subclass("pi052")
@dataclass
class PI052Config(PI05Config):
"""PI0.5 configuration for recipe-driven text and action supervision."""
# Recipe / language stack ---------------------------------------------
recipe_path: str | None = "recipes/subtask_mem.yaml"
"""Recipe path relative to ``src/lerobot/configs/``, or ``None`` for the plain PI0.5 prompt."""
apply_chat_template: bool = False
"""Whether to apply a tokenizer chat template.
PaliGemma defaults to plain recipe-rendered prefixes because it is not chat-pretrained.
"""
# Balance frequent recipe text supervision against the paper's α=10 flow weight.
text_loss_weight: float = 1.0
"""LM-head cross-entropy weight; ``0`` disables text training."""
flow_loss_weight: float = 10.0
"""Weight on action-expert flow matching relative to text supervision."""
# Backbone training ---------------------------------------------------
unfreeze_lm_head: bool = True
"""Keep PaliGemma's language head trainable for hierarchical inference."""
# Optional context dropout improves tolerance to missing or stale language state.
plan_dropout_prob: float = 0.0
memory_dropout_prob: float = 0.0
subtask_dropout_prob: float = 0.0
# FAST adds discrete-action CE to the text and flow objectives from paper §III.B-C.
enable_fast_action_loss: bool = True
"""Add FAST-tokenized action cross-entropy to text CE and flow matching."""
action_tokenizer_name: str = "physical-intelligence/fast"
"""HF identifier for the FAST action tokenizer."""
max_action_tokens: int = 256
"""Maximum number of FAST tokens per action chunk."""
fast_skip_tokens: int = 1152
"""Number of top-of-vocab tokens the FAST id mapping skips.
1152 skips PaliGemma's 128 ``<seg>`` and 1024 ``<loc>`` special tokens so
FAST codes land in plain-text ids below 256000 and never collide with the
``<loc>`` targets used for VQA. openpi's pi0-FAST convention is 128 (FAST
occupies the ``<loc>`` range); use 128 only to stay weight-compatible with
checkpoints trained that way."""
fast_action_loss_weight: float = 1.0
"""Weight on FAST action-token CE relative to continuous-flow supervision."""
subtask_replan_steps: int = 0
"""Environment steps between subtask generations during evaluation.
Non-positive values regenerate each action chunk while still refreshing the action prompt every chunk.
"""
joint_subtask_conditioning: bool = False
"""Condition low-level action inference on the task plus the generated subtask.
Matches paper-style joint-sequence recipes (``recipes/subtask_joint.yaml``)
where one sample supervises the subtask text and conditions the action
losses on it: the inference prefix becomes
``User: {task}, State: ...;\\nAssistant: {subtask}<eos>`` with the subtask
span attended causally, exactly as trained. Leave ``False`` for the blend
recipes, whose low-level samples use ``User: {subtask}, State: ...;``."""
auto_fit_fast_tokenizer: bool = False
"""Fit and cache a dataset-specific FAST tokenizer before training.
Disabled by default to avoid the extra dataset pass and use the universal tokenizer.
"""
fast_tokenizer_cache_dir: str = "~/.cache/lerobot/fast_tokenizers"
"""Where fitted FAST tokenizers are stored. ``~`` expands."""
fast_tokenizer_fit_samples: int = 1024
"""Number of action chunks sampled when fitting FAST."""
fast_tokenizer_validation_samples: int = 256
"""Held-out action chunks used to validate tokenizer reconstruction."""
fast_tokenizer_max_reconstruction_rmse: float = 0.10
"""Maximum normalized RMSE allowed across held-out action chunks."""
fast_tokenizer_max_dim_rmse: float = 0.20
"""Maximum normalized RMSE allowed for any nonconstant action dimension."""
# Knowledge insulation detaches VLM K/V from action-loss gradients (paper §III.B).
knowledge_insulation: bool = True
"""Block action-loss gradients through VLM keys and values."""
# Optional training backends. Defaults preserve the eager/SDPA path.
use_flashrt_adarms: bool = False
"""Use FlashRT adaptive RMSNorm kernels when available."""
use_compiled_text_ce: bool = False
"""Compile the materialized-logits text and FAST CE path."""
use_compiled_vision: bool = False
"""Compile the SigLIP tower for no-grad flow and inference passes."""
use_flex_attention: bool = False
"""Use FlexAttention for amortized KI, with SDPA fallback where unsupported."""
use_manual_attention: bool = False
"""Use materialized-logits attention for explicitly profiled KI shapes."""
manual_attention_scope: str = "all"
"""Apply manual attention to all KI queries or only action queries."""
# Scale language-head updates relative to the base optimizer schedule.
lm_head_lr_scale: float = 1.0
# Scale backbone and action-expert optimizer groups independently.
backbone_lr_scale: float = 1.0
action_expert_lr_scale: float = 1.0
# Reuse each VLM prefix across independent denoising draws; 1 restores single-draw flow.
flow_num_repeats: int = 5
# PaLM-style z-loss stabilizes large-vocabulary CE; 0 disables it.
text_ce_z_loss_weight: float = 1e-4
use_flashrt_fp8_mlp: bool = False
"""Enable calibrated FlashRT FP8 kernels for Gemma and SigLIP MLPs.
Apply after loading with ``PI052Policy.apply_flashrt_fp8_mlp``; unavailable kernels keep BF16.
"""
# Keep serialized PI052 AdamW options local because PI05Config lacks them.
optimizer_foreach: bool | None = False
optimizer_fused: bool | None = True
def get_optimizer_preset(self) -> AdamWConfig:
return AdamWConfig(
lr=self.optimizer_lr,
betas=self.optimizer_betas,
eps=self.optimizer_eps,
weight_decay=self.optimizer_weight_decay,
grad_clip_norm=self.optimizer_grad_clip_norm,
foreach=self.optimizer_foreach,
fused=self.optimizer_fused,
)
def __post_init__(self) -> None:
super().__post_init__()
if self.enable_fast_action_loss and not self.recipe_path:
raise ValueError("PI052 FAST action loss requires recipe_path to build action supervision.")
if self.text_loss_weight > 0 and self.unfreeze_lm_head:
self.train_expert_only = False
if self.flow_num_repeats < 1:
raise ValueError(f"flow_num_repeats must be >= 1, got {self.flow_num_repeats}")
if self.fast_tokenizer_validation_samples < 1:
raise ValueError("fast_tokenizer_validation_samples must be >= 1")
if self.fast_tokenizer_max_reconstruction_rmse <= 0 or self.fast_tokenizer_max_dim_rmse <= 0:
raise ValueError("FAST tokenizer reconstruction thresholds must be positive")
if self.manual_attention_scope not in {"all", "action"}:
raise ValueError(
f"manual_attention_scope must be 'all' or 'action', got {self.manual_attention_scope!r}"
)
if self.use_flex_attention and self.use_manual_attention:
raise ValueError("use_flex_attention and use_manual_attention are mutually exclusive")
if self.use_flex_attention and self.flow_num_repeats == 1:
raise ValueError("use_flex_attention requires flow_num_repeats > 1")
if not self.knowledge_insulation and (
self.use_flex_attention or self.use_manual_attention or self.use_flashrt_adarms
):
raise ValueError("KI attention and AdaRMS optimizations require knowledge_insulation=True")
src/lerobot/policies/pi052/fit_fast_tokenizer.py
+522
View File
@@ -0,0 +1,522 @@
# 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.
"""Fit and cache a FAST tokenizer for a dataset's action distribution.
Training invokes this automatically when FAST loss and automatic fitting are enabled.
"""
from __future__ import annotations
import hashlib
import json
import logging
import os
import shutil
import time
from pathlib import Path
from typing import Any
import numpy as np
logger = logging.getLogger(__name__)
# ``ProcessorMixin.save_pretrained`` writes this shared cache sentinel.
_CACHE_SENTINEL = "processor_config.json"
def _is_global_leader() -> bool:
return int(os.environ.get("RANK", "0")) == 0
def _jsonable(value: Any) -> Any:
if hasattr(value, "detach"):
value = value.detach().cpu().numpy()
if isinstance(value, np.ndarray):
return value.tolist()
if isinstance(value, dict):
return {key: _jsonable(item) for key, item in sorted(value.items())}
if isinstance(value, (list, tuple)):
return [_jsonable(item) for item in value]
return value
def _dataset_signature(
dataset_repo_id: str,
base_tokenizer_name: str,
n_samples: int,
chunk_size: int,
normalization_mode: str,
dataset_revision: str | None = None,
episodes: list[int] | None = None,
exclude_episodes: list[int] | None = None,
action_stats: dict | None = None,
use_relative_actions: bool = False,
relative_action_mask: list[bool] | None = None,
validation_samples: int = 256,
max_reconstruction_rmse: float = 0.10,
max_dim_rmse: float = 0.20,
) -> str:
"""Hash every input that changes the fitted action distribution."""
payload = {
"dataset_repo_id": dataset_repo_id,
"dataset_revision": dataset_revision,
"base_tokenizer_name": base_tokenizer_name,
"n_samples": n_samples,
"chunk_size": chunk_size,
"normalization_mode": normalization_mode,
"episodes": episodes,
"exclude_episodes": exclude_episodes,
"action_stats": action_stats,
"use_relative_actions": use_relative_actions,
"relative_action_mask": relative_action_mask,
"validation_samples": validation_samples,
"max_reconstruction_rmse": max_reconstruction_rmse,
"max_dim_rmse": max_dim_rmse,
}
encoded = json.dumps(_jsonable(payload), sort_keys=True, separators=(",", ":")).encode()
return hashlib.sha256(encoded).hexdigest()[:16]
def _select_episode_indices(
available_episodes: list[int],
episodes: list[int] | None,
exclude_episodes: list[int] | None,
) -> list[int]:
allowed = set(episodes) if episodes is not None else set(available_episodes)
excluded = set(exclude_episodes or [])
return [episode for episode in available_episodes if episode in allowed and episode not in excluded]
def _apply_relative_actions(
actions: np.ndarray,
states: np.ndarray,
relative_action_mask: list[bool] | None,
) -> np.ndarray:
"""Match RelativeActionsProcessorStep before tokenizer fitting."""
action_dim = actions.shape[-1]
mask = list(relative_action_mask) if relative_action_mask is not None else [True] * action_dim
if len(mask) < action_dim:
mask.extend([True] * (action_dim - len(mask)))
mask_array = np.asarray(mask[:action_dim], dtype=np.float32)
relative = actions.copy()
relative -= states[:, None, :action_dim] * mask_array
return relative
def _normalize_actions(
actions: np.ndarray,
normalization_mode: str,
action_stats: dict | None = None,
) -> np.ndarray:
"""Match the action normalization applied by the training preprocessor."""
mode = getattr(normalization_mode, "value", normalization_mode).upper()
flat = actions.reshape(-1, actions.shape[-1])
stats = action_stats or {}
def stat(name: str, fallback) -> np.ndarray:
value = stats.get(name)
if value is None:
value = fallback()
if hasattr(value, "detach"):
value = value.detach().cpu().numpy()
return np.asarray(value, dtype=np.float32)
if mode == "IDENTITY":
return actions
if mode == "MEAN_STD":
mean = stat("mean", lambda: flat.mean(axis=0))
std = stat("std", lambda: flat.std(axis=0))
return ((actions - mean) / np.where(std == 0, 1e-8, std)).astype(np.float32)
if mode in {"QUANTILES", "QUANTILE10"}:
low_name, high_name, low_q, high_q = (
("q01", "q99", 0.01, 0.99) if mode == "QUANTILES" else ("q10", "q90", 0.10, 0.90)
)
low = stat(low_name, lambda: np.quantile(flat, low_q, axis=0))
high = stat(high_name, lambda: np.quantile(flat, high_q, axis=0))
elif mode == "MIN_MAX":
low = stat("min", lambda: flat.min(axis=0))
high = stat("max", lambda: flat.max(axis=0))
else:
raise ValueError(f"Unsupported FAST tokenizer normalization mode: {mode}")
return (2.0 * (actions - low) / np.where(high == low, 1e-8, high - low) - 1.0).astype(np.float32)
def _validate_fast_reconstruction(
tokenizer: Any,
actions: np.ndarray,
max_reconstruction_rmse: float,
max_dim_rmse: float,
) -> tuple[dict[str, Any], np.ndarray]:
"""Decode held-out chunks and reject tokenizers with excessive quantization error."""
decoded = np.asarray(tokenizer.decode(tokenizer(actions)), dtype=np.float32)
if decoded.shape != actions.shape:
raise RuntimeError(
f"FAST tokenizer reconstruction shape mismatch: expected {actions.shape}, got {decoded.shape}."
)
if not np.isfinite(decoded).all():
raise RuntimeError("FAST tokenizer reconstruction contains non-finite values.")
squared_error = np.square(decoded - actions)
rmse = float(np.sqrt(squared_error.mean()))
dim_rmse = np.sqrt(squared_error.mean(axis=(0, 1)))
nonconstant_dims = np.ptp(actions, axis=(0, 1)) > 1e-8
max_observed_dim_rmse = float(dim_rmse[nonconstant_dims].max(initial=0.0))
report = {
"num_validation_chunks": int(actions.shape[0]),
"reconstruction_rmse": rmse,
"max_dim_rmse": max_observed_dim_rmse,
"dim_rmse": dim_rmse.tolist(),
"max_reconstruction_rmse": max_reconstruction_rmse,
"max_allowed_dim_rmse": max_dim_rmse,
}
if rmse > max_reconstruction_rmse or max_observed_dim_rmse > max_dim_rmse:
raise RuntimeError(
"FAST tokenizer reconstruction error exceeds the configured limit: "
f"rmse={rmse:.4f} (max {max_reconstruction_rmse:.4f}), "
f"max_dim_rmse={max_observed_dim_rmse:.4f} (max {max_dim_rmse:.4f})."
)
return report, decoded
def _load_fast_fitter(base_tokenizer_name: str) -> Any:
"""Load FAST's fitting implementation without requiring its universal BPE weights."""
from transformers import AutoProcessor # noqa: PLC0415
try:
return AutoProcessor.from_pretrained(base_tokenizer_name, trust_remote_code=True)
except ValueError as error:
if base_tokenizer_name != "physical-intelligence/fast":
raise
logger.warning(
"Could not load the universal FAST tokenizer backend; loading its fitting class directly: %s",
error,
)
from transformers.dynamic_module_utils import get_class_from_dynamic_module # noqa: PLC0415
return get_class_from_dynamic_module(
"processing_action_tokenizer.UniversalActionProcessor",
base_tokenizer_name,
)
def fit_fast_tokenizer(
*,
dataset_repo_id: str,
cache_dir: str | Path,
base_tokenizer_name: str = "physical-intelligence/fast",
n_samples: int = 1024,
chunk_size: int = 50,
seed: int = 42,
dataset_root: str | Path | None = None,
dataset_revision: str | None = None,
episodes: list[int] | None = None,
exclude_episodes: list[int] | None = None,
normalization_mode: str = "QUANTILES",
action_stats: dict | None = None,
use_relative_actions: bool = False,
relative_action_mask: list[bool] | None = None,
validation_samples: int = 256,
max_reconstruction_rmse: float = 0.10,
max_dim_rmse: float = 0.20,
) -> str:
"""Fit a FAST tokenizer on a LeRobot dataset's action distribution.
Args:
dataset_repo_id: HF Hub repo id of the LeRobotDataset to fit on.
cache_dir: Directory under which to save (and look up) fitted
tokenizers. The actual save path is
``{cache_dir}/{signature}``.
base_tokenizer_name: HF identifier for the base FAST tokenizer
to finetune from. ``physical-intelligence/fast`` is the
universal one.
n_samples: Number of action chunks to sample for the fit. The
FAST paper uses a few thousand; ``1024`` is a good default
for medium datasets.
chunk_size: Length of each action chunk (matches
``policy.chunk_size``). The FAST tokenizer is fit on
sequences of this length.
seed: RNG seed for sample selection.
Returns:
The local path to the fitted tokenizer. Passed directly to
``--policy.action_tokenizer_name`` for the training run.
Raises:
ImportError: If the ``transformers`` library doesn't expose
``AutoProcessor`` or the FAST tokenizer doesn't have a
``.fit()`` method (then you're on an older FAST snapshot —
update to the current published model).
FileNotFoundError: If the dataset can't be loaded.
"""
cache_dir = Path(cache_dir)
normalization_mode = getattr(normalization_mode, "value", normalization_mode).upper()
sig = _dataset_signature(
dataset_repo_id,
base_tokenizer_name,
n_samples,
chunk_size,
normalization_mode,
dataset_revision,
episodes,
exclude_episodes,
action_stats,
use_relative_actions,
relative_action_mask,
validation_samples,
max_reconstruction_rmse,
max_dim_rmse,
)
out_dir = cache_dir / sig
if out_dir.exists() and (out_dir / _CACHE_SENTINEL).exists():
logger.info(
"FAST tokenizer cache hit: %s — re-using fitted tokenizer for dataset=%s base=%s n_samples=%d",
out_dir,
dataset_repo_id,
base_tokenizer_name,
n_samples,
)
return str(out_dir)
# One global rank populates the shared cache; every other rank waits for the atomic publish.
is_leader = _is_global_leader()
if not is_leader:
timeout_s = 1800.0 # 30 min — covers ~1024-sample fits on cold caches
start = time.monotonic()
while not (out_dir / _CACHE_SENTINEL).exists():
if time.monotonic() - start > timeout_s:
raise RuntimeError(
f"FAST tokenizer fit: non-leader rank timed out after "
f"{timeout_s:.0f}s waiting for {out_dir / _CACHE_SENTINEL}. "
"Leader rank likely crashed during the fit."
)
time.sleep(2.0)
logger.info("FAST tokenizer ready (leader populated cache): %s", out_dir)
return str(out_dir)
logger.info(
"FAST tokenizer cache miss — fitting on dataset=%s base=%s n_samples=%d chunk_size=%d → %s",
dataset_repo_id,
base_tokenizer_name,
n_samples,
chunk_size,
out_dir,
)
# Read action columns directly to avoid video decoding and bound memory to sampled episodes.
rng = np.random.default_rng(seed)
actions_buf: list[np.ndarray] = []
# Read v3 parquet shards directly to avoid split lookup failures and repeated metadata parsing.
import pyarrow as _pa # noqa: PLC0415
import pyarrow.parquet as _pq # noqa: PLC0415
if dataset_root is not None:
snap = Path(dataset_root)
else:
from huggingface_hub import snapshot_download # noqa: PLC0415
snap = Path(
snapshot_download(repo_id=dataset_repo_id, repo_type="dataset", revision=dataset_revision)
)
data_files = sorted((snap / "data").glob("chunk-*/file-*.parquet"))
if not data_files:
raise RuntimeError(f"FAST fit: no ``data/chunk-*/file-*.parquet`` shards found under {snap!s}.")
columns = ["episode_index", "action"]
if use_relative_actions:
columns.append("observation.state")
tables = [_pq.read_table(f, columns=columns) for f in data_files]
table = _pa.concat_tables(tables)
eps = table["episode_index"].to_numpy()
acts_col = table["action"]
# Normalize Arrow action representations into an (N, D) array.
try:
acts = np.stack(acts_col.to_numpy(zero_copy_only=False)).astype(np.float32)
except Exception: # noqa: BLE001
# Fallback path for nested-list types: flatten via to_pylist().
acts = np.asarray(acts_col.to_pylist(), dtype=np.float32)
if acts.ndim != 2:
raise RuntimeError(f"FAST fit: expected ``action`` rows to be 1-D vectors; got shape {acts.shape}.")
states = None
if use_relative_actions:
try:
states = np.stack(table["observation.state"].to_numpy(zero_copy_only=False)).astype(np.float32)
except Exception: # noqa: BLE001
states = np.asarray(table["observation.state"].to_pylist(), dtype=np.float32)
if states.ndim != 2:
raise RuntimeError(
f"FAST fit: expected ``observation.state`` rows to be 1-D vectors; got {states.shape}."
)
# Sort once because episode order is only guaranteed within each shard.
order = np.argsort(eps, kind="stable")
eps_sorted = eps[order]
boundaries = np.searchsorted(eps_sorted, np.arange(int(eps_sorted.max()) + 2))
ep_to_slice: dict[int, tuple[int, int]] = {
int(ep): (int(boundaries[ep]), int(boundaries[ep + 1]))
for ep in range(len(boundaries) - 1)
if boundaries[ep] < boundaries[ep + 1]
}
num_episodes = len(ep_to_slice)
# ``acts`` is in original (un-sorted-by-episode) row order; reorder
# so per-episode slices are contiguous.
acts = acts[order]
if states is not None:
states = states[order]
ep_indices = _select_episode_indices(list(ep_to_slice), episodes, exclude_episodes)
if not ep_indices:
raise RuntimeError("FAST fit: episode selection is empty after applying exclusions.")
total_samples = n_samples + validation_samples
samples_per_episode = max(1, (total_samples + len(ep_indices) - 1) // len(ep_indices))
collected = 0
eps_visited = 0
short_episodes = 0
states_buf: list[np.ndarray] = []
for ep_idx in rng.permutation(ep_indices):
if collected >= total_samples:
break
start, stop = ep_to_slice[int(ep_idx)]
ep_actions = acts[start:stop]
if ep_actions.shape[0] < chunk_size:
short_episodes += 1
continue
starts = rng.integers(0, ep_actions.shape[0] - chunk_size + 1, size=samples_per_episode)
for s in starts:
actions_buf.append(ep_actions[int(s) : int(s) + chunk_size])
if states is not None:
states_buf.append(states[start + int(s)])
collected += 1
if collected >= total_samples:
break
eps_visited += 1
if not actions_buf:
raise RuntimeError(
f"FAST fit collected zero action chunks from {dataset_repo_id!r}: "
f"all {num_episodes} episodes were shorter than chunk_size="
f"{chunk_size} ({short_episodes} too short) or had an unreadable "
"``action`` column. Lower ``chunk_size`` to match your episode "
"lengths."
)
actions = np.stack(actions_buf, axis=0).astype(np.float32) # (N, H, D)
if states is not None:
actions = _apply_relative_actions(actions, np.stack(states_buf), relative_action_mask)
logger.info(
"FAST fit: collected %d chunks of shape %s from %d episodes",
actions.shape[0],
actions.shape[1:],
eps_visited,
)
actions = _normalize_actions(actions, normalization_mode, action_stats)
base = _load_fast_fitter(base_tokenizer_name)
if not hasattr(base, "fit"):
raise ImportError(
f"Base FAST tokenizer {base_tokenizer_name!r} has no ``.fit()`` "
"method — your transformers / model snapshot is too old. Update "
"to the current ``physical-intelligence/fast`` revision."
)
if actions.shape[0] < total_samples:
raise RuntimeError(
f"FAST fit collected {actions.shape[0]} chunks, but {total_samples} are required "
f"for {n_samples} fit and {validation_samples} validation chunks."
)
fit_actions = actions[:n_samples]
validation_actions = actions[n_samples:total_samples]
fitted = base.fit(fit_actions)
validation_report, decoded_actions = _validate_fast_reconstruction(
fitted,
validation_actions,
max_reconstruction_rmse,
max_dim_rmse,
)
cache_dir.mkdir(parents=True, exist_ok=True)
staging_dir = cache_dir / f".{sig}.tmp-{os.getpid()}"
shutil.rmtree(staging_dir, ignore_errors=True)
fitted.save_pretrained(str(staging_dir))
(staging_dir / "reconstruction_validation.json").write_text(
json.dumps(validation_report, indent=2) + "\n"
)
np.savez_compressed(
staging_dir / "reconstruction_examples.npz",
original=validation_actions[:8],
decoded=decoded_actions[:8],
)
if out_dir.exists():
shutil.rmtree(out_dir)
staging_dir.replace(out_dir)
logger.info("FAST fit: saved fitted tokenizer to %s", out_dir)
return str(out_dir)
def resolve_fast_tokenizer(
config: Any,
dataset_repo_id: str | None,
dataset_root: str | Path | None = None,
dataset_stats: dict | None = None,
dataset_revision: str | None = None,
episodes: list[int] | None = None,
exclude_episodes: list[int] | None = None,
) -> str:
"""Return the configured tokenizer, fitting a cached dataset-specific one when requested."""
if not getattr(config, "auto_fit_fast_tokenizer", False) or dataset_repo_id is None:
return config.action_tokenizer_name
relative_action_mask = None
if getattr(config, "use_relative_actions", False):
action_names = getattr(config, "action_feature_names", None)
exclude_tokens = [
str(name).lower() for name in getattr(config, "relative_exclude_joints", []) if name
]
if action_names is not None and exclude_tokens:
relative_action_mask = [
not any(token == str(name).lower() or token in str(name).lower() for token in exclude_tokens)
for name in action_names
]
fit_kwargs = {
"dataset_repo_id": dataset_repo_id,
"cache_dir": Path(config.fast_tokenizer_cache_dir).expanduser(),
"base_tokenizer_name": config.action_tokenizer_name,
"n_samples": config.fast_tokenizer_fit_samples,
"chunk_size": config.chunk_size,
"dataset_root": dataset_root,
"dataset_revision": dataset_revision,
"episodes": episodes,
"exclude_episodes": exclude_episodes,
"normalization_mode": config.normalization_mapping.get("ACTION", "QUANTILES"),
"action_stats": (dataset_stats or {}).get("action"),
"use_relative_actions": getattr(config, "use_relative_actions", False),
"relative_action_mask": relative_action_mask,
}
validation_fields = {
"validation_samples": "fast_tokenizer_validation_samples",
"max_reconstruction_rmse": "fast_tokenizer_max_reconstruction_rmse",
"max_dim_rmse": "fast_tokenizer_max_dim_rmse",
}
fit_kwargs.update(
{
argument: getattr(config, attribute)
for argument, attribute in validation_fields.items()
if hasattr(config, attribute)
}
)
return fit_fast_tokenizer(**fit_kwargs)
src/lerobot/policies/pi052/flashrt_fp8.py
+263
View File
@@ -0,0 +1,263 @@
# 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.
"""Optional FlashRT FP8 MLP kernels with one-pass calibration and BF16 fallback."""
from __future__ import annotations
import logging
import torch
import torch.nn as nn
import torch.nn.functional as F # noqa: N812
logger = logging.getLogger(__name__)
_FP8_MAX = 448.0
def _roundtrip_fp8(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
"""Quantize->dequantize an activation through FP8 E4M3 at ``scale`` (f32)."""
q = torch.clamp(x.float() / scale.float(), -_FP8_MAX, _FP8_MAX).to(torch.float8_e4m3fn)
return q.float() * scale.float()
_SWIGLU_REPO = "flashrt/flashrt-fp8-swiglu-ffn"
_GELU_REPO = "flashrt/flashrt-fp8-ffn"
_GEMM_REPO = "flashrt/flashrt-gemm-epilogues"
def _get_kernel(repo: str):
"""Load a cached FlashRT Hub package."""
from kernels import get_kernel
return get_kernel(repo, version=1)
def _quantize_fp8(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
scale = max(weight.detach().float().abs().max().item(), 1e-12) / _FP8_MAX
fp8 = torch.clamp(weight.float() / scale, -_FP8_MAX, _FP8_MAX).to(torch.float8_e4m3fn)
return fp8.contiguous(), torch.tensor([scale], dtype=torch.float32)
def _static_scale(amax: float, safety: float) -> torch.Tensor:
return torch.tensor([max(amax, 1e-12) / _FP8_MAX * safety], dtype=torch.float32)
class _FlashRTGeGLU(nn.Module):
"""FP8 Gemma GeGLU MLP."""
def __init__(self, mlp, in_amax, hid_amax, ffn_ops, quant_ops, safety, fuse_weight=None):
super().__init__()
self.ffn_ops = ffn_ops
self.quant_ops = quant_ops
self.in_features = mlp.gate_proj.weight.shape[1]
device = mlp.gate_proj.weight.device
gate_up = torch.cat([mlp.gate_proj.weight, mlp.up_proj.weight], dim=0).float()
# Fold fixed RMSNorm weights into GEMM; adaptive norms use identity scaling.
if fuse_weight is not None:
f = 1.0 + fuse_weight.detach().float()
gate_up = gate_up * f[None, :]
channel_scale = (1.0 / f).to(torch.bfloat16)
else:
channel_scale = torch.ones(self.in_features, dtype=torch.bfloat16)
gate_up_fp8, gate_up_scale = _quantize_fp8(gate_up)
down_fp8, down_scale = _quantize_fp8(mlp.down_proj.weight)
self.register_buffer("gate_up_fp8", gate_up_fp8.to(device))
self.register_buffer("down_fp8", down_fp8.to(device))
self.register_buffer("gate_up_scale", gate_up_scale.to(device))
self.register_buffer("down_scale", down_scale.to(device))
self.register_buffer("input_scale", _static_scale(in_amax, safety).to(device))
self.register_buffer("hidden_scale", _static_scale(hid_amax, safety).to(device))
self.register_buffer("channel_scale", channel_scale.to(device))
self.safety = safety
self.calibrating = False
self._ia = 0.0
self._ha = 0.0
def _calibrate_step(self, x):
# Track input and hidden maxima on live FP8-propagated activations.
flat = x.reshape(-1, self.in_features).to(torch.bfloat16)
xq = flat.float() * self.channel_scale.float()
self._ia = max(self._ia, xq.abs().max().item())
self.input_scale.copy_(_static_scale(self._ia, self.safety).to(self.input_scale.device))
xdq = _roundtrip_fp8(xq, self.input_scale)
wdq = self.gate_up_fp8.float() * self.gate_up_scale.float()
gate, up = (xdq @ wdq.t()).chunk(2, dim=-1)
hidden = F.gelu(gate, approximate="tanh") * up
self._ha = max(self._ha, hidden.abs().max().item())
self.hidden_scale.copy_(_static_scale(self._ha, self.safety).to(self.hidden_scale.device))
def forward(self, x):
if self.calibrating:
self._calibrate_step(x)
shape = x.shape
flat = x.reshape(-1, self.in_features).to(torch.bfloat16)
x_fp8 = self.quant_ops.channel_scale_quantize_fp8_static_bf16(
flat, self.channel_scale, self.input_scale
)
out = self.ffn_ops.fp8_geglu_mlp_bf16(
x_fp8,
self.gate_up_fp8,
self.down_fp8,
self.input_scale,
self.gate_up_scale,
self.hidden_scale,
self.down_scale,
)
return out.reshape(shape)
class _FlashRTGeluMLP(nn.Module):
"""FP8 SigLIP GELU MLP."""
def __init__(self, mlp, in_amax, hid_amax, ffn_ops, quant_ops, safety):
super().__init__()
self.ffn_ops = ffn_ops
self.quant_ops = quant_ops
self.in_features = mlp.fc1.weight.shape[1]
self.out_features = mlp.fc2.weight.shape[0]
device = mlp.fc1.weight.device
up_fp8, up_scale = _quantize_fp8(mlp.fc1.weight)
down_fp8, down_scale = _quantize_fp8(mlp.fc2.weight)
self.register_buffer("up_fp8", up_fp8.to(device))
self.register_buffer("down_fp8", down_fp8.to(device))
self.register_buffer("up_scale", up_scale.to(device))
self.register_buffer("down_scale", down_scale.to(device))
self.register_buffer("up_bias", mlp.fc1.bias.detach().to(torch.bfloat16))
self.register_buffer("down_bias", mlp.fc2.bias.detach().to(torch.bfloat16))
self.register_buffer("input_scale", _static_scale(in_amax, safety).to(device))
self.register_buffer("hidden_scale", _static_scale(hid_amax, safety).to(device))
self.register_buffer(
"channel_scale", torch.ones(self.in_features, device=device, dtype=torch.bfloat16)
)
self.safety = safety
self.calibrating = False
self._ia = 0.0
self._ha = 0.0
def _calibrate_step(self, x):
flat = x.reshape(-1, self.in_features).to(torch.bfloat16)
self._ia = max(self._ia, flat.float().abs().max().item())
self.input_scale.copy_(_static_scale(self._ia, self.safety).to(self.input_scale.device))
xdq = _roundtrip_fp8(flat.float(), self.input_scale)
hid = (xdq @ (self.up_fp8.float() * self.up_scale.float()).t()) + self.up_bias.float()
hid = F.gelu(hid, approximate="tanh")
self._ha = max(self._ha, hid.abs().max().item())
self.hidden_scale.copy_(_static_scale(self._ha, self.safety).to(self.hidden_scale.device))
def forward(self, x):
if self.calibrating:
self._calibrate_step(x)
shape = x.shape
dtype = x.dtype
flat = x.reshape(-1, self.in_features).to(torch.bfloat16)
x_fp8 = self.quant_ops.channel_scale_quantize_fp8_static_bf16(
flat, self.channel_scale, self.input_scale
)
out = self.ffn_ops.fp8_gelu_mlp_bf16(
x_fp8,
self.up_fp8,
self.up_bias,
self.down_fp8,
self.down_bias,
self.input_scale,
self.up_scale,
self.hidden_scale,
self.down_scale,
)
return out.reshape(*shape[:-1], self.out_features).to(dtype)
def _siglip_mlps(model) -> list:
tower = model.paligemma_with_expert.paligemma.model.vision_tower
return [m for _, m in tower.named_modules() if type(m).__name__ == "SiglipMLP"]
def _run_forward(policy, batches) -> None:
"""Run eager action prediction so calibration reaches Python module forwards."""
model = policy.model
saved = {name: vars(model).pop(name) for name in ("sample_actions", "forward") if name in vars(model)}
with torch.inference_mode():
for batch in batches:
policy.predict_action_chunk(
{k: (v.clone() if torch.is_tensor(v) else v) for k, v in batch.items()}
)
torch.cuda.synchronize()
vars(model).update(saved)
def _fixed_norm_weight(norm):
"""Return a fixed RMSNorm fold weight, or ``None`` for adaptive norms."""
return norm.weight if getattr(norm, "dense", None) is None else None
def _fp8_supported(device) -> bool:
"""Return whether the device supports FP8 E4M3 tensor cores (CUDA SM >= 8.9)."""
if device.type != "cuda" or not torch.cuda.is_available():
return False
major, minor = torch.cuda.get_device_capability(device)
return (major, minor) >= (8, 9)
def apply_fp8_mlp(policy, batch, *, safety: float = 1.05) -> bool:
"""Replace Gemma and SigLIP MLPs with FlashRT FP8 kernels calibrated on the supplied batch.
Returns ``False`` without modifying BF16 execution when FP8 or its kernels are unavailable.
"""
device = next(policy.parameters()).device
if not _fp8_supported(device):
logger.warning(
"PI052: device %s has no FP8 (E4M3) support (needs CUDA SM>=8.9); keeping BF16.",
device,
)
return False
batches = batch if isinstance(batch, (list, tuple)) else [batch]
try:
ffn_ops = _get_kernel(_SWIGLU_REPO)
gelu_ops = _get_kernel(_GELU_REPO)
quant_ops = _get_kernel(_GEMM_REPO)
except Exception as exc: # noqa: BLE001
logger.warning("PI052: FlashRT FP8 kernels unavailable (%s); keeping BF16.", exc)
return False
model = policy.model
calibrating = []
gemma_layers = list(model.paligemma_with_expert.gemma_expert.model.layers) + list(
model.paligemma_with_expert.paligemma.model.language_model.layers
)
for layer in gemma_layers:
fw = _fixed_norm_weight(layer.post_attention_layernorm)
layer.mlp = _FlashRTGeGLU(layer.mlp, 1.0, 1.0, ffn_ops, quant_ops, safety, fuse_weight=fw).to(device)
calibrating.append(layer.mlp)
siglip = _siglip_mlps(model)
for mlp_parent in model.paligemma_with_expert.paligemma.model.vision_tower.vision_model.encoder.layers:
mlp_parent.mlp = _FlashRTGeluMLP(mlp_parent.mlp, 1.0, 1.0, gelu_ops, quant_ops, safety).to(device)
calibrating.append(mlp_parent.mlp)
# Calibrate every swapped module in one FP8-propagated forward.
for m in calibrating:
m.calibrating = True
_run_forward(policy, batches)
for m in calibrating:
m.calibrating = False
logger.info(
"PI052: FlashRT FP8 enabled (%d Gemma + %d SigLIP MLPs).",
len(gemma_layers),
len(siglip),
)
return True
src/lerobot/policies/pi052/inference/__init__.py
+19
View File
@@ -0,0 +1,19 @@
# 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.
"""PI052 adapter for the policy-agnostic language runtime."""
from .pi052_adapter import PI052PolicyAdapter
__all__ = ["PI052PolicyAdapter"]
src/lerobot/policies/pi052/inference/pi052_adapter.py
+254
View File
@@ -0,0 +1,254 @@
# 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.
"""PI052 actions and text generation for the generic language runtime."""
from __future__ import annotations
import logging
from typing import Any
from lerobot.runtime import RuntimeState
from lerobot.runtime.adapter import BaseLanguageAdapter
logger = logging.getLogger(__name__)
_LOC_TOKENIZER_CACHE: dict[str, Any] = {}
class PI052PolicyAdapter(BaseLanguageAdapter):
"""Runtime bridge for PI052 policies."""
def select_action(self, observation: dict[str, Any], state: RuntimeState) -> Any:
import torch # noqa: PLC0415
from lerobot.utils.constants import ( # noqa: PLC0415
OBS_LANGUAGE_ATTENTION_MASK,
OBS_LANGUAGE_TOKENS,
OBS_STATE,
)
subtask = state.language_context.get("subtask") or state.task or ""
# Match the training prompt by conditioning on both subtask and discretized state.
state_str = None
obs_state = observation.get(OBS_STATE)
if isinstance(obs_state, torch.Tensor) and obs_state.numel() > 0:
from lerobot.policies.pi052.text_processor_pi052 import discretize_state_str # noqa: PLC0415
state_row = obs_state[0] if obs_state.ndim > 1 else obs_state
state_str = discretize_state_str(state_row)
batch = dict(observation)
if getattr(self.policy.config, "joint_subtask_conditioning", False):
# Joint sequences keep the task turn (with state) and render the
# subtask as a causal assistant turn, exactly as trained.
from transformers import AutoTokenizer # noqa: PLC0415
from lerobot.policies.pi052.text_processor_pi052 import ( # noqa: PLC0415
encode_prompt_with_targets,
register_paligemma_loc_tokens,
)
from lerobot.utils.constants import OBS_LANGUAGE_CAUSAL_MARKS # noqa: PLC0415
task = state.task or ""
task_content = task if state_str is None else f"{task}, State: {state_str};"
tok_name = getattr(self.policy.config, "tokenizer_name", None) or "google/paligemma-3b-pt-224"
tokenizer = _get_loc_tokenizer(tok_name, AutoTokenizer, register_paligemma_loc_tokens)
ids, attn, marks = encode_prompt_with_targets(
tokenizer,
[
{"role": "user", "content": task_content},
{"role": "assistant", "content": subtask},
],
target_indices=[1],
)
device = getattr(self.policy.config, "device", None)
if device is not None:
ids, attn, marks = ids.to(device), attn.to(device), marks.to(device)
batch[OBS_LANGUAGE_TOKENS] = ids
batch[OBS_LANGUAGE_ATTENTION_MASK] = attn
batch[OBS_LANGUAGE_CAUSAL_MARKS] = marks
else:
content = subtask if state_str is None else f"{subtask}, State: {state_str};"
text_batch = _build_text_batch(
self.policy,
[{"role": "user", "content": content}],
add_generation_prompt=False,
)
batch[OBS_LANGUAGE_TOKENS] = text_batch["lang_tokens"]
batch[OBS_LANGUAGE_ATTENTION_MASK] = text_batch["lang_masks"]
return self.policy.predict_action_chunk(batch)
def generate_text(
self,
kind: str,
observation: dict[str, Any] | None,
state: RuntimeState,
user_text: str | None = None,
) -> str:
messages = self.build_messages(kind, state, user_text=user_text)
if kind == "subtask" and getattr(self.policy.config, "joint_subtask_conditioning", False):
# Joint samples carry state on the task turn, so the subtask must be
# generated from the same state-bearing prompt.
import torch # noqa: PLC0415
from lerobot.policies.pi052.text_processor_pi052 import discretize_state_str # noqa: PLC0415
from lerobot.utils.constants import OBS_STATE # noqa: PLC0415
obs_state = (observation or {}).get(OBS_STATE)
if isinstance(obs_state, torch.Tensor) and obs_state.numel() > 0:
state_row = obs_state[0] if obs_state.ndim > 1 else obs_state
for m in reversed(messages):
if m.get("role") == "user":
m["content"] = f"{m.get('content', '')}, State: {discretize_state_str(state_row)};"
break
return _generate_with_policy(
self.policy,
messages,
observation=observation,
state=state,
label=f"{kind} gen",
min_new_tokens=self.gen.min_new_tokens,
temperature=self.gen.temperature,
top_p=self.gen.top_p,
suppress_loc_tokens=True, # all runtime text is prose; never emit <loc>
)
def build_messages(
self,
kind: str,
state: RuntimeState,
*,
user_text: str | None = None,
) -> list[dict[str, Any]]:
if kind in ("subtask", "plan"):
return [{"role": "user", "content": state.task or ""}]
if kind == "memory":
messages = [{"role": "user", "content": state.task or ""}]
if state.language_context.get("memory"):
messages.append(
{"role": "assistant", "content": f"Previous memory: {state.language_context['memory']}"}
)
if state.extra.get("prior_subtask"):
messages.append(
{"role": "user", "content": f"Completed subtask: {state.extra['prior_subtask']}"}
)
return messages
if kind == "interjection":
messages = [{"role": "user", "content": state.task or ""}]
if state.language_context.get("plan"):
messages.append(
{"role": "assistant", "content": f"Previous plan:\n{state.language_context['plan']}"}
)
if user_text:
messages.append({"role": "user", "content": user_text})
return messages
raise ValueError(f"Unknown PI052 text kind: {kind}")
def _get_loc_tokenizer(tok_name: str, auto_tokenizer_cls: Any, register_loc_fn: Any) -> Any:
tokenizer = _LOC_TOKENIZER_CACHE.get(tok_name)
if tokenizer is None:
tokenizer = register_loc_fn(auto_tokenizer_cls.from_pretrained(tok_name))
_LOC_TOKENIZER_CACHE[tok_name] = tokenizer
return tokenizer
def _build_text_batch(
policy: Any,
prompt_messages: list[dict[str, Any]],
*,
add_generation_prompt: bool = True,
) -> dict[str, Any]:
import torch # noqa: PLC0415
from transformers import AutoTokenizer # noqa: PLC0415
from lerobot.policies.pi052.text_processor_pi052 import ( # noqa: PLC0415
_flatten_say_tool_calls,
_format_messages,
_strip_blocks,
register_paligemma_loc_tokens,
)
tok_name = getattr(policy.config, "tokenizer_name", None) or "google/paligemma-3b-pt-224"
tokenizer = _get_loc_tokenizer(tok_name, AutoTokenizer, register_paligemma_loc_tokens)
messages = [_strip_blocks(_flatten_say_tool_calls(m)) for m in prompt_messages]
prompt, _spans = _format_messages(messages)
if add_generation_prompt:
# No trailing space: SentencePiece folds it into the first target token
# ("▁move"), so a space-suffixed prefill ends in a lone "▁" the model
# never saw at this position during training.
prompt = prompt + "Assistant:"
encoded = tokenizer(prompt, return_tensors="pt")
ids = encoded["input_ids"]
attn = encoded.get("attention_mask")
if attn is None and tokenizer.pad_token_id is not None:
attn = ids != tokenizer.pad_token_id
if attn is not None and hasattr(attn, "dtype") and attn.dtype != torch.bool:
attn = attn.bool()
device = getattr(getattr(policy, "config", None), "device", None)
if device is not None:
try:
ids = ids.to(device)
if attn is not None and hasattr(attn, "to"):
attn = attn.to(device)
except Exception as exc: # noqa: BLE001
logger.debug("could not move pi052 lang tokens to %s: %s", device, exc)
return {"lang_tokens": ids, "lang_masks": attn, "tokenizer": tokenizer}
def _generate_with_policy(
policy: Any,
messages: list[dict[str, Any]],
*,
observation: dict[str, Any] | None = None,
state: RuntimeState | None = None,
label: str = "select_message",
min_new_tokens: int = 0,
temperature: float = 0.0,
top_p: float = 1.0,
suppress_loc_tokens: bool = False,
) -> str:
if not hasattr(policy, "select_message"):
if state is not None:
state.log(f" [warn] policy has no select_message — skipping {label}")
return ""
text_batch = _build_text_batch(policy, messages)
try:
from lerobot.utils.constants import OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS # noqa: PLC0415
batch: dict[str, Any] = {
OBS_LANGUAGE_TOKENS: text_batch["lang_tokens"],
OBS_LANGUAGE_ATTENTION_MASK: text_batch["lang_masks"],
}
if observation:
for k, v in observation.items():
if isinstance(k, str) and k.startswith("observation.") and k not in batch:
batch[k] = v
return policy.select_message(
batch,
tokenizer=text_batch["tokenizer"],
min_new_tokens=min_new_tokens,
temperature=temperature,
top_p=top_p,
suppress_loc_tokens=suppress_loc_tokens,
)
except Exception as exc: # noqa: BLE001
logger.warning("%s failed: %s", label, exc, exc_info=logger.isEnabledFor(logging.DEBUG))
if state is not None:
state.log(f" [warn] {label} failed: {type(exc).__name__}: {exc}")
return ""
src/lerobot/policies/pi052/modeling_pi052.py
+2037
View File
File diff suppressed because it is too large Load Diff
src/lerobot/policies/pi052/processor_pi052.py
+164
View File
@@ -0,0 +1,164 @@
# 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.
"""PI052 processor factory with optional recipe rendering and text tokenization.
Without a recipe it delegates to the standard PI0.5 pipeline.
"""
from __future__ import annotations
from pathlib import Path
from typing import Any
import torch
from lerobot.configs.recipe import TrainingRecipe
from lerobot.processor import (
AbsoluteActionsProcessorStep,
ActionTokenizerProcessorStep,
AddBatchDimensionProcessorStep,
DeviceProcessorStep,
NormalizerProcessorStep,
PolicyAction,
PolicyProcessorPipeline,
RelativeActionsProcessorStep,
RenameObservationsProcessorStep,
UnnormalizerProcessorStep,
policy_action_to_transition,
transition_to_policy_action,
)
# Import directly to keep optional language dependencies out of ``lerobot.processor``.
from lerobot.processor.render_messages_processor import RenderMessagesStep
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
from ..pi05.processor_pi05 import make_pi05_pre_post_processors
from .configuration_pi052 import PI052Config
from .text_processor_pi052 import PI052TextTokenizerStep
def make_pi052_pre_post_processors(
config: PI052Config,
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
dataset_repo_id: str | None = None,
dataset_root: str | None = None,
dataset_revision: str | None = None,
episodes: list[int] | None = None,
exclude_episodes: list[int] | None = None,
) -> tuple[
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
PolicyProcessorPipeline[PolicyAction, PolicyAction],
]:
"""Build PI0.5-v2's pre/post-processor pipelines.
Falls through to π0.5's stock pipeline when ``recipe_path`` is unset.
"""
if not config.recipe_path:
if getattr(config, "enable_fast_action_loss", False):
raise ValueError("PI052 FAST action loss requires recipe_path to build action supervision.")
return make_pi05_pre_post_processors(config, dataset_stats=dataset_stats)
recipe = _load_recipe(config.recipe_path)
relative_step = RelativeActionsProcessorStep(
enabled=config.use_relative_actions,
exclude_joints=getattr(config, "relative_exclude_joints", []),
action_names=getattr(config, "action_feature_names", None),
)
input_steps = [
RenameObservationsProcessorStep(rename_map={}),
AddBatchDimensionProcessorStep(),
relative_step,
NormalizerProcessorStep(
features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping,
stats=dataset_stats,
),
RenderMessagesStep(recipe=recipe),
PI052TextTokenizerStep(
tokenizer_name="google/paligemma-3b-pt-224",
max_length=config.tokenizer_max_length,
plan_dropout_prob=getattr(config, "plan_dropout_prob", 0.0),
memory_dropout_prob=getattr(config, "memory_dropout_prob", 0.0),
subtask_dropout_prob=getattr(config, "subtask_dropout_prob", 0.0),
),
]
# Add FAST action-token supervision only when explicitly enabled.
if getattr(config, "enable_fast_action_loss", False):
from .fit_fast_tokenizer import resolve_fast_tokenizer # noqa: PLC0415
input_steps.append(
ActionTokenizerProcessorStep(
action_tokenizer_name=resolve_fast_tokenizer(
config,
dataset_repo_id,
dataset_root,
dataset_stats,
dataset_revision,
episodes,
exclude_episodes,
),
max_action_tokens=config.max_action_tokens,
fast_skip_tokens=config.fast_skip_tokens,
paligemma_tokenizer_name="google/paligemma-3b-pt-224",
allow_truncation=False,
)
)
input_steps.append(DeviceProcessorStep(device=config.device))
output_steps = [
UnnormalizerProcessorStep(
features=config.output_features,
norm_map=config.normalization_mapping,
stats=dataset_stats,
),
AbsoluteActionsProcessorStep(
enabled=config.use_relative_actions,
relative_step=relative_step,
),
DeviceProcessorStep(device="cpu"),
]
return (
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
steps=input_steps,
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
),
PolicyProcessorPipeline[PolicyAction, PolicyAction](
steps=output_steps,
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
to_transition=policy_action_to_transition,
to_output=transition_to_policy_action,
),
)
def _load_recipe(path_str: str) -> TrainingRecipe:
"""Resolve ``path_str`` to a ``TrainingRecipe``.
Accepts an absolute path or a path relative to
``src/lerobot/configs/``.
"""
p = Path(path_str)
if not p.is_absolute() and not p.exists():
from lerobot.configs import recipe as _recipe_module # noqa: PLC0415
configs_dir = Path(_recipe_module.__file__).resolve().parent
candidate = configs_dir / path_str
if candidate.exists():
p = candidate
return TrainingRecipe.from_yaml(p)
src/lerobot/policies/pi052/text_processor_pi052.py
+521
View File
@@ -0,0 +1,521 @@
# 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.
"""Tokenize PI052 messages and build text/action supervision masks."""
from __future__ import annotations
import json
import logging
from dataclasses import dataclass
from typing import Any
import numpy as np
import torch
from torch import Tensor
from lerobot.configs import PipelineFeatureType, PolicyFeature
from lerobot.processor.pipeline import ProcessorStep, ProcessorStepRegistry
from lerobot.types import EnvTransition, TransitionKey
from lerobot.utils.constants import OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS, OBS_STATE
logger = logging.getLogger(__name__)
def discretize_state_str(state_row: Any) -> str:
"""Format one normalized state row with PI0.5's 256-bin convention."""
arr = state_row.detach().cpu().numpy() if hasattr(state_row, "detach") else np.asarray(state_row)
disc = np.digitize(arr, bins=np.linspace(-1, 1, 256 + 1)[:-1]) - 1
return " ".join(str(int(x)) for x in disc.reshape(-1).tolist())
def _state_row_at(state_all: Any, pos: int) -> Any:
"""Select the per-sample state row from a (possibly batched) state tensor."""
if state_all is None:
return None
if hasattr(state_all, "ndim") and state_all.ndim >= 2:
return state_all[pos]
return state_all
def _content_to_text(content: Any) -> str:
"""Collapse a message's ``content`` (string or multimodal blocks) to text."""
if isinstance(content, str):
return content
if isinstance(content, list):
parts = [
b["text"]
for b in content
if isinstance(b, dict) and b.get("type") == "text" and isinstance(b.get("text"), str)
]
return "\n".join(parts)
return ""
def _flatten_say_tool_calls(message: dict[str, Any]) -> dict[str, Any]:
"""Move ``say`` tool calls into text markers that PaliGemma can learn."""
tool_calls = message.get("tool_calls")
if not tool_calls:
return message
say_texts: list[str] = []
for call in tool_calls:
if not isinstance(call, dict):
continue
fn = call.get("function") or {}
if fn.get("name") != "say":
continue
args = fn.get("arguments")
if isinstance(args, str):
try:
import json # noqa: PLC0415
args = json.loads(args)
except (ValueError, TypeError):
args = {}
text = args.get("text", "") if isinstance(args, dict) else ""
if text:
say_texts.append(str(text))
new = dict(message)
new.pop("tool_calls", None)
if not say_texts:
return new
base = _content_to_text(new.get("content")).strip()
marker = "".join(f"<say>{t}</say>" for t in say_texts)
new["content"] = f"{base}\n{marker}" if base else marker
return new
def _strip_blocks(message: dict[str, Any]) -> dict[str, Any]:
"""Flatten text blocks and drop image blocks handled by observation inputs."""
new = dict(message)
new.pop("stream", None)
new.pop("target", None)
content = new.get("content")
if content is None:
new["content"] = ""
elif isinstance(content, str):
pass
elif isinstance(content, list):
parts: list[str] = []
for block in content:
if not isinstance(block, dict):
continue
if block.get("type") == "text":
t = block.get("text", "")
if isinstance(t, str):
parts.append(t)
new["content"] = "\n".join(parts)
else:
new["content"] = str(content)
return new
def _is_batched_messages(messages: Any) -> bool:
return isinstance(messages, list) and bool(messages) and isinstance(messages[0], list)
def _sample_indices(value: Any, batch_size: int) -> list[int | None]:
if value is None:
return [None] * batch_size
if isinstance(value, torch.Tensor):
if value.numel() == 1:
return [int(value.item())] * batch_size
values = value.reshape(-1).tolist()
return [int(v) for v in values[:batch_size]]
if isinstance(value, (list, tuple)):
if len(value) == 1:
return _sample_indices(value[0], batch_size)
return [int(v.item() if hasattr(v, "item") else v) for v in value[:batch_size]]
return [int(value)] * batch_size
_VQA_COORD_SCALE = 1000.0
def register_paligemma_loc_tokens(tokenizer: Any) -> Any:
"""Register PaliGemma's reserved ``<locDDDD>`` strings as single tokens.
Without registration, the stock tokenizer splits each location into generic text pieces.
"""
if "<loc0000>" in getattr(tokenizer, "added_tokens_encoder", {}):
return tokenizer
tokenizer.add_tokens([f"<loc{i:04d}>" for i in range(1024)])
return tokenizer
def _loc_token(coord: float, scale: float = _VQA_COORD_SCALE) -> str:
"""PaliGemma ``<locNNNN>`` for a coord on a ``[0, scale]`` axis."""
idx = round(float(coord) / scale * 1023) if scale > 0 else 0
return f"<loc{max(0, min(1023, idx)):04d}>"
def _vqa_answer_to_loc(answer: dict[str, Any]) -> str | None:
"""Convert normalized bbox/keypoint answers to label-first PaliGemma locations.
Label-first targets prevent location tokens from dominating every assistant turn; non-spatial answers return ``None``.
"""
point = answer.get("point")
if isinstance(point, list | tuple) and len(point) == 2 and "point_format" in answer:
try:
x, y = float(point[0]), float(point[1])
except (TypeError, ValueError):
return None
label = str(answer.get("label", "")).strip()
if not label:
return None
return f"{label} {_loc_token(y)}{_loc_token(x)}"
detections = answer.get("detections")
if isinstance(detections, list) and detections:
parts: list[str] = []
for det in detections:
if not isinstance(det, dict):
continue
box = det.get("bbox")
if not (isinstance(box, list | tuple) and len(box) == 4):
continue
try:
x1, y1, x2, y2 = (float(v) for v in box)
except (TypeError, ValueError):
continue
label = str(det.get("label", "")).strip()
if not label:
continue
toks = f"{_loc_token(y1)}{_loc_token(x1)}{_loc_token(y2)}{_loc_token(x2)}"
parts.append(f"{label} {toks}")
return " ; ".join(parts) if parts else None
return None
def _messages_vqa_to_loc(
messages: list[dict[str, Any]],
target_indices: list[int],
) -> list[dict[str, Any]]:
"""Rewrite spatial VQA target JSON as camera-independent ``<loc>`` text."""
if not target_indices:
return messages
out = list(messages)
for idx in target_indices:
if not (0 <= idx < len(out)):
continue
content = out[idx].get("content")
if not isinstance(content, str) or not content.strip():
continue
try:
answer = json.loads(content)
except (ValueError, TypeError):
continue
if not isinstance(answer, dict):
continue
loc_text = _vqa_answer_to_loc(answer)
if loc_text is not None:
out[idx] = {**out[idx], "content": loc_text}
return out
def _format_messages(
messages: list[dict[str, Any]],
target_indices: list[int] | None = None,
eos_token: str | None = None,
) -> tuple[str, list[tuple[int, int]]]:
"""Build the flat PI0.5 prompt and each message's payload span.
Supervised targets include EOS so generation learns when to stop.
"""
targets = set(target_indices or [])
parts: list[str] = []
spans: list[tuple[int, int]] = []
cursor = 0
for i, m in enumerate(messages):
role = m.get("role", "user")
content = m.get("content", "") or ""
header = f"{role.capitalize()}: "
body = content + eos_token if (eos_token and i in targets) else content
full = header + body + "\n"
start = cursor + len(header)
end = start + len(body)
parts.append(full)
spans.append((start, end))
cursor += len(full)
return "".join(parts), spans
def encode_prompt_with_targets(
tokenizer: Any, messages: list[dict[str, Any]], target_indices: list[int]
) -> tuple[Tensor, Tensor, Tensor]:
"""Tokenize a flat prompt and mark the token positions of target spans.
Inference-side twin of ``PI052TextTokenizerStep._encode_messages``: same
serialization (role headers, target EOS) and the same offset-overlap span
arithmetic, but unpadded and returning a boolean target mask instead of
labels. Used to rebuild joint-sequence prompts whose target spans must be
attended causally, matching ``_mark_target_span_causal`` at train time.
Returns ``(input_ids, attention_mask, target_marks)``, each ``(1, L)``.
"""
prompt, spans = _format_messages(messages, target_indices, getattr(tokenizer, "eos_token", None))
encoded = tokenizer(prompt, return_tensors="pt", return_offsets_mapping=True)
input_ids = encoded["input_ids"][0]
attention_mask = encoded.get("attention_mask")
if attention_mask is None:
attention_mask = torch.ones_like(input_ids, dtype=torch.bool)
else:
attention_mask = attention_mask[0].bool()
offsets = encoded["offset_mapping"][0]
marks = torch.zeros_like(input_ids, dtype=torch.bool)
for idx in target_indices:
if idx >= len(spans):
continue
char_start, char_end = spans[idx]
for token_pos in range(input_ids.shape[0]):
if not attention_mask[token_pos]:
continue
tok_start, tok_end = int(offsets[token_pos, 0]), int(offsets[token_pos, 1])
if tok_end <= char_start or tok_start >= char_end:
continue
marks[token_pos] = True
return input_ids.unsqueeze(0), attention_mask.unsqueeze(0), marks.unsqueeze(0)
@dataclass
@ProcessorStepRegistry.register(name="pi052_text_tokenizer")
class PI052TextTokenizerStep(ProcessorStep):
"""Convert flat role-delimited messages into tokens and supervision masks."""
tokenizer_name: str = "google/paligemma-3b-pt-224"
max_length: int = 200
padding: str = "max_length"
padding_side: str = "right"
plan_dropout_prob: float = 0.0
memory_dropout_prob: float = 0.0
subtask_dropout_prob: float = 0.0
interjection_dropout_prob: float = 0.0
dropout_seed: int | None = None
def __post_init__(self) -> None:
self._tokenizer: Any = None
def get_config(self) -> dict[str, Any]:
return {
"tokenizer_name": self.tokenizer_name,
"max_length": self.max_length,
"padding": self.padding,
"padding_side": self.padding_side,
"plan_dropout_prob": self.plan_dropout_prob,
"memory_dropout_prob": self.memory_dropout_prob,
"subtask_dropout_prob": self.subtask_dropout_prob,
"interjection_dropout_prob": self.interjection_dropout_prob,
"dropout_seed": self.dropout_seed,
}
def _ensure_tokenizer(self) -> Any:
if self._tokenizer is not None:
return self._tokenizer
from transformers import AutoTokenizer # noqa: PLC0415
self._tokenizer = register_paligemma_loc_tokens(AutoTokenizer.from_pretrained(self.tokenizer_name))
return self._tokenizer
def __call__(self, transition: EnvTransition) -> EnvTransition | None:
transition = transition.copy()
complementary = transition.get(TransitionKey.COMPLEMENTARY_DATA, {}) or {}
messages = complementary.get("messages") or []
if not messages:
return transition
tokenizer = self._ensure_tokenizer()
state_all = (transition.get(TransitionKey.OBSERVATION) or {}).get(OBS_STATE)
if _is_batched_messages(messages):
indices_iter = _sample_indices(complementary.get("index"), len(messages))
encoded = [
self._encode_messages(
tokenizer,
msg,
list(streams),
list(tgt_indices),
complementary,
sample_idx=int(s_idx) if s_idx is not None else None,
state_row=_state_row_at(state_all, pos),
)
for pos, (msg, streams, tgt_indices, s_idx) in enumerate(
zip(
messages,
complementary.get("message_streams") or [[] for _ in messages],
complementary.get("target_message_indices") or [[] for _ in messages],
indices_iter,
strict=False,
)
)
]
else:
sample_idx = _sample_indices(complementary.get("index"), 1)[0]
encoded = [
self._encode_messages(
tokenizer,
messages,
list(complementary.get("message_streams") or []),
list(complementary.get("target_message_indices") or []),
complementary,
sample_idx=sample_idx,
state_row=_state_row_at(state_all, 0),
)
]
obs = dict(transition.get(TransitionKey.OBSERVATION) or {})
obs[OBS_LANGUAGE_TOKENS] = torch.stack([ids for ids, _, _, _, _ in encoded])
obs[OBS_LANGUAGE_ATTENTION_MASK] = torch.stack([attn for _, attn, _, _, _ in encoded])
transition[TransitionKey.OBSERVATION] = obs
transition[TransitionKey.COMPLEMENTARY_DATA] = {
**complementary,
"text_labels": torch.stack([labels for _, _, labels, _, _ in encoded]),
"predict_actions": torch.stack([pred for _, _, _, pred, _ in encoded]),
}
return transition
def _encode_messages(
self,
tokenizer: Any,
messages: list[dict[str, Any]],
message_streams: list[str | None],
target_indices: list[int],
complementary: dict[str, Any],
sample_idx: int | None = None,
state_row: Any = None,
) -> tuple[Tensor, Tensor, Tensor, Tensor, str]:
if (
self.plan_dropout_prob
or self.memory_dropout_prob
or self.subtask_dropout_prob
or self.interjection_dropout_prob
):
messages, target_indices = self._apply_prompt_dropout(
messages,
target_indices,
complementary,
sample_idx=sample_idx,
)
messages = _messages_vqa_to_loc(messages, target_indices)
messages = [_strip_blocks(_flatten_say_tool_calls(m)) for m in messages]
# Only low-level prompts carry PI0.5-style proprioception.
if state_row is not None and any(s == "low_level" for s in message_streams):
state_str = discretize_state_str(state_row)
for m in reversed(messages):
if m.get("role") == "user":
base = _content_to_text(m.get("content", ""))
m["content"] = f"{base}, State: {state_str};"
break
prompt, spans = _format_messages(messages, target_indices, getattr(tokenizer, "eos_token", None))
encoded = tokenizer(
prompt,
max_length=self.max_length,
padding=self.padding,
truncation=True,
return_tensors="pt",
return_offsets_mapping=True,
padding_side=self.padding_side,
)
input_ids = encoded["input_ids"][0]
attention_mask = encoded["attention_mask"][0].bool()
offsets = encoded["offset_mapping"][0]
labels = torch.full_like(input_ids, fill_value=-100)
for idx in target_indices:
if idx >= len(spans):
continue
char_start, char_end = spans[idx]
for token_pos in range(input_ids.shape[0]):
if not attention_mask[token_pos]:
continue
tok_start, tok_end = int(offsets[token_pos, 0]), int(offsets[token_pos, 1])
if tok_end <= char_start or tok_start >= char_end:
continue
labels[token_pos] = input_ids[token_pos]
predict_actions = torch.tensor(
bool(any(s == "low_level" for s in message_streams)),
dtype=torch.bool,
)
return input_ids, attention_mask, labels, predict_actions, prompt
def _apply_prompt_dropout(
self,
messages: list[dict[str, Any]],
target_indices: list[int],
complementary: dict[str, Any],
sample_idx: int | None = None,
) -> tuple[list[dict[str, Any]], list[int]]:
"""Drop sampled context messages and remap the retained target positions."""
import random # noqa: PLC0415
seed = self.dropout_seed
if seed is None:
seed_src = sample_idx if sample_idx is not None else complementary.get("index", 0)
try:
if hasattr(seed_src, "item"):
seed_src = seed_src.item()
seed = int(seed_src)
except (TypeError, ValueError):
seed = 0
rng = random.Random(seed)
keep_indices: list[int] = []
for idx, msg in enumerate(messages):
if idx in target_indices:
keep_indices.append(idx)
continue
kind = _classify_for_dropout(msg)
prob = {
"plan": self.plan_dropout_prob,
"memory": self.memory_dropout_prob,
"subtask": self.subtask_dropout_prob,
"interjection": self.interjection_dropout_prob,
}.get(kind, 0.0)
if prob > 0.0 and rng.random() < prob:
continue
keep_indices.append(idx)
new_messages = [messages[i] for i in keep_indices]
old_to_new = {old: new for new, old in enumerate(keep_indices)}
new_targets = [old_to_new[t] for t in target_indices if t in old_to_new]
return new_messages, new_targets
def transform_features(
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
return features
def _classify_for_dropout(message: dict[str, Any]) -> str | None:
"""Classify context from its rendered text prefix."""
content = message.get("content")
if isinstance(content, list):
text_parts = [b.get("text", "") for b in content if isinstance(b, dict) and b.get("type") == "text"]
content = " ".join(text_parts)
elif content is None or not isinstance(content, str):
return None
s = content.strip()
if s.startswith("Plan:") or s.startswith("Previous plan"):
return "plan"
if s.startswith("Memory:") or s.startswith("Previous memory"):
return "memory"
if s.startswith("Current subtask") or s.startswith("Completed subtask"):
return "subtask"
return None
src/lerobot/policies/pi0_fast/configuration_pi0_fast.py
+3
View File
@@ -61,6 +61,9 @@ class PI0FastConfig(PreTrainedConfig):
tokenizer_max_length: int = 200 # see openpi `__post_init__`
text_tokenizer_name: str = "google/paligemma-3b-pt-224"
action_tokenizer_name: str = "lerobot/fast-action-tokenizer"
auto_fit_fast_tokenizer: bool = False
fast_tokenizer_cache_dir: str = "~/.cache/lerobot/fast_tokenizers"
fast_tokenizer_fit_samples: int = 1024
temperature: float = 0.0
max_decoding_steps: int = 256
fast_skip_tokens: int = 128
src/lerobot/policies/pi0_fast/processor_pi0_fast.py
+18 -1
View File
@@ -92,6 +92,11 @@ class Pi0FastPrepareStateAndLanguageTokenizerProcessorStep(ProcessorStep):
def make_pi0_fast_pre_post_processors(
config: PI0FastConfig,
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
dataset_repo_id: str | None = None,
dataset_root: str | None = None,
dataset_revision: str | None = None,
episodes: list[int] | None = None,
exclude_episodes: list[int] | None = None,
) -> tuple[
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
PolicyProcessorPipeline[PolicyAction, PolicyAction],
@@ -136,6 +141,18 @@ def make_pi0_fast_pre_post_processors(
# state from the observation but does not change it. NormalizerProcessorStep still runs
# before Pi0FastPrepareStateAndLanguageTokenizerProcessorStep, so the state tokenizer
# continues to receive normalized state in [-1, 1] as expected.
from ..pi052.fit_fast_tokenizer import resolve_fast_tokenizer # noqa: PLC0415
action_tokenizer_path = resolve_fast_tokenizer(
config,
dataset_repo_id,
dataset_root,
dataset_stats,
dataset_revision,
episodes,
exclude_episodes,
)
input_steps: list[ProcessorStep] = [
steps.rename_observations, # To mimic the same processor as pretrained one
steps.add_batch_dim,
@@ -149,7 +166,7 @@ def make_pi0_fast_pre_post_processors(
padding="max_length",
),
ActionTokenizerProcessorStep(
action_tokenizer_name=config.action_tokenizer_name,
action_tokenizer_name=action_tokenizer_path,
max_action_tokens=config.max_action_tokens,
fast_skip_tokens=config.fast_skip_tokens,
paligemma_tokenizer_name=config.text_tokenizer_name,
src/lerobot/policies/pi_gemma.py
+389 -2
View File
@@ -14,18 +14,27 @@
from __future__ import annotations
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Literal
import torch
from torch import nn
from torch.nn import functional as F # noqa: N812
from lerobot.utils.import_utils import _transformers_available
# Default PaliGemma SigLIP input resolution. Mirrors
# ``pi05.configuration_pi05.DEFAULT_IMAGE_SIZE``; duplicated as a plain constant
# to avoid importing the pi05 package here (which would create an import cycle:
# pi_gemma -> pi05.__init__ -> modeling_pi05 -> pi_gemma).
DEFAULT_IMAGE_SIZE = 224
if TYPE_CHECKING or _transformers_available:
from transformers.cache_utils import DynamicCache
from transformers.masking_utils import create_causal_mask
from transformers.modeling_layers import GradientCheckpointingLayer
from transformers.modeling_outputs import BaseModelOutputWithPast
from transformers.models.auto import CONFIG_MAPPING
from transformers.models.gemma import modeling_gemma
from transformers.models.gemma.modeling_gemma import (
GemmaAttention,
GemmaConfig,
@@ -49,6 +58,8 @@ else:
GradientCheckpointingLayer = None
BaseModelOutputWithPast = None
create_causal_mask = None
CONFIG_MAPPING = None
modeling_gemma = None
def _gated_residual(
@@ -121,7 +132,10 @@ class PiGemmaRMSNorm(nn.Module):
if cond.shape[-1] != self.cond_dim:
raise ValueError(f"Expected cond dim {self.cond_dim}, got {cond.shape[-1]}")
modulation = self.dense(cond)
if len(x.shape) == 3:
# Per-sample cond (B, cond_dim) → broadcast over the sequence. A
# per-token cond (B, T, cond_dim) is already aligned with x and must
# not be unsqueezed (used by pi052's amortized K_repeat path).
if len(x.shape) == 3 and modulation.dim() == 2:
modulation = modulation.unsqueeze(1)
scale, shift, gate = modulation.chunk(3, dim=-1)
normed = normed * (1 + scale.float()) + shift.float()
@@ -275,6 +289,8 @@ class PiGemmaModel(GemmaModel): # type: ignore[misc]
# Convert to bfloat16 if the first layer uses bfloat16
if len(self.layers) > 0 and self.layers[0].self_attn.q_proj.weight.dtype == torch.bfloat16:
hidden_states = hidden_states.to(torch.bfloat16)
if causal_mask is not None and torch.is_floating_point(causal_mask):
causal_mask = causal_mask.to(dtype=hidden_states.dtype)
# create position embeddings to be shared across the decoder layers
position_embeddings = self.rotary_emb(hidden_states, position_ids)
@@ -367,3 +383,374 @@ __all__ = [
"PaliGemmaModelWithPiGemma",
"PaliGemmaForConditionalGenerationWithPiGemma",
]
# PI0.5 / PI052 dual-expert backbone: generic PaliGemma + Gemma action-expert
# transformer machinery used by the pi052 policy. GemmaVariantConfig is openpi's
# width/depth variant config (renamed from GemmaConfig to avoid clashing with
# transformers' GemmaConfig).
def sdpa_attention_forward(
module,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
attention_mask: torch.Tensor | None,
scaling: float,
dropout: float = 0.0,
):
"""Drop-in for ``modeling_gemma.eager_attention_forward`` using
``torch.nn.functional.scaled_dot_product_attention``.
PyTorch SDPA picks the memory-efficient kernel for arbitrary additive
bias masks (the FA backend only accepts causal/sliding-window). On
H100 that is ~1.3-1.7x faster and uses ~30-40% less attention memory
than the eager softmax(QK^T)+matmul path. Mirrors eager's signature
and output shape (``(B, Lq, H, D)``) so call sites are unchanged.
"""
n_rep = module.num_key_value_groups
if n_rep > 1:
key = key.repeat_interleave(n_rep, dim=1)
value = value.repeat_interleave(n_rep, dim=1)
if attention_mask is not None and attention_mask.dtype != query.dtype:
attention_mask = attention_mask.to(dtype=query.dtype)
attn_output = F.scaled_dot_product_attention(
query,
key,
value,
attn_mask=attention_mask,
dropout_p=dropout if module.training else 0.0,
is_causal=False,
scale=scaling,
)
return attn_output.transpose(1, 2).contiguous(), None
# Define the complete layer computation function for gradient checkpointing
def compute_layer_complete(
layer_idx, inputs_embeds, attention_mask, position_ids, adarms_cond, paligemma, gemma_expert
):
models = [paligemma.model.language_model, gemma_expert.model]
query_states = []
key_states = []
value_states = []
gates = []
for i, hidden_states in enumerate(inputs_embeds):
layer = models[i].layers[layer_idx]
hidden_states, gate = layernorm_forward(layer.input_layernorm, hidden_states, adarms_cond[i])
gates.append(gate)
input_shape = hidden_states.shape[:-1]
hidden_shape = (*input_shape, -1, layer.self_attn.head_dim)
query_state = layer.self_attn.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
key_state = layer.self_attn.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
value_state = layer.self_attn.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
query_states.append(query_state)
key_states.append(key_state)
value_states.append(value_state)
# Concatenate and process attention
query_states = torch.cat(query_states, dim=2)
key_states = torch.cat(key_states, dim=2)
value_states = torch.cat(value_states, dim=2)
dummy_tensor = torch.zeros(
query_states.shape[0],
query_states.shape[2],
query_states.shape[-1],
device=query_states.device,
dtype=query_states.dtype,
)
cos, sin = paligemma.model.language_model.rotary_emb(dummy_tensor, position_ids)
query_states, key_states = modeling_gemma.apply_rotary_pos_emb(
query_states, key_states, cos, sin, unsqueeze_dim=1
)
batch_size = query_states.shape[0]
scaling = paligemma.model.language_model.layers[layer_idx].self_attn.scaling
att_output, _ = sdpa_attention_forward(
paligemma.model.language_model.layers[layer_idx].self_attn,
query_states,
key_states,
value_states,
attention_mask,
scaling,
)
# Get head_dim from the current layer, not from the model
head_dim = paligemma.model.language_model.layers[layer_idx].self_attn.head_dim
att_output = att_output.reshape(batch_size, -1, 1 * 8 * head_dim)
# Process layer outputs
outputs_embeds = []
start_pos = 0
for i, hidden_states in enumerate(inputs_embeds):
layer = models[i].layers[layer_idx]
end_pos = start_pos + hidden_states.shape[1]
if att_output.dtype != layer.self_attn.o_proj.weight.dtype:
att_output = att_output.to(layer.self_attn.o_proj.weight.dtype)
out_emb = layer.self_attn.o_proj(att_output[:, start_pos:end_pos])
# first residual
out_emb = _gated_residual(hidden_states, out_emb, gates[i])
after_first_residual = out_emb.clone()
out_emb, gate = layernorm_forward(layer.post_attention_layernorm, out_emb, adarms_cond[i])
# Convert to bfloat16 if the next layer (mlp) uses bfloat16
if layer.mlp.up_proj.weight.dtype == torch.bfloat16:
out_emb = out_emb.to(dtype=torch.bfloat16)
out_emb = layer.mlp(out_emb)
# second residual
out_emb = _gated_residual(after_first_residual, out_emb, gate)
outputs_embeds.append(out_emb)
start_pos = end_pos
return outputs_embeds
class GemmaVariantConfig: # see openpi `gemma.py: Config`
"""Configuration for Gemma model variants."""
def __init__(self, width, depth, mlp_dim, num_heads, num_kv_heads, head_dim):
self.width = width
self.depth = depth
self.mlp_dim = mlp_dim
self.num_heads = num_heads
self.num_kv_heads = num_kv_heads
self.head_dim = head_dim
def get_gemma_config(variant: str) -> GemmaVariantConfig: # see openpi `gemma.py: get_config`
"""Returns config for specified gemma variant."""
if variant == "gemma_300m":
return GemmaVariantConfig(
width=1024,
depth=18,
mlp_dim=4096,
num_heads=8,
num_kv_heads=1,
head_dim=256,
)
elif variant == "gemma_2b":
return GemmaVariantConfig(
width=2048,
depth=18,
mlp_dim=16_384,
num_heads=8,
num_kv_heads=1,
head_dim=256,
)
else:
raise ValueError(f"Unknown variant: {variant}")
class PaliGemmaWithExpertModel(
nn.Module
): # see openpi `gemma_pytorch.py: PaliGemmaWithExpertModel` this class is almost a exact copy of PaliGemmaWithExpertModel in openpi
"""PaliGemma model with action expert for PI05."""
def __init__(
self,
vlm_config,
action_expert_config,
use_adarms=None,
precision: Literal["bfloat16", "float32"] = "bfloat16",
image_size: int = DEFAULT_IMAGE_SIZE,
freeze_vision_encoder: bool = False,
train_expert_only: bool = False,
):
if use_adarms is None:
use_adarms = [False, False]
super().__init__()
self.freeze_vision_encoder = freeze_vision_encoder
self.train_expert_only = train_expert_only
vlm_config_hf = CONFIG_MAPPING["paligemma"]()
vlm_config_hf._vocab_size = 257152 # noqa: SLF001
vlm_config_hf.image_token_index = 257152
vlm_config_hf.text_config.hidden_size = vlm_config.width
vlm_config_hf.text_config.intermediate_size = vlm_config.mlp_dim
vlm_config_hf.text_config.num_attention_heads = vlm_config.num_heads
vlm_config_hf.text_config.head_dim = vlm_config.head_dim
vlm_config_hf.text_config.num_hidden_layers = vlm_config.depth
vlm_config_hf.text_config.num_key_value_heads = vlm_config.num_kv_heads
vlm_config_hf.text_config.hidden_activation = "gelu_pytorch_tanh"
vlm_config_hf.text_config.dtype = "float32"
vlm_config_hf.text_config.vocab_size = 257152
vlm_config_hf.text_config.use_adarms = use_adarms[0]
vlm_config_hf.text_config.adarms_cond_dim = vlm_config.width if use_adarms[0] else None
vlm_config_hf.vision_config.image_size = image_size
vlm_config_hf.vision_config.intermediate_size = 4304
vlm_config_hf.vision_config.projection_dim = 2048
vlm_config_hf.vision_config.projector_hidden_act = "gelu_fast"
vlm_config_hf.vision_config.dtype = "float32"
action_expert_config_hf = CONFIG_MAPPING["gemma"](
head_dim=action_expert_config.head_dim,
hidden_size=action_expert_config.width,
intermediate_size=action_expert_config.mlp_dim,
num_attention_heads=action_expert_config.num_heads,
num_hidden_layers=action_expert_config.depth,
num_key_value_heads=action_expert_config.num_kv_heads,
vocab_size=257152,
hidden_activation="gelu_pytorch_tanh",
dtype="float32",
use_adarms=use_adarms[1],
adarms_cond_dim=action_expert_config.width if use_adarms[1] else None,
)
self.paligemma = PaliGemmaForConditionalGenerationWithPiGemma(config=vlm_config_hf)
self.gemma_expert = PiGemmaForCausalLM(config=action_expert_config_hf)
self.gemma_expert.model.embed_tokens = None
self.to_bfloat16_for_selected_params(precision)
self._set_requires_grad()
def to_bfloat16_for_selected_params(self, precision: Literal["bfloat16", "float32"] = "bfloat16"):
if precision == "bfloat16":
self.to(dtype=torch.bfloat16)
elif precision == "float32":
self.to(dtype=torch.float32)
return
else:
raise ValueError(f"Invalid precision: {precision}")
# Keep full vision path in float32 so we never toggle (toggle causes optimizer
# "same dtype" error). Saves memory vs full float32; more memory than only 3 params.
params_to_keep_float32 = [
"vision_tower",
"multi_modal_projector",
"lm_head",
"input_layernorm",
"post_attention_layernorm",
"model.norm",
]
for name, param in self.named_parameters():
if any(selector in name for selector in params_to_keep_float32):
param.data = param.data.to(dtype=torch.float32)
def _set_requires_grad(self):
if self.freeze_vision_encoder:
self.paligemma.model.vision_tower.eval()
for param in self.paligemma.model.vision_tower.parameters():
param.requires_grad = False
if self.train_expert_only:
self.paligemma.eval()
for param in self.paligemma.parameters():
param.requires_grad = False
def train(self, mode: bool = True):
super().train(mode)
if self.freeze_vision_encoder:
self.paligemma.model.vision_tower.eval()
if self.train_expert_only:
self.paligemma.eval()
def embed_image(self, image: torch.Tensor):
# Vision tower and multi_modal_projector are kept in float32 (params_to_keep_float32).
out_dtype = image.dtype
if image.dtype != torch.float32:
image = image.to(torch.float32)
image_outputs = self.paligemma.model.get_image_features(image)
# OpenPI / big_vision convention: image (soft) tokens are NOT scaled by the
# Gemma embedder normalizer (sqrt(hidden_size)) — only text tokens are. lerobot/pi05_base
# was trained in this regime, so scaling image features here over-scales them ~45x and
# breaks the pretrained vision-language alignment. Keep image features un-normalized.
features = image_outputs.pooler_output
if features.dtype != out_dtype:
features = features.to(out_dtype)
return features
def embed_language_tokens(self, tokens: torch.Tensor):
return self.paligemma.model.language_model.embed_tokens(tokens)
def forward(
self,
attention_mask: torch.Tensor | None = None,
position_ids: torch.LongTensor | None = None,
past_key_values: list[torch.FloatTensor] | None = None,
inputs_embeds: list[torch.FloatTensor] | None = None,
use_cache: bool | None = None,
adarms_cond: list[torch.Tensor] | None = None,
):
if adarms_cond is None:
adarms_cond = [None, None]
if inputs_embeds[1] is None:
prefix_output = self.paligemma.model.language_model.forward(
inputs_embeds=inputs_embeds[0],
attention_mask=attention_mask,
position_ids=position_ids,
past_key_values=past_key_values,
use_cache=use_cache,
adarms_cond=adarms_cond[0] if adarms_cond is not None else None,
)
prefix_past_key_values = prefix_output.past_key_values
prefix_output = prefix_output.last_hidden_state
suffix_output = None
elif inputs_embeds[0] is None:
suffix_output = self.gemma_expert.model.forward(
inputs_embeds=inputs_embeds[1],
attention_mask=attention_mask,
position_ids=position_ids,
past_key_values=past_key_values,
use_cache=use_cache,
adarms_cond=adarms_cond[1] if adarms_cond is not None else None,
)
suffix_output = suffix_output.last_hidden_state
prefix_output = None
prefix_past_key_values = None
else:
models = [self.paligemma.model.language_model, self.gemma_expert.model]
num_layers = self.paligemma.config.text_config.num_hidden_layers
# Check if gradient checkpointing is enabled for any of the models
use_gradient_checkpointing = (
hasattr(self.gemma_expert.model, "gradient_checkpointing")
and self.gemma_expert.model.gradient_checkpointing
and self.training
) or (hasattr(self, "gradient_checkpointing") and self.gradient_checkpointing and self.training)
# Process all layers with gradient checkpointing if enabled
for layer_idx in range(num_layers):
if use_gradient_checkpointing:
inputs_embeds = torch.utils.checkpoint.checkpoint(
compute_layer_complete,
layer_idx,
inputs_embeds,
attention_mask,
position_ids,
adarms_cond,
use_reentrant=False,
preserve_rng_state=False,
paligemma=self.paligemma,
gemma_expert=self.gemma_expert,
)
else:
inputs_embeds = compute_layer_complete(
layer_idx,
inputs_embeds,
attention_mask,
position_ids,
adarms_cond,
paligemma=self.paligemma,
gemma_expert=self.gemma_expert,
)
# final norm
def compute_final_norms(inputs_embeds, adarms_cond):
outputs_embeds = []
for i, hidden_states in enumerate(inputs_embeds):
out_emb, _ = layernorm_forward(models[i].norm, hidden_states, adarms_cond[i])
outputs_embeds.append(out_emb)
return outputs_embeds
# Apply gradient checkpointing to final norm if enabled
if use_gradient_checkpointing:
outputs_embeds = torch.utils.checkpoint.checkpoint(
compute_final_norms,
inputs_embeds,
adarms_cond,
use_reentrant=False,
preserve_rng_state=False,
)
else:
outputs_embeds = compute_final_norms(inputs_embeds, adarms_cond)
prefix_output = outputs_embeds[0]
suffix_output = outputs_embeds[1]
prefix_past_key_values = None
return [prefix_output, suffix_output], prefix_past_key_values
src/lerobot/scripts/lerobot_train.py
+56 -11
View File
@@ -20,9 +20,11 @@ Requires: pip install 'lerobot[training]' (includes dataset + accelerate + wand
import dataclasses
import logging
import os
import sys
import time
from contextlib import nullcontext
from datetime import timedelta
from pprint import pformat
from typing import TYPE_CHECKING, Any
@@ -91,6 +93,7 @@ def update_policy(
lr_scheduler=None,
lock=None,
sample_weighter=None,
log_metrics: bool = True,
) -> tuple[MetricsTracker, dict | None]:
"""
Performs a single training step to update the policy's weights.
@@ -108,6 +111,7 @@ def update_policy(
lr_scheduler: An optional learning rate scheduler.
lock: An optional lock for thread-safe optimizer updates.
sample_weighter: Optional SampleWeighter instance for per-sample loss weighting.
log_metrics: Whether to synchronize and record GPU metrics this step.
Returns:
A tuple containing:
@@ -175,12 +179,20 @@ def update_policy(
if has_method(accelerator.unwrap_model(policy, keep_fp32_wrapper=True), "update"):
accelerator.unwrap_model(policy, keep_fp32_wrapper=True).update()
train_metrics.loss = loss.item()
train_metrics.grad_norm = grad_norm.item()
train_metrics.lr = optimizer.param_groups[0]["lr"]
train_metrics.update_s = time.perf_counter() - start_time
if torch.cuda.is_available():
train_metrics.gpu_mem_gb = torch.cuda.max_memory_allocated() / (1024**3)
train_metrics.accumulate_tensor("loss", loss)
train_metrics.accumulate_tensor("grad_norm", grad_norm)
train_metrics.update_s = time.perf_counter() - start_time
# Synchronize accumulated GPU metrics only when logging.
if log_metrics:
train_metrics.materialize_tensors()
# Materialize detached loss components during the same logging synchronization.
if output_dict:
output_dict = {
k: (v.item() if isinstance(v, torch.Tensor) else v) for k, v in output_dict.items()
}
# Aggregate the policy's scalar outputs for logging and rank-reduction across the log window.
if output_dict:
train_metrics.update_metrics(output_dict)
@@ -211,7 +223,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
require_package("accelerate", extra="training")
from accelerate import Accelerator
from accelerate.utils import DistributedDataParallelKwargs, DistributedType
from accelerate.utils import DistributedDataParallelKwargs, DistributedType, InitProcessGroupKwargs
cfg.validate()
@@ -220,7 +232,16 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
# We set step_scheduler_with_optimizer=False to prevent accelerate from adjusting the lr_scheduler steps based on the num_processes
# We set find_unused_parameters=True to handle models with conditional computation
if accelerator is None:
ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
# Static graphs restore DDP overlap when conditional parameter usage is stable.
# Environment flags retain the existing defaults.
ddp_find_unused = os.environ.get("LEROBOT_DDP_FIND_UNUSED", "1") == "1"
ddp_static_graph = os.environ.get("LEROBOT_DDP_STATIC_GRAPH", "0") == "1"
ddp_kwargs = DistributedDataParallelKwargs(
find_unused_parameters=ddp_find_unused and not ddp_static_graph,
static_graph=ddp_static_graph,
)
# Allow rank 0 enough time to index large datasets before other ranks leave the barrier.
ipg_kwargs = InitProcessGroupKwargs(timeout=timedelta(hours=2))
# Accelerate auto-detects the device based on the available hardware and ignores the policy.device setting.
# Force the device to be CPU when the active config's device is set to CPU (works for both policy and reward model training).
force_cpu = cfg.trainable_config.device == "cpu"
@@ -230,7 +251,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
accelerator = Accelerator(
step_scheduler_with_optimizer=False,
mixed_precision=mixed_precision,
kwargs_handlers=[ddp_kwargs],
kwargs_handlers=[ddp_kwargs, ipg_kwargs],
cpu=force_cpu,
)
@@ -326,6 +347,14 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
active_cfg = cfg.trainable_config
processor_pretrained_path = active_cfg.pretrained_path
# A weight checkpoint may contain PI05 or differently configured PI052 processors.
if cfg.policy.type == "pi052" and processor_pretrained_path is not None and not cfg.resume:
logging.warning(
"pi052 is loading pretrained weights from %s, but building processors from the current "
"pi052 config so recipe text labels and FAST action labels are generated.",
processor_pretrained_path,
)
processor_pretrained_path = None
processor_kwargs = {}
if (processor_pretrained_path and not cfg.resume) or not processor_pretrained_path:
@@ -334,6 +363,13 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if cfg.is_reward_model_training:
processor_kwargs["dataset_meta"] = dataset.meta
if cfg.policy.type in {"pi0_fast", "pi052"}:
processor_kwargs["dataset_repo_id"] = cfg.dataset.repo_id
processor_kwargs["dataset_revision"] = cfg.dataset.revision
processor_kwargs["dataset_episodes"] = cfg.dataset.episodes
processor_kwargs["dataset_exclude_episodes"] = cfg.dataset.exclude_episodes
processor_kwargs["dataset_root"] = cfg.dataset.root
if not cfg.is_reward_model_training and processor_pretrained_path is not None:
preprocessor_overrides = {
"device_processor": {"device": device.type},
@@ -430,13 +466,17 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
# same permutation. accelerate then shards it disjointly across ranks via BatchSamplerShard
# without needing a `generator` attribute to synchronize an RNG, and resume is sample-exact.
shuffle = False
from_indices = dataset.meta.episodes["dataset_from_index"]
to_indices = dataset.meta.episodes["dataset_to_index"]
seed = cfg.seed if cfg.seed is not None else 0
sampler = EpisodeAwareSampler(
dataset.meta.episodes["dataset_from_index"],
dataset.meta.episodes["dataset_to_index"],
from_indices,
to_indices,
episode_indices_to_use=dataset.episodes,
drop_n_last_frames=getattr(active_cfg, "drop_n_last_frames", 0),
shuffle=True,
seed=cfg.seed if cfg.seed is not None else 0,
seed=seed,
absolute_to_relative_idx=dataset.absolute_to_relative_idx,
)
if cfg.resume and step > 0:
@@ -583,7 +623,10 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
batch = preprocessor(batch)
train_tracker.dataloading_s = time.perf_counter() - start_time
train_tracker, _ = update_policy(
# Synchronize GPU metrics only for updates that will be logged.
log_metrics = cfg.log_freq > 0 and (step + 1) % cfg.log_freq == 0
train_tracker, output_dict = update_policy(
train_tracker,
policy,
batch,
@@ -592,6 +635,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
accelerator=accelerator,
lr_scheduler=lr_scheduler,
sample_weighter=sample_weighter,
log_metrics=log_metrics,
)
# Note: eval and checkpoint happens *after* the `step`th training update has completed, so we
@@ -692,10 +736,11 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if is_main_process:
step_id = get_step_identifier(step, cfg.steps)
logging.info(f"Eval policy at step {step}")
eval_target_policy = accelerator.unwrap_model(policy)
with torch.no_grad(), accelerator.autocast():
eval_info = eval_policy_all(
envs=eval_env, # dict[suite][task_id] -> vec_env
policy=accelerator.unwrap_model(policy),
policy=eval_target_policy,
env_preprocessor=env_preprocessor,
env_postprocessor=env_postprocessor,
preprocessor=preprocessor,
src/lerobot/utils/logging_utils.py
+23 -1
View File
@@ -105,6 +105,8 @@ class MetricsTracker:
"epochs",
"accelerator",
"_caller_metrics",
"_tensor_sums",
"_tensor_counts",
]
def __init__(
@@ -133,6 +135,8 @@ class MetricsTracker:
# Meter names the caller registered up front. update_metrics() leaves these untouched, so a
# policy that echoes e.g. "loss" in its output dict can't clobber the aggregated meter.
self._caller_metrics: set[str] = set(self.metrics)
self._tensor_sums: dict[str, torch.Tensor] = {}
self._tensor_counts: dict[str, int] = {}
def __getattr__(self, name: str) -> int | dict[str, AverageMeter] | AverageMeter | Any:
if name in self.__dict__:
@@ -160,6 +164,22 @@ class MetricsTracker:
self.episodes = self.samples / self._avg_samples_per_ep
self.epochs = self.samples / self._num_frames
def accumulate_tensor(self, name: str, value: torch.Tensor) -> None:
"""Accumulate a detached metric on-device until the next logging step."""
if name not in self.metrics:
raise KeyError(f"Unknown metric {name!r}.")
value = value.detach()
self._tensor_sums[name] = self._tensor_sums.get(name, torch.zeros_like(value)) + value
self._tensor_counts[name] = self._tensor_counts.get(name, 0) + 1
def materialize_tensors(self) -> None:
"""Transfer pending tensor averages to their meters with one sync per metric."""
for name, total in self._tensor_sums.items():
count = self._tensor_counts[name]
self.metrics[name].update((total / count).item(), n=count)
self._tensor_sums.clear()
self._tensor_counts.clear()
def update_metrics(self, values: dict[str, Any]) -> None:
"""Accumulate a dict of scalar metrics, auto-registering a meter for each new key.
@@ -167,7 +187,7 @@ class MetricsTracker:
Caller-registered metrics (those passed to the constructor) are never overridden.
"""
for name, value in values.items():
if isinstance(value, bool) or not isinstance(value, (int, float)):
if isinstance(value, bool) or not isinstance(value, int | float):
continue
if name in self._caller_metrics:
continue
@@ -235,3 +255,5 @@ class MetricsTracker:
"""Resets average meters."""
for m in self.metrics.values():
m.reset()
self._tensor_sums.clear()
self._tensor_counts.clear()
tests/policies/pi052/test_pi052_attention_masking.py
+151
View File
@@ -0,0 +1,151 @@
#!/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.
"""Attention-masking tests for the PI052 (π0.5 v2) text head.
Regression coverage for the text-CE collapse bug: PaliGemma's
``embed_prefix`` flags every language token ``att=0``, which
``make_att_2d_masks`` turns into one fully *bidirectional* block. Under
that mask the text cross-entropy degenerates into a copy task — a
supervised target token attends to the tokens it is trained to predict —
and the LM head never learns causal generation, so ``select_message``
collapses at inference.
``_mark_target_span_causal`` sets ``att=1`` on the supervised target
language positions so each target token attends causally among the
targets while staying bidirectional to images + the user prompt. These
tests pin that behaviour for the PaliGemma prefix layout.
"""
import pytest
import torch
# modeling_pi052 / modeling_pi05 import transformers transitively.
pytest.importorskip("transformers")
from lerobot.policies.pi05.modeling_pi05 import make_att_2d_masks # noqa: E402
from lerobot.policies.pi052.modeling_pi052 import ( # noqa: E402
_mark_target_span_causal,
_shifted_lin_ce,
)
def _shifted_ce(logits, labels):
"""Adapter: ``_shifted_lin_ce`` is Liger-fused (hidden @ lm_head_weightᵀ).
An identity ``lm_head_weight`` makes the computed logits equal ``logits``.
Liger's Triton kernel is GPU-only, so inputs run on CUDA; the loss is
returned on CPU so grad still flows back to the CPU ``logits`` leaf.
"""
if not torch.cuda.is_available():
pytest.skip("Liger fused CE requires CUDA")
vocab_size = logits.shape[-1]
eye = torch.eye(vocab_size, dtype=logits.dtype, device="cuda")
return _shifted_lin_ce(logits.cuda(), eye, labels.cuda()).cpu()
# Synthetic prefix: two image tokens, three prompt tokens, and four supervised target tokens.
# Text labels mask the prompt with -100 and cover the target through the prefix end.
N_IMAGE = 2
N_PROMPT = 3
N_TARGET = 4
LANG_START = N_IMAGE
LANG_END = N_IMAGE + N_PROMPT + N_TARGET # = prefix length
PREFIX_LEN = LANG_END
def _embed_prefix_att_masks() -> torch.Tensor:
"""Mimic PaliGemma ``embed_prefix``: images + lang all att=0."""
return torch.zeros(1, PREFIX_LEN, dtype=torch.bool)
def _text_labels() -> torch.Tensor:
"""-100 over the prompt span, real ids over the target span."""
labels = torch.full((1, N_PROMPT + N_TARGET), -100, dtype=torch.long)
labels[0, N_PROMPT:] = torch.arange(10, 10 + N_TARGET)
return labels
def _attends(prefix_att_masks: torch.Tensor) -> torch.Tensor:
"""2D boolean attendance matrix; ``[i, j]`` True ⇒ i attends to j."""
pad = torch.ones(1, PREFIX_LEN, dtype=torch.bool)
return make_att_2d_masks(pad, prefix_att_masks)[0]
def test_mark_sets_att_on_targets_only():
"""Only the supervised target language positions flip to att=1."""
marked = _mark_target_span_causal(_embed_prefix_att_masks(), _text_labels(), LANG_START, LANG_END)
expected = [False] * PREFIX_LEN
for i in range(LANG_START + N_PROMPT, LANG_END): # target span
expected[i] = True
assert marked[0].tolist() == expected
def test_target_tokens_attend_causally_among_themselves():
"""A target token must NOT attend to later targets, but must attend
to earlier ones — genuine causal next-token prediction."""
marked = _mark_target_span_causal(_embed_prefix_att_masks(), _text_labels(), LANG_START, LANG_END)
attends = _attends(marked)
tgt = range(LANG_START + N_PROMPT, LANG_END)
for i in tgt:
for j in tgt:
if j > i:
assert not attends[i, j], f"target {i} must not see future target {j}"
else:
assert attends[i, j], f"target {i} must see earlier/self target {j}"
def test_target_tokens_attend_prompt_and_images_bidirectionally():
"""Targets keep full visibility of images + the user prompt."""
marked = _mark_target_span_causal(_embed_prefix_att_masks(), _text_labels(), LANG_START, LANG_END)
attends = _attends(marked)
context = list(range(0, LANG_START + N_PROMPT)) # images + prompt
for i in range(LANG_START + N_PROMPT, LANG_END):
for j in context:
assert attends[i, j], f"target {i} must attend context {j}"
def test_non_target_subtask_stays_bidirectional():
"""A flow-only / non-target language span (all -100 labels) leaves the
mask untouched — the action expert reads it bidirectionally."""
all_ignored = torch.full((1, N_PROMPT + N_TARGET), -100, dtype=torch.long)
marked = _mark_target_span_causal(_embed_prefix_att_masks(), all_ignored, LANG_START, LANG_END)
assert torch.equal(marked, _embed_prefix_att_masks())
def test_unmarked_mask_is_bidirectional_the_bug():
"""Documents the bug the fix prevents: without ``_mark_target_span_causal``
a target token attends *bidirectionally* to later targets — the
text-CE can copy the answer it is trained to predict."""
attends = _attends(_embed_prefix_att_masks())
first_tgt = LANG_START + N_PROMPT
last_tgt = LANG_END - 1
assert attends[first_tgt, last_tgt], (
"raw embed_prefix mask is bidirectional over language — the first "
"target token can see the last, which is the collapse bug"
)
def test_shifted_ce_returns_zero_when_no_text_positions_are_supervised():
pytest.importorskip("liger_kernel")
logits = torch.randn(2, 4, 8, requires_grad=True)
labels = torch.full((2, 4), -100, dtype=torch.long)
loss = _shifted_ce(logits, labels)
assert loss.item() == 0
loss.backward()
assert logits.grad is not None
tests/policies/pi052/test_pi052_bucketed_ce.py
+146
View File
@@ -0,0 +1,146 @@
#!/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 pytest
import torch
pytest.importorskip("transformers")
from lerobot.policies.pi052.modeling_pi052 import _lin_ce_flat, _shifted_lin_ce
def test_shifted_ce_none_retains_distinct_per_sample_losses():
hidden = torch.tensor(
[
[[8.0, 0.0], [0.0, 8.0], [0.0, 0.0]],
[[0.0, 8.0], [8.0, 0.0], [0.0, 0.0]],
]
)
labels = torch.tensor([[0, 0, 1], [0, 0, 1]])
losses = _shifted_lin_ce(hidden, torch.eye(2), labels, reduction="none")
assert losses.shape == (2,)
assert losses[0] < losses[1]
def test_checkpoint_resolution_forwards_explicit_hub_options(monkeypatch, tmp_path):
import lerobot.policies.pi05.modeling_pi05 as modeling_pi05
checkpoint = tmp_path / "model.safetensors"
checkpoint.touch()
calls = []
def fake_cached_file(model_id, filename, **kwargs):
calls.append((model_id, filename, kwargs))
return None if filename.endswith("index.json") else str(checkpoint)
monkeypatch.setattr(modeling_pi05, "cached_file", fake_cached_file)
files = modeling_pi05._resolve_weight_files(
"org/model",
force_download=True,
resume_download=True,
proxies={"https": "proxy"},
token="secret",
cache_dir=tmp_path / "cache",
local_files_only=True,
revision="commit",
)
assert files == [checkpoint]
for _model_id, _filename, kwargs in calls:
assert kwargs["revision"] == "commit"
assert kwargs["cache_dir"] == tmp_path / "cache"
assert kwargs["force_download"] is True
assert kwargs["resume_download"] is True
assert kwargs["proxies"] == {"https": "proxy"}
assert kwargs["token"] == "secret"
assert kwargs["local_files_only"] is True
def test_checkpoint_resolution_rejects_local_directory_without_weights(tmp_path):
import lerobot.policies.pi05.modeling_pi05 as modeling_pi05
with pytest.raises(FileNotFoundError, match="model.safetensors"):
modeling_pi05._resolve_weight_files(
tmp_path,
force_download=False,
resume_download=None,
proxies=None,
token=None,
cache_dir=None,
local_files_only=False,
revision=None,
)
@pytest.mark.parametrize("z_loss_weight", [0.0, 1e-4])
@pytest.mark.parametrize("rows,valid_rows", [(24, 9), (48, 25)])
def test_bucketed_ce_matches_dense_loss_and_gradients(z_loss_weight, rows, valid_rows):
generator = torch.Generator().manual_seed(23)
hidden_size, vocab_size = 7, 19
hidden_ref = torch.randn(rows, hidden_size, generator=generator, dtype=torch.float64, requires_grad=True)
weight_ref = torch.randn(
vocab_size, hidden_size, generator=generator, dtype=torch.float64, requires_grad=True
)
labels = torch.full((rows,), -100, dtype=torch.long)
valid_indices = torch.randperm(rows, generator=generator)[:valid_rows]
labels[valid_indices] = torch.randint(0, vocab_size, (valid_rows,), generator=generator)
hidden_bucketed = hidden_ref.detach().clone().requires_grad_(True)
weight_bucketed = weight_ref.detach().clone().requires_grad_(True)
import lerobot.policies.pi052.modeling_pi052 as modeling_pi052
loss_ref = _lin_ce_flat(hidden_ref, weight_ref, labels, z_loss_weight=z_loss_weight)
old_limit = modeling_pi052._LOGITS_CE_MAX_POSITIONS
modeling_pi052._LOGITS_CE_MAX_POSITIONS = 16
try:
loss_bucketed = _lin_ce_flat(
hidden_bucketed,
weight_bucketed,
labels,
z_loss_weight=z_loss_weight,
)
finally:
modeling_pi052._LOGITS_CE_MAX_POSITIONS = old_limit
loss_ref.backward()
loss_bucketed.backward()
torch.testing.assert_close(loss_bucketed, loss_ref, rtol=1e-6, atol=1e-6)
torch.testing.assert_close(hidden_bucketed.grad, hidden_ref.grad, rtol=1e-12, atol=1e-12)
torch.testing.assert_close(weight_bucketed.grad, weight_ref.grad, rtol=1e-12, atol=1e-12)
def test_bucketed_ce_all_ignored_preserves_zero_gradients():
hidden = torch.randn(24, 7, dtype=torch.float64, requires_grad=True)
weight = torch.randn(19, 7, dtype=torch.float64, requires_grad=True)
labels = torch.full((24,), -100, dtype=torch.long)
import lerobot.policies.pi052.modeling_pi052 as modeling_pi052
old_limit = modeling_pi052._LOGITS_CE_MAX_POSITIONS
modeling_pi052._LOGITS_CE_MAX_POSITIONS = 16
try:
loss = _lin_ce_flat(hidden, weight, labels)
finally:
modeling_pi052._LOGITS_CE_MAX_POSITIONS = old_limit
loss.backward()
assert loss.item() == 0.0
assert hidden.grad is not None
assert weight.grad is not None
assert torch.count_nonzero(hidden.grad) == 0
assert torch.count_nonzero(weight.grad) == 0
tests/policies/pi052/test_pi052_checkpoint_portability.py
+152
View File
@@ -0,0 +1,152 @@
#!/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 json
import shutil
from dataclasses import asdict
from types import SimpleNamespace
import numpy as np
import pytest
import torch
from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature
from lerobot.configs.recipe import MessageTurn, TrainingRecipe
from lerobot.policies import make_pre_post_processors
from lerobot.processor import ActionTokenizerProcessorStep, DataProcessorPipeline, NormalizerProcessorStep
from lerobot.processor.converters import identity_transition
from lerobot.processor.render_messages_processor import RenderMessagesStep
from lerobot.utils.constants import ACTION
class _ActionTokenizer:
def __call__(self, actions):
return np.asarray(actions).round().astype(np.int64)
def save_pretrained(self, path):
path.mkdir(parents=True)
(path / "processor_config.json").write_text('{"processor_class": "_ActionTokenizer"}\n')
class _PaligemmaTokenizer:
vocab_size = 4096
bos_token_id = 2
def encode(self, text, **kwargs):
return [10, 11] if text == "Action: " else [12]
def _make_pipeline(action_tokenizer_path):
recipe = TrainingRecipe(
messages=[
MessageTurn(role="user", content="${task}", stream="high_level"),
MessageTurn(role="assistant", content="${subtask}", stream="low_level", target=True),
]
)
stats = {ACTION: {"min": torch.tensor([-1.0, -2.0]), "max": torch.tensor([1.0, 2.0])}}
normalizer = NormalizerProcessorStep(
features={ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(2,))},
norm_map={FeatureType.ACTION: NormalizationMode.MIN_MAX},
stats=stats,
)
action_tokenizer = ActionTokenizerProcessorStep(
action_tokenizer_name=str(action_tokenizer_path),
max_action_tokens=16,
fast_skip_tokens=128,
)
return DataProcessorPipeline(
[normalizer, RenderMessagesStep(recipe), action_tokenizer],
name="policy_preprocessor",
to_transition=identity_transition,
to_output=identity_transition,
)
def test_pi052_pipeline_embeds_and_loads_fitted_action_tokenizer(tmp_path, monkeypatch):
original_cache = tmp_path / "original_fast_cache"
original_cache.mkdir()
tokenizer = _ActionTokenizer()
monkeypatch.setattr(
"lerobot.processor.tokenizer_processor.AutoProcessor.from_pretrained",
lambda path, **kwargs: tokenizer,
)
monkeypatch.setattr(
"lerobot.processor.tokenizer_processor.AutoTokenizer.from_pretrained",
lambda *args, **kwargs: _PaligemmaTokenizer(),
)
monkeypatch.setattr(
"lerobot.policies.pi052.fit_fast_tokenizer.fit_fast_tokenizer",
lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("FAST fitting must not run")),
)
pipeline = _make_pipeline(original_cache)
expected_tokens = pipeline.steps[-1]._tokenize_action(torch.tensor([[[0.2, 0.8]]]))[0]
expected_recipe = asdict(pipeline.steps[1].recipe)
expected_state = pipeline.steps[0].state_dict()
checkpoint = tmp_path / "checkpoint"
pipeline.save_pretrained(checkpoint)
DataProcessorPipeline(
[],
name="policy_postprocessor",
to_transition=identity_transition,
to_output=identity_transition,
).save_pretrained(checkpoint)
saved_config = json.loads((checkpoint / "policy_preprocessor.json").read_text())
tokenizer_step = saved_config["steps"][2]
assert tokenizer_step["config"]["action_tokenizer_name"] == "action_tokenizer"
assert tokenizer_step["artifacts"] == {"action_tokenizer_name": "action_tokenizer"}
assert (checkpoint / "action_tokenizer" / "processor_config.json").is_file()
shutil.rmtree(original_cache)
loaded, _ = make_pre_post_processors(
SimpleNamespace(type="pi052", auto_fit_fast_tokenizer=True),
pretrained_path=str(checkpoint),
dataset_repo_id="org/dataset-that-must-not-be-read",
)
assert asdict(loaded.steps[1].recipe) == expected_recipe
for key, tensor in expected_state.items():
torch.testing.assert_close(loaded.steps[0].state_dict()[key], tensor)
torch.testing.assert_close(
loaded.steps[-1]._tokenize_action(torch.tensor([[[0.2, 0.8]]]))[0],
expected_tokens,
)
def test_pi052_pipeline_rejects_missing_fitted_action_tokenizer(tmp_path, monkeypatch):
tokenizer = _ActionTokenizer()
monkeypatch.setattr(
"lerobot.processor.tokenizer_processor.AutoProcessor.from_pretrained",
lambda path, **kwargs: tokenizer,
)
monkeypatch.setattr(
"lerobot.processor.tokenizer_processor.AutoTokenizer.from_pretrained",
lambda *args, **kwargs: _PaligemmaTokenizer(),
)
pipeline = _make_pipeline(tmp_path / "original_fast_cache")
checkpoint = tmp_path / "checkpoint"
pipeline.save_pretrained(checkpoint)
shutil.rmtree(checkpoint / "action_tokenizer")
with pytest.raises(FileNotFoundError, match="Checkpoint artifacts are incomplete"):
DataProcessorPipeline.from_pretrained(
checkpoint,
config_filename="policy_preprocessor.json",
to_transition=identity_transition,
to_output=identity_transition,
)
tests/policies/pi052/test_pi052_fast_action_loss.py
+162
View File
@@ -0,0 +1,162 @@
#!/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.
"""Regression tests for PI052 FAST action-code supervision."""
from types import SimpleNamespace
import pytest
import torch
from torch import nn
from torch.nn import functional as F # noqa: N812
pytest.importorskip("transformers")
pytest.importorskip("liger_kernel")
from lerobot.policies.pi052.modeling_pi052 import PI052Policy, _fast_lin_ce # noqa: E402
from lerobot.policies.pi052.processor_pi052 import make_pi052_pre_post_processors # noqa: E402
def _fast_ce(logits, action_tokens, action_code_mask, predict_actions_t):
"""Adapter: ``_fast_lin_ce`` is Liger-fused (hidden @ lm_head_weightᵀ).
Feeding an identity ``lm_head_weight`` makes the computed logits equal the
provided ``logits``, so these regression tests exercise the masking/gating
logic exactly as before the fused-CE refactor. Liger's Triton kernel is
GPU-only, so inputs are moved to CUDA and the loss is returned on CPU
(keeping grad flowing back to the CPU ``logits`` leaf).
"""
if not torch.cuda.is_available():
pytest.skip("Liger fused CE requires CUDA")
vocab_size = logits.shape[-1]
eye = torch.eye(vocab_size, dtype=logits.dtype, device="cuda")
predict = predict_actions_t.cuda() if predict_actions_t is not None else None
loss = _fast_lin_ce(logits.cuda(), eye, action_tokens.cuda(), action_code_mask.cuda(), predict)
return loss.cpu()
def test_fast_ce_supervises_only_discrete_action_codes():
"""Wrapper tokens can be wrong without affecting the FAST action-code loss."""
vocab_size = 8
action_tokens = torch.tensor([[1, 2, 3, 4, 5, 0]])
action_code_mask = torch.tensor([[False, False, True, True, False, False]])
logits = torch.zeros(1, action_tokens.shape[1], vocab_size)
# Deliberately bad wrapper-token predictions. These should be ignored.
logits[0, 0, 7] = 10.0 # target would be token 2
logits[0, 3, 7] = 10.0 # target would be delimiter token 5
# Correct action-code predictions: hidden t predicts target t + 1.
logits[0, 1, 3] = 10.0
logits[0, 2, 4] = 10.0
loss = _fast_ce(logits, action_tokens, action_code_mask, predict_actions_t=None)
expected = F.cross_entropy(
torch.stack([logits[0, 1], logits[0, 2]]),
torch.tensor([3, 4]),
reduction="mean",
)
# Allow the fused GPU kernel's ~1e-7 difference on small losses.
assert torch.allclose(loss, expected, atol=1e-5, rtol=1e-3)
def test_fast_ce_masks_non_action_samples():
"""Recipe samples with predict_actions=False do not contribute FAST loss."""
vocab_size = 8
action_tokens = torch.tensor([[1, 2, 3, 4], [1, 2, 5, 6]])
action_code_mask = torch.tensor([[False, False, True, True], [False, False, True, True]])
predict_actions = torch.tensor([True, False])
logits = torch.zeros(2, action_tokens.shape[1], vocab_size)
logits[0, 1, 3] = 10.0
logits[0, 2, 4] = 10.0
# Bad predictions in the masked sample should not matter.
logits[1, 1, 7] = 10.0
logits[1, 2, 7] = 10.0
loss = _fast_ce(logits, action_tokens, action_code_mask, predict_actions)
expected = F.cross_entropy(
torch.stack([logits[0, 1], logits[0, 2]]),
torch.tensor([3, 4]),
reduction="mean",
)
# Allow the fused GPU kernel's ~1e-7 difference on small losses.
assert torch.allclose(loss, expected, atol=1e-5, rtol=1e-3)
def test_fast_ce_returns_zero_when_no_action_code_positions_are_valid():
logits = torch.randn(2, 4, 8, requires_grad=True)
action_tokens = torch.tensor([[1, 2, 3, 4], [1, 2, 5, 6]])
action_code_mask = torch.zeros_like(action_tokens, dtype=torch.bool)
loss = _fast_ce(logits, action_tokens, action_code_mask, predict_actions_t=None)
assert loss.item() == 0
loss.backward()
assert logits.grad is not None
def test_fast_ce_averages_each_action_sample_equally():
torch.manual_seed(0)
hidden = torch.randn(2, 5, 8)
lm_head_weight = torch.eye(8)
action_tokens = torch.tensor([[1, 2, 0, 0, 0], [1, 3, 4, 5, 6]])
action_code_mask = torch.tensor([[False, True, False, False, False], [False, True, True, True, True]])
loss = _fast_lin_ce(
hidden,
lm_head_weight,
action_tokens,
action_code_mask,
predict_actions_t=None,
reduction="mean",
)
per_sample = _fast_lin_ce(
hidden,
lm_head_weight,
action_tokens,
action_code_mask,
predict_actions_t=None,
reduction="none",
)
assert torch.allclose(loss, per_sample.mean())
def test_pi052_rejects_fast_loss_without_recipe():
config = SimpleNamespace(recipe_path=None, enable_fast_action_loss=True)
with pytest.raises(ValueError, match="recipe_path"):
make_pi052_pre_post_processors(config)
def test_pi052_rejects_missing_fast_batch_keys():
policy = PI052Policy.__new__(PI052Policy)
nn.Module.__init__(policy)
policy.config = SimpleNamespace(
enable_fast_action_loss=True,
fast_action_loss_weight=1.0,
flow_loss_weight=0.0,
text_loss_weight=1.0,
)
batch = {
"text_labels": torch.tensor([[1, 2]]),
"predict_actions": torch.tensor([True]),
}
with pytest.raises(ValueError, match="FAST action loss is enabled"):
policy.forward(batch)
tests/policies/pi052/test_pi052_fit_fast_tokenizer.py
+122
View File
@@ -0,0 +1,122 @@
#!/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.policies.pi052.fit_fast_tokenizer import (
_apply_relative_actions,
_dataset_signature,
_is_global_leader,
_normalize_actions,
_select_episode_indices,
_validate_fast_reconstruction,
)
def test_fast_tokenizer_fit_uses_training_mean_std_normalization():
actions = np.array([[[1.0, 7.0], [3.0, 3.0]]], dtype=np.float32)
stats = {"mean": [2.0, 5.0], "std": [0.5, 2.0]}
normalized = _normalize_actions(actions, "MEAN_STD", stats)
np.testing.assert_allclose(normalized, [[[-2.0, 1.0], [2.0, -1.0]]])
def test_fast_tokenizer_fit_quantiles_match_training_without_clipping():
actions = np.array([[[-1.0], [3.0]]], dtype=np.float32)
stats = {"q01": [0.0], "q99": [2.0]}
normalized = _normalize_actions(actions, "QUANTILES", stats)
np.testing.assert_allclose(normalized, [[[-2.0], [2.0]]])
def test_fast_tokenizer_cache_signature_tracks_stats_and_episode_selection():
kwargs = {
"dataset_repo_id": "org/dataset",
"base_tokenizer_name": "physical-intelligence/fast",
"n_samples": 100,
"chunk_size": 20,
"normalization_mode": "QUANTILES",
"dataset_revision": "main",
"episodes": [1, 2, 3],
"exclude_episodes": [2],
"use_relative_actions": False,
"relative_action_mask": None,
}
first = _dataset_signature(**kwargs, action_stats={"q01": [0.0], "q99": [1.0]})
changed_stats = _dataset_signature(**kwargs, action_stats={"q01": [0.0], "q99": [2.0]})
changed_selection = _dataset_signature(
**{**kwargs, "exclude_episodes": [2, 3]},
action_stats={"q01": [0.0], "q99": [1.0]},
)
assert first != changed_stats
assert first != changed_selection
def test_fast_tokenizer_uses_only_global_rank_zero(monkeypatch):
monkeypatch.setenv("RANK", "8")
monkeypatch.setenv("LOCAL_RANK", "0")
assert not _is_global_leader()
monkeypatch.setenv("RANK", "0")
assert _is_global_leader()
def test_fast_tokenizer_episode_selection_applies_allowlist_and_exclusions():
selected = _select_episode_indices([0, 1, 2, 3], episodes=[1, 2, 3], exclude_episodes=[2])
assert selected == [1, 3]
def test_fast_tokenizer_relative_actions_match_training_transform():
actions = np.array([[[2.0, 10.0], [3.0, 11.0]]], dtype=np.float32)
states = np.array([[1.0, 4.0]], dtype=np.float32)
relative = _apply_relative_actions(actions, states, [True, False])
np.testing.assert_allclose(relative, [[[1.0, 10.0], [2.0, 11.0]]])
class _RoundTripTokenizer:
def __init__(self, offset: float = 0.0):
self.offset = offset
def __call__(self, actions):
return actions
def decode(self, tokens):
return tokens + self.offset
def test_fast_tokenizer_reconstruction_validation_reports_error():
actions = np.arange(24, dtype=np.float32).reshape(2, 3, 4) / 24
report, decoded = _validate_fast_reconstruction(_RoundTripTokenizer(0.05), actions, 0.1, 0.1)
np.testing.assert_allclose(decoded, actions + 0.05)
assert report["reconstruction_rmse"] == pytest.approx(0.05)
assert report["max_dim_rmse"] == pytest.approx(0.05)
def test_fast_tokenizer_reconstruction_validation_rejects_large_error():
actions = np.arange(24, dtype=np.float32).reshape(2, 3, 4) / 24
with pytest.raises(RuntimeError, match="exceeds the configured limit"):
_validate_fast_reconstruction(_RoundTripTokenizer(0.25), actions, 0.1, 0.2)
tests/policies/pi052/test_pi052_flex_attention.py
+65
View File
@@ -0,0 +1,65 @@
#!/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 logging
import pytest
import torch
pytest.importorskip("transformers")
import lerobot.policies.pi052.modeling_pi052 as modeling_pi052 # noqa: E402
from lerobot.policies.pi052.configuration_pi052 import PI052Config # noqa: E402
def test_flex_backend_skips_non_cuda_without_initializing(monkeypatch):
monkeypatch.setattr(modeling_pi052, "_flex_fns", None)
monkeypatch.setattr(torch, "compile", lambda *args, **kwargs: pytest.fail("torch.compile was called"))
monkeypatch.setattr(
torch.cuda,
"get_device_properties",
lambda *args, **kwargs: pytest.fail("CUDA properties were queried"),
)
assert modeling_pi052._get_flex_fns(torch.device("cpu")) is None
assert modeling_pi052._get_flex_kernel_options(torch.device("cpu")) is None
assert modeling_pi052._flex_fns is None
def test_flex_initialization_failure_falls_back(monkeypatch, caplog):
monkeypatch.setattr(modeling_pi052, "_flex_fns", None)
monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
def fail_compile(*args, **kwargs):
raise RuntimeError("compile failed")
monkeypatch.setattr(torch, "compile", fail_compile)
with caplog.at_level(logging.WARNING, logger=modeling_pi052.__name__):
assert modeling_pi052._get_flex_fns(torch.device("cuda", 0)) is None
assert modeling_pi052._flex_fns is False
assert "FlexAttention unavailable" in caplog.text
def test_flex_rejects_single_repeat_configuration():
with pytest.raises(ValueError, match="use_flex_attention requires flow_num_repeats > 1"):
PI052Config(use_flex_attention=True, flow_num_repeats=1)
def test_flex_accepts_amortized_repeat_configuration():
config = PI052Config(use_flex_attention=True, flow_num_repeats=5)
assert config.use_flex_attention
tests/policies/pi052/test_pi052_import.py
+27
View File
@@ -0,0 +1,27 @@
# 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 subprocess
import sys
def test_pi052_config_import_does_not_load_model_or_dataset_processor():
code = """
import sys
from lerobot.policies import PI052Config
assert PI052Config.__name__ == "PI052Config"
assert "lerobot.policies.pi052.modeling_pi052" not in sys.modules
assert "lerobot.policies.pi052.processor_pi052" not in sys.modules
"""
subprocess.run([sys.executable, "-c", code], check=True)
tests/policies/pi052/test_pi052_joint_sequence.py
+148
View File
@@ -0,0 +1,148 @@
#!/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.
"""Tests for PI052 joint-sequence (paper-style) subtask conditioning.
Joint recipes train the subtask text and the action losses in one sequence,
with the supervised subtask span attended causally. At inference the same
layout is rebuilt around the *generated* subtask, so these tests pin:
- the inference-side encoder produces the same token ids and target positions
as the training-time tokenizer step for the same messages;
- OR-ing causal marks into a prefix reproduces the training-time attention
pattern (prompt cannot see the subtask; subtask is causal over itself);
- the joint recipe file stays a valid message recipe;
- the FAST id mapping with the default ``fast_skip_tokens`` stays clear of
PaliGemma's ``<loc>`` range so VQA and FAST supervision never collide.
"""
from pathlib import Path
import torch
from lerobot.configs.recipe import TrainingRecipe
from lerobot.policies.pi052.text_processor_pi052 import (
PI052TextTokenizerStep,
encode_prompt_with_targets,
)
class _CharTokenizer:
"""Char-level stub: 1 char = 1 token, so offsets are trivially aligned."""
pad_token_id = 0
eos_token = "\x1f" # unit separator — a 1-char "EOS" for testing
def __call__(self, text, max_length=None, padding=None, return_tensors=None, **kwargs):
limit = max_length if max_length is not None else len(text)
ids = [ord(c) % 251 + 1 for c in text[:limit]]
offsets = [(i, i + 1) for i in range(len(ids))]
attention = [1] * len(ids)
if padding == "max_length" and max_length is not None and len(ids) < max_length:
pad = max_length - len(ids)
ids += [self.pad_token_id] * pad
offsets += [(0, 0)] * pad
attention += [0] * pad
return {
"input_ids": torch.tensor([ids], dtype=torch.long),
"attention_mask": torch.tensor([attention], dtype=torch.long),
"offset_mapping": torch.tensor([offsets], dtype=torch.long),
}
_MESSAGES = [
{"role": "user", "content": "fold the towel"},
{"role": "assistant", "content": "grab the near corner"},
]
def test_encode_prompt_with_targets_matches_training_labels():
tokenizer = _CharTokenizer()
step = PI052TextTokenizerStep(max_length=120)
step._tokenizer = tokenizer
train_ids, train_attn, labels, predict_actions, _prompt = step._encode_messages(
tokenizer,
[dict(m) for m in _MESSAGES],
message_streams=["low_level", "low_level"],
target_indices=[1],
complementary={},
)
assert bool(predict_actions)
ids, attn, marks = encode_prompt_with_targets(tokenizer, [dict(m) for m in _MESSAGES], [1])
n = int(attn.sum())
assert n == int(train_attn.sum())
assert torch.equal(ids[0, :n], train_ids[:n])
# Causal marks at inference must cover exactly the supervised label span.
assert torch.equal(marks[0, :n], labels[:n] != -100)
assert marks.any(), "the assistant target span must be marked"
# The user turn must stay unmarked (bidirectional prompt).
user_len = len("User: fold the towel\n")
assert not marks[0, :user_len].any()
def test_apply_causal_language_marks_reproduces_training_mask():
from lerobot.policies.pi05.modeling_pi05 import make_att_2d_masks
from lerobot.policies.pi052.modeling_pi052 import _apply_causal_language_marks
n_img, n_lang = 4, 8
prefix_len = n_img + n_lang
pad = torch.ones((1, prefix_len), dtype=torch.bool)
att = torch.zeros((1, prefix_len), dtype=torch.bool)
# Subtask span = language positions 5..7 (prefix positions 9..11).
marks = torch.zeros((1, n_lang), dtype=torch.bool)
marks[0, 5:8] = True
att_marked = _apply_causal_language_marks(att, marks)
att_2d = make_att_2d_masks(pad, att_marked)[0]
subtask = [n_img + 5, n_img + 6, n_img + 7]
# Prompt and images never see the subtask.
for q in range(n_img + 5):
for k in subtask:
assert not att_2d[q, k], f"prompt position {q} must not attend subtask position {k}"
# Subtask tokens see the full prompt and earlier subtask tokens only.
for qi, q in enumerate(subtask):
for k in range(n_img + 5):
assert att_2d[q, k]
for ki, k in enumerate(subtask):
assert bool(att_2d[q, k]) == (ki <= qi)
def test_joint_recipe_is_a_valid_message_recipe():
recipe_path = Path(__file__).parents[3] / "src" / "lerobot" / "configs" / "recipes" / "subtask_joint.yaml"
recipe = TrainingRecipe.from_yaml(recipe_path)
assert recipe.messages is not None and len(recipe.messages) == 2
assert all(turn.stream == "low_level" for turn in recipe.messages)
assert not recipe.messages[0].target
assert recipe.messages[1].target
assert recipe.messages[1].if_present == "subtask"
def test_default_fast_mapping_clears_loc_and_seg_ranges():
from lerobot.policies.pi052.configuration_pi052 import PI052Config
from lerobot.policies.pi052.modeling_pi052 import _FAST_ACTION_VOCAB_SIZE
skip = PI052Config.__dataclass_fields__["fast_skip_tokens"].default
assert skip == 1152
paligemma_vocab = 257152
fast_ids = paligemma_vocab - 1 - skip - torch.arange(_FAST_ACTION_VOCAB_SIZE)
# Below the <loc> range [256000, 257024) and the <seg> range [257024, 257152).
assert int(fast_ids.max()) < 256000
assert int(fast_ids.min()) >= 0
tests/policies/pi052/test_pi052_runtime_adapter.py
+85
View File
@@ -0,0 +1,85 @@
# 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 types import SimpleNamespace
from lerobot.policies.pi052.inference.pi052_adapter import PI052PolicyAdapter
from lerobot.runtime import RuntimeState
from lerobot.runtime.adapter import split_plan_and_say
def test_pi052_adapter_builds_recipe_prompts_from_runtime_state():
adapter = PI052PolicyAdapter(policy=object())
state = RuntimeState(
task="clean the kitchen",
language_context={"memory": "cup moved", "plan": "pick then place"},
extra={"prior_subtask": "pick the cup"},
)
assert adapter.build_messages("subtask", state) == [{"role": "user", "content": "clean the kitchen"}]
assert adapter.build_messages("memory", state) == [
{"role": "user", "content": "clean the kitchen"},
{"role": "assistant", "content": "Previous memory: cup moved"},
{"role": "user", "content": "Completed subtask: pick the cup"},
]
assert adapter.build_messages("interjection", state, user_text="wait") == [
{"role": "user", "content": "clean the kitchen"},
{"role": "assistant", "content": "Previous plan:\npick then place"},
{"role": "user", "content": "wait"},
]
def test_pi052_adapter_strips_say_markers_from_plan_text():
adapter = PI052PolicyAdapter(policy=object())
text = "Move to the sink. <say>heading to the sink</say>"
assert split_plan_and_say(text) == ("Move to the sink.", "heading to the sink")
assert adapter.plan_from_text(text) == "Move to the sink."
def test_rollout_language_cli_smoke_does_not_load_model(monkeypatch):
"""lerobot-rollout dispatches language flags to the adapter-based runtime."""
from lerobot.runtime import cli
from lerobot.scripts import lerobot_rollout
fake_policy = SimpleNamespace(config=SimpleNamespace(device="cpu", type="pi052"))
monkeypatch.setattr(
cli,
"_load_policy_and_preprocessor",
lambda policy_path, **kwargs: (fake_policy, None, None),
)
monkeypatch.setattr(cli, "_run_repl", lambda runtime, **kwargs: 0)
assert lerobot_rollout.main(["--policy.path=fake", "--no_robot", "--task=clean", "--max_ticks=0"]) == 0
def test_rollout_language_dispatch_preserves_standard_molmoact2_path(monkeypatch):
"""MolmoAct2 only opts into open prompting when a language flag is present."""
from lerobot.scripts import lerobot_rollout
standard = [
"--policy.path=lerobot/MolmoAct2-SO100_101-LeRobot",
"--robot.type=so101_follower",
"--task=pick up the cube",
]
assert not lerobot_rollout._uses_language_runtime(standard)
assert lerobot_rollout._uses_language_runtime([*standard, "--direct_subtask"])
assert lerobot_rollout._uses_language_runtime(["--policy.path=lerobot/pi052_robocasa", "--sim"])
standard_calls = []
monkeypatch.setattr(lerobot_rollout, "register_third_party_plugins", lambda: None)
monkeypatch.setattr(lerobot_rollout, "rollout", lambda: standard_calls.append(True))
lerobot_rollout.main(standard)
assert standard_calls == [True]
tests/policies/pi052/test_pi052_sdpa_attention.py
+147
View File
@@ -0,0 +1,147 @@
#!/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.
"""Numerical-parity tests for the SDPA attention port.
``pi05`` / ``pi052`` replaced the per-layer call from
``modeling_gemma.eager_attention_forward`` with
``sdpa_attention_forward`` (PyTorch SDPA + GQA repeat). The forward
output must be bit-equivalent (within bf16 tolerance) on the masks
this model actually uses — block-bidirectional with an arbitrary
additive bias — otherwise we silently change training behaviour.
"""
from types import SimpleNamespace
import pytest
import torch
pytest.importorskip("transformers")
from transformers.models.gemma import modeling_gemma # noqa: E402
from lerobot.policies.pi052.modeling_pi052 import make_att_2d_masks # noqa: E402
from lerobot.policies.pi_gemma import sdpa_attention_forward # noqa: E402
from lerobot.utils.constants import OPENPI_ATTENTION_MASK_VALUE # noqa: E402
def _mock_self_attn(num_kv_groups: int, training: bool = False):
"""Bare module surface that both forwards read."""
return SimpleNamespace(
num_key_value_groups=num_kv_groups,
training=training,
)
def _build_inputs(
bsize: int,
num_heads: int,
num_kv_heads: int,
seq_len: int,
head_dim: int,
dtype: torch.dtype,
seed: int = 0,
):
g = torch.Generator(device="cpu").manual_seed(seed)
q = torch.randn(bsize, num_heads, seq_len, head_dim, dtype=dtype, generator=g)
k = torch.randn(bsize, num_kv_heads, seq_len, head_dim, dtype=dtype, generator=g)
v = torch.randn(bsize, num_kv_heads, seq_len, head_dim, dtype=dtype, generator=g)
return q, k, v
def _block_bidirectional_mask(
bsize: int, seq_len: int, block_sizes: list[int], dtype: torch.dtype
) -> torch.Tensor:
"""Mimic ``_prepare_attention_masks_4d`` on a block layout that
matches ``[images, language, suffix]`` from ``embed_prefix`` +
``embed_suffix``: every block bidirectional internally, later
blocks visible to earlier ones via the cumulative-block rule.
"""
assert sum(block_sizes) == seq_len
att_marks = []
for i, n in enumerate(block_sizes):
att_marks += [1 if i > 0 else 0] + [0] * (n - 1)
pad = torch.ones(bsize, seq_len, dtype=torch.bool)
att = torch.tensor(att_marks, dtype=torch.bool)[None].expand(bsize, seq_len)
att_2d = make_att_2d_masks(pad, att)
bias = torch.where(
att_2d[:, None, :, :],
torch.zeros((), dtype=dtype),
torch.tensor(OPENPI_ATTENTION_MASK_VALUE, dtype=dtype),
)
return bias
@pytest.mark.parametrize(
"num_heads,num_kv_heads,head_dim",
[
(8, 1, 256), # gemma_2b / paligemma config
(8, 8, 64), # MHA control (no GQA repeat)
],
)
def test_sdpa_parity_with_eager_block_bidirectional(num_heads, num_kv_heads, head_dim):
"""SDPA forward output matches the eager softmax(QK^T)@V on the
block-bidirectional mask layout pi05 actually uses."""
bsize, seq_len = 2, 13
block_sizes = [4, 5, 4] # images, language, suffix-style blocks
dtype = torch.float32 # cpu math kernel — keep fp32 for tight tol
scaling = head_dim**-0.5
q, k, v = _build_inputs(bsize, num_heads, num_kv_heads, seq_len, head_dim, dtype)
mask = _block_bidirectional_mask(bsize, seq_len, block_sizes, dtype)
module = _mock_self_attn(num_heads // num_kv_heads)
out_eager, _ = modeling_gemma.eager_attention_forward(module, q, k, v, mask, scaling)
out_sdpa, _ = sdpa_attention_forward(module, q, k, v, mask, scaling)
assert out_eager.shape == out_sdpa.shape
torch.testing.assert_close(out_sdpa, out_eager, atol=1e-5, rtol=1e-4)
def test_sdpa_parity_bf16():
"""bf16 path — looser tolerance, must still match eager."""
bsize, num_heads, num_kv_heads, seq_len, head_dim = 2, 8, 1, 17, 256
scaling = head_dim**-0.5
q, k, v = _build_inputs(bsize, num_heads, num_kv_heads, seq_len, head_dim, torch.bfloat16)
mask = _block_bidirectional_mask(bsize, seq_len, [5, 6, 6], torch.bfloat16)
module = _mock_self_attn(num_heads // num_kv_heads)
out_eager, _ = modeling_gemma.eager_attention_forward(module, q, k, v, mask, scaling)
out_sdpa, _ = sdpa_attention_forward(module, q, k, v, mask, scaling)
torch.testing.assert_close(out_sdpa, out_eager, atol=2e-2, rtol=2e-2)
def test_sdpa_parity_backward():
"""Gradients flow through SDPA and match the eager path within
bf16 tolerance — critical for any training-side parity claim."""
bsize, num_heads, num_kv_heads, seq_len, head_dim = 1, 4, 2, 9, 32
scaling = head_dim**-0.5
q, k, v = _build_inputs(bsize, num_heads, num_kv_heads, seq_len, head_dim, torch.float32)
q.requires_grad_(True)
k.requires_grad_(True)
v.requires_grad_(True)
mask = _block_bidirectional_mask(bsize, seq_len, [3, 3, 3], torch.float32)
module = _mock_self_attn(num_heads // num_kv_heads)
out_e, _ = modeling_gemma.eager_attention_forward(module, q, k, v, mask, scaling)
g_q_e, g_k_e, g_v_e = torch.autograd.grad(out_e.sum(), [q, k, v])
out_s, _ = sdpa_attention_forward(module, q, k, v, mask, scaling)
g_q_s, g_k_s, g_v_s = torch.autograd.grad(out_s.sum(), [q, k, v])
torch.testing.assert_close(g_q_s, g_q_e, atol=1e-5, rtol=1e-4)
torch.testing.assert_close(g_k_s, g_k_e, atol=1e-5, rtol=1e-4)
torch.testing.assert_close(g_v_s, g_v_e, atol=1e-5, rtol=1e-4)
tests/policies/pi052/test_pi052_text_processor.py
+223
View File
@@ -0,0 +1,223 @@
#!/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.
"""Tests for PI052's text tokenizer.
Covers ``say`` tool-call flattening (PaliGemma's flat prompt has no
structured tool calls, so a ``say`` call must be serialized into a
``<say>...</say>`` text marker) and EOS-termination supervision (the
supervised target span must end with an EOS token so the LM head learns
to stop instead of rambling to ``max_length`` at inference).
"""
import torch
from lerobot.configs.recipe import MessageTurn, TrainingRecipe
from lerobot.policies.pi052.text_processor_pi052 import (
PI052TextTokenizerStep,
_flatten_say_tool_calls,
_format_messages,
)
from lerobot.processor import PolicyProcessorPipeline
from lerobot.processor.render_messages_processor import RenderMessagesStep
from lerobot.types import TransitionKey
from lerobot.utils.constants import (
OBS_LANGUAGE_ATTENTION_MASK,
OBS_LANGUAGE_TOKENS,
POLICY_PREPROCESSOR_DEFAULT_NAME,
)
def _say_call(text):
return {"type": "function", "function": {"name": "say", "arguments": {"text": text}}}
def test_flatten_appends_say_marker_and_drops_tool_calls():
msg = {"role": "assistant", "content": "Heading to the cube.", "tool_calls": [_say_call("On it!")]}
out = _flatten_say_tool_calls(msg)
assert "tool_calls" not in out
assert out["content"] == "Heading to the cube.\n<say>On it!</say>"
def test_flatten_marker_only_when_content_empty_or_none():
out = _flatten_say_tool_calls({"role": "assistant", "tool_calls": [_say_call("hi")]})
assert out["content"] == "<say>hi</say>"
def test_flatten_accepts_json_string_arguments():
call = {"type": "function", "function": {"name": "say", "arguments": '{"text": "hello there"}'}}
out = _flatten_say_tool_calls({"role": "assistant", "content": "p", "tool_calls": [call]})
assert out["content"] == "p\n<say>hello there</say>"
def test_flatten_leaves_messages_without_tool_calls_untouched():
msg = {"role": "assistant", "content": "just a plan"}
assert _flatten_say_tool_calls(msg) == msg
def test_flatten_drops_non_say_tool_calls_but_keeps_content():
weather = {"type": "function", "function": {"name": "check_weather", "arguments": {}}}
out = _flatten_say_tool_calls({"role": "assistant", "content": "plan only", "tool_calls": [weather]})
assert out["content"] == "plan only"
assert "tool_calls" not in out
def test_format_messages_appends_eos_to_target_turns_only():
msgs = [
{"role": "user", "content": "pick cube"},
{"role": "assistant", "content": "move to cube"},
]
prompt, spans = _format_messages(msgs, target_indices=[1], eos_token="<eos>")
# EOS is appended to the supervised target (assistant) turn only.
assert prompt == "User: pick cube\nAssistant: move to cube<eos>\n"
# The user span is unchanged; the target span covers content + EOS.
assert prompt[spans[0][0] : spans[0][1]] == "pick cube"
assert prompt[spans[1][0] : spans[1][1]] == "move to cube<eos>"
def test_format_messages_without_eos_args_is_unchanged():
"""Inference callers omit target_indices / eos_token — no EOS baked in."""
prompt, spans = _format_messages([{"role": "user", "content": "hi"}])
assert prompt == "User: hi\n"
assert prompt[spans[0][0] : spans[0][1]] == "hi"
def test_pi052_steps_roundtrip_through_standard_pipeline_loader(tmp_path):
recipe = TrainingRecipe(messages=[MessageTurn(role="user", content="${task}", stream="low_level")])
pipeline = PolicyProcessorPipeline(
steps=[
RenderMessagesStep(recipe),
PI052TextTokenizerStep(
tokenizer_name="custom-tokenizer",
max_length=77,
plan_dropout_prob=0.2,
dropout_seed=3,
),
],
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
)
pipeline.save_pretrained(tmp_path)
loaded = PolicyProcessorPipeline.from_pretrained(
tmp_path, config_filename=f"{POLICY_PREPROCESSOR_DEFAULT_NAME}.json"
)
assert loaded.steps[0].recipe == recipe
assert loaded.steps[1].tokenizer_name == "custom-tokenizer"
assert loaded.steps[1].max_length == 77
assert loaded.steps[1].plan_dropout_prob == 0.2
assert loaded.steps[1].dropout_seed == 3
def _eos_char_id() -> int:
"""Token id _CharTokenizer assigns to its 1-char EOS."""
return ord("\x1f") % 251 + 1
def test_pi052_text_tokenizer_supervises_eos_at_target_end():
"""The appended EOS is the last supervised label on a target turn —
that's the signal that teaches the LM head to stop. The trailing
newline right after it stays unsupervised (-100)."""
step = PI052TextTokenizerStep(max_length=64)
step._tokenizer = _CharTokenizer()
transition = {
TransitionKey.OBSERVATION: {},
TransitionKey.COMPLEMENTARY_DATA: {
"messages": [
{"role": "user", "content": "pick cube"},
{"role": "assistant", "content": "move to cube"},
],
"target_message_indices": [1],
"message_streams": ["high_level", "high_level"],
"index": torch.tensor(10),
},
}
out = step(transition)
ids = out[TransitionKey.OBSERVATION][OBS_LANGUAGE_TOKENS][0]
labels = out[TransitionKey.COMPLEMENTARY_DATA]["text_labels"][0]
supervised = (labels != -100).nonzero().flatten().tolist()
assert supervised, "target turn produced no supervised labels"
last = supervised[-1]
# The last supervised token is the appended EOS.
assert int(ids[last]) == _eos_char_id()
assert int(labels[last]) == _eos_char_id()
# The token right after the EOS (the trailing newline) is NOT supervised.
assert int(labels[last + 1]) == -100
class _CharTokenizer:
pad_token_id = 0
eos_token = "\x1f" # unit separator — a 1-char "EOS" for testing
def __call__(
self,
text,
max_length,
padding,
truncation,
return_tensors,
return_offsets_mapping,
padding_side,
):
ids = [ord(c) % 251 + 1 for c in text[:max_length]]
offsets = [(i, i + 1) for i in range(len(ids))]
attention = [1] * len(ids)
if padding == "max_length" and len(ids) < max_length:
pad = max_length - len(ids)
ids += [self.pad_token_id] * pad
offsets += [(0, 0)] * pad
attention += [0] * pad
return {
"input_ids": torch.tensor([ids], dtype=torch.long),
"attention_mask": torch.tensor([attention], dtype=torch.long),
"offset_mapping": torch.tensor([offsets], dtype=torch.long),
}
def decode(self, token_ids, skip_special_tokens=False):
return "".join(chr(max(int(i) - 1, 0)) for i in token_ids if int(i) != self.pad_token_id)
def test_pi052_text_tokenizer_handles_batched_rendered_messages():
step = PI052TextTokenizerStep(max_length=64)
step._tokenizer = _CharTokenizer()
transition = {
TransitionKey.OBSERVATION: {},
TransitionKey.COMPLEMENTARY_DATA: {
"messages": [
[
{"role": "user", "content": "pick cube"},
{"role": "assistant", "content": "move to cube"},
],
[{"role": "user", "content": "open drawer"}],
],
"target_message_indices": [[1], []],
"message_streams": [["high_level", "high_level"], ["low_level"]],
"index": torch.tensor([10, 11]),
},
}
out = step(transition)
obs = out[TransitionKey.OBSERVATION]
comp = out[TransitionKey.COMPLEMENTARY_DATA]
assert obs[OBS_LANGUAGE_TOKENS].shape == (2, 64)
assert obs[OBS_LANGUAGE_ATTENTION_MASK].shape == (2, 64)
assert comp["text_labels"].shape == (2, 64)
assert comp["predict_actions"].tolist() == [False, True]
assert (comp["text_labels"][0] != -100).any()
assert not (comp["text_labels"][1] != -100).any()
tests/policies/pi052/test_pi052_vision_checkpointing.py
+141
View File
@@ -0,0 +1,141 @@
#!/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 types import MethodType, SimpleNamespace
import pytest
import torch
from torch import nn
pytest.importorskip("transformers")
from lerobot.policies.pi052.modeling_pi052 import PI05Pytorch
class _MockVisionTower:
def __init__(self):
self.enable_kwargs = None
self.disable_calls = 0
def gradient_checkpointing_enable(self, **kwargs):
self.enable_kwargs = kwargs
def gradient_checkpointing_disable(self):
self.disable_calls += 1
def _checkpoint_model():
tower = _MockVisionTower()
language_model = SimpleNamespace(gradient_checkpointing=False)
expert_model = SimpleNamespace(gradient_checkpointing=False)
model = PI05Pytorch.__new__(PI05Pytorch)
nn.Module.__init__(model)
model.gradient_checkpointing_enabled = False
model.paligemma_with_expert = SimpleNamespace(
paligemma=SimpleNamespace(model=SimpleNamespace(language_model=language_model, vision_tower=tower)),
gemma_expert=SimpleNamespace(model=expert_model),
)
return model, tower, language_model, expert_model
def test_gradient_checkpointing_uses_vision_tower_layer_api():
model, tower, language_model, expert_model = _checkpoint_model()
PI05Pytorch.gradient_checkpointing_enable(model)
assert model.gradient_checkpointing_enabled
assert language_model.gradient_checkpointing
assert expert_model.gradient_checkpointing
assert tower.enable_kwargs == {"gradient_checkpointing_kwargs": {"use_reentrant": False}}
PI05Pytorch.gradient_checkpointing_disable(model)
assert not model.gradient_checkpointing_enabled
assert not language_model.gradient_checkpointing
assert not expert_model.gradient_checkpointing
assert tower.disable_calls == 1
def test_siglip_layers_recompute_individually():
from transformers.models.siglip.configuration_siglip import SiglipVisionConfig
from transformers.models.siglip.modeling_siglip import SiglipVisionModel
config = SiglipVisionConfig(
hidden_size=16,
intermediate_size=32,
num_hidden_layers=2,
num_attention_heads=2,
num_channels=3,
image_size=16,
patch_size=8,
)
tower = SiglipVisionModel(config).train()
tower.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
calls = [0] * config.num_hidden_layers
for index, layer in enumerate(tower.vision_model.encoder.layers):
original_forward = layer.forward
def counted_forward(self, *args, _index=index, _forward=original_forward, **kwargs):
calls[_index] += 1
return _forward(*args, **kwargs)
layer.forward = MethodType(counted_forward, layer)
pixels = torch.randn(2, config.num_channels, config.image_size, config.image_size)
tower(pixels).last_hidden_state.sum().backward()
assert calls == [2] * config.num_hidden_layers
def test_embed_prefix_does_not_wrap_the_whole_vision_tower_checkpoint():
model = PI05Pytorch.__new__(PI05Pytorch)
nn.Module.__init__(model)
model.config = SimpleNamespace()
model.gradient_checkpointing_enabled = True
model.train()
image_calls = []
def embed_image(image):
image_calls.append(image.shape)
return image[:, :1, 0, :2]
def embed_language_tokens(tokens):
return tokens.to(torch.float32).unsqueeze(-1).expand(*tokens.shape, 2)
model.paligemma_with_expert = SimpleNamespace(
embed_image=embed_image,
embed_language_tokens=embed_language_tokens,
)
outer_checkpoint_calls = []
def apply_checkpoint(func, value):
outer_checkpoint_calls.append(value.shape)
return func(value)
model._apply_checkpoint = apply_checkpoint
images = [torch.randn(2, 3, 4, 4), torch.randn(2, 3, 4, 4)]
image_masks = [torch.ones(2, dtype=torch.bool) for _ in images]
tokens = torch.ones(2, 3, dtype=torch.long)
token_masks = torch.ones_like(tokens, dtype=torch.bool)
embeddings, _, _ = model.embed_prefix(images, image_masks, tokens, token_masks)
assert image_calls == [image.shape for image in images]
assert outer_checkpoint_calls == [tokens.shape]
assert embeddings.shape == (2, 5, 2)
tests/policies/pi0_fast/test_pi0_fast_tokenizer_fit.py
+106
View File
@@ -0,0 +1,106 @@
#!/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 types import SimpleNamespace
import pytest
from lerobot.policies import factory
from lerobot.policies.pi0_fast.configuration_pi0_fast import PI0FastConfig
from lerobot.policies.pi052 import fit_fast_tokenizer as fit_module
def test_pi0_fast_resolves_dataset_specific_tokenizer(monkeypatch, tmp_path):
config = PI0FastConfig(
auto_fit_fast_tokenizer=True,
action_tokenizer_name="base-tokenizer",
fast_tokenizer_cache_dir=str(tmp_path),
fast_tokenizer_fit_samples=17,
chunk_size=12,
n_action_steps=12,
)
received = {}
def fake_fit(**kwargs):
received.update(kwargs)
return "/cache/fitted-tokenizer"
monkeypatch.setattr(fit_module, "fit_fast_tokenizer", fake_fit)
assert fit_module.resolve_fast_tokenizer(config, "user/dataset") == "/cache/fitted-tokenizer"
assert received == {
"dataset_repo_id": "user/dataset",
"cache_dir": tmp_path,
"base_tokenizer_name": "base-tokenizer",
"n_samples": 17,
"chunk_size": 12,
"dataset_root": None,
"dataset_revision": None,
"episodes": None,
"exclude_episodes": None,
"normalization_mode": config.normalization_mapping["ACTION"],
"action_stats": None,
"use_relative_actions": False,
"relative_action_mask": None,
}
def test_fast_fit_failure_is_not_silently_replaced(monkeypatch, tmp_path):
config = PI0FastConfig(auto_fit_fast_tokenizer=True, fast_tokenizer_cache_dir=str(tmp_path))
monkeypatch.setattr(
fit_module,
"fit_fast_tokenizer",
lambda **kwargs: (_ for _ in ()).throw(RuntimeError("fit failed")),
)
with pytest.raises(RuntimeError, match="fit failed"):
fit_module.resolve_fast_tokenizer(config, "user/dataset")
def test_only_global_rank_zero_fits_shared_tokenizer(monkeypatch):
monkeypatch.setenv("RANK", "8")
monkeypatch.setenv("LOCAL_RANK", "0")
assert not fit_module._is_global_leader()
monkeypatch.setenv("RANK", "0")
assert fit_module._is_global_leader()
def test_pretrained_pi0_fast_overrides_only_fitted_tokenizer(monkeypatch):
config = PI0FastConfig(auto_fit_fast_tokenizer=True)
calls = []
monkeypatch.setattr(
fit_module,
"resolve_fast_tokenizer",
lambda config, dataset_repo_id, *args: "/cache/fitted-tokenizer",
)
def fake_from_pretrained(cls, *args, **kwargs):
calls.append(kwargs)
return SimpleNamespace(steps=[])
monkeypatch.setattr(factory.PolicyProcessorPipeline, "from_pretrained", classmethod(fake_from_pretrained))
factory.make_pre_post_processors(
config,
pretrained_path="checkpoint",
dataset_repo_id="user/dataset",
)
assert calls[0]["overrides"] == {
"action_tokenizer_processor": {"action_tokenizer_name": "/cache/fitted-tokenizer"}
}
tests/policies/pi0_pi05/test_pi05.py
+24
View File
@@ -16,8 +16,12 @@
"""Test script to verify PI0.5 (pi05) support in PI0 policy"""
from types import SimpleNamespace
import pytest
import torch
from safetensors.torch import save_file
from torch import nn
pytest.importorskip("transformers")
@@ -31,6 +35,26 @@ from lerobot.utils.random_utils import set_seed
from tests.utils import require_cuda, require_hf_token # noqa: E402
class _CheckpointPolicy(PI05Policy):
def __init__(self, config, **kwargs):
nn.Module.__init__(self)
self.config = config
self.loaded_state_dict = None
def load_state_dict(self, state_dict, strict=True, assign=False):
self.loaded_state_dict = state_dict
return [], []
def test_from_pretrained_loads_existing_single_file_checkpoint(tmp_path):
save_file({"weight": torch.tensor([1.0])}, tmp_path / "model.safetensors")
policy = _CheckpointPolicy.from_pretrained(tmp_path, config=SimpleNamespace())
assert policy.loaded_state_dict is not None
torch.testing.assert_close(policy.loaded_state_dict["model.weight"], torch.tensor([1.0]))
@require_cuda
@require_hf_token
def test_policy_instantiation():
tests/scripts/test_backfill_pi052_action_tokenizer.py
+67
View File
@@ -0,0 +1,67 @@
#!/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 scripts.backfill_pi052_action_tokenizer import (
CHECKPOINT_DIRECTORIES,
DEFAULT_REPOSITORIES,
artifact_fingerprint,
make_portable_preprocessor,
)
def test_atomic4_backfill_covers_every_repository_and_checkpoint():
assert len(DEFAULT_REPOSITORIES) == 6
assert CHECKPOINT_DIRECTORIES == (
"",
"checkpoints/003000/pretrained_model",
"checkpoints/006000/pretrained_model",
"checkpoints/009000/pretrained_model",
"checkpoints/012000/pretrained_model",
)
def test_backfill_embeds_recipe_and_declares_relative_tokenizer():
recipe = {"messages": [{"role": "user", "content": "${task}", "stream": "low_level"}]}
preprocessor = {
"name": "policy_preprocessor",
"steps": [
{
"registry_name": "normalizer_processor",
"config": {},
"state_file": "normalizer.safetensors",
},
{"registry_name": "render_messages_processor", "config": {"recipe": recipe}},
{
"registry_name": "action_tokenizer_processor",
"config": {"action_tokenizer_name": "/fsx/original/tokenizer"},
},
],
}
portable = make_portable_preprocessor(preprocessor)
assert portable["steps"][1]["config"]["recipe"] == recipe
assert portable["steps"][2]["config"]["action_tokenizer_name"] == "action_tokenizer"
assert portable["steps"][2]["artifacts"] == {"action_tokenizer_name": "action_tokenizer"}
assert preprocessor["steps"][2]["config"]["action_tokenizer_name"].startswith("/fsx/")
def test_artifact_fingerprint_includes_paths_and_contents():
first = artifact_fingerprint([("a/file", b"same"), ("b/file", b"content")])
assert first == artifact_fingerprint([("b/file", b"content"), ("a/file", b"same")])
assert first != artifact_fingerprint([("a/renamed", b"same"), ("b/file", b"content")])
assert first != artifact_fingerprint([("a/file", b"changed"), ("b/file", b"content")])
tests/utils/test_logging_utils.py
+18
View File
@@ -37,6 +37,12 @@ class MockAccelerator:
return self._reduce_fn(tensor, reduction)
return tensor
def gather(self, tensor):
if self._reduce_fn is None:
return tensor.repeat(self.num_processes)
reduced = self._reduce_fn(tensor, "max")
return torch.cat([tensor.repeat(self.num_processes - 1), reduced])
def test_average_meter_initialization():
meter = AverageMeter("loss", ":.2f")
@@ -168,6 +174,18 @@ def test_metrics_tracker_reset_averages(mock_metrics):
assert tracker.accuracy.avg == 0.0
def test_metrics_tracker_materializes_full_tensor_window(mock_metrics):
tracker = MetricsTracker(batch_size=2, num_frames=10, num_episodes=2, metrics=mock_metrics)
tracker.accumulate_tensor("loss", torch.tensor(1.0))
tracker.accumulate_tensor("loss", torch.tensor(3.0))
assert tracker.loss.count == 0
tracker.materialize_tensors()
assert tracker.loss.avg == pytest.approx(2.0)
assert tracker.loss.count == 2
def test_average_meter_invalid_reduction():
with pytest.raises(ValueError):
AverageMeter("loss", reduction="median")
uv.lock
Generated
+184 -134
View File
@@ -402,10 +402,10 @@ name = "bddl"
version = "1.0.1"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "jupytext" },
{ name = "networkx" },
{ name = "numpy" },
{ name = "pytest" },
{ name = "jupytext", marker = "sys_platform == 'linux'" },
{ name = "networkx", marker = "sys_platform == 'linux'" },
{ name = "numpy", marker = "sys_platform == 'linux'" },
{ name = "pytest", marker = "sys_platform == 'linux'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/5c/37/0211f82891a9f14efcfd2b2096f8d9e4351398ad637fdd1ee59cfc580b0e/bddl-1.0.1.tar.gz", hash = "sha256:1fa4e6e5050b93888ff6fd8455c39bfb29d3864ce06b4c37c0f781f513a2ae26", size = 164809, upload-time = "2022-03-08T01:48:23.564Z" }
@@ -1010,7 +1010,7 @@ name = "cuda-bindings"
version = "12.9.7"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "cuda-pathfinder" },
{ name = "cuda-pathfinder", marker = "sys_platform == 'linux'" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/32/45/557d4ed1fa54f0c7db8aee083229f624990d69f7d00f55477eed5c7e169a/cuda_bindings-12.9.7-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0666d3c082ef8f4b2d670950589373550e9f3bf564d635dd883f24a0b40402ff", size = 7071026, upload-time = "2026-05-27T18:44:13.356Z" },
@@ -1043,37 +1043,37 @@ wheels = [
[package.optional-dependencies]
cublas = [
{ name = "nvidia-cublas-cu12" },
{ name = "nvidia-cublas-cu12", marker = "sys_platform == 'linux'" },
]
cudart = [
{ name = "nvidia-cuda-runtime-cu12" },
{ name = "nvidia-cuda-runtime-cu12", marker = "sys_platform == 'linux'" },
]
cufft = [
{ name = "nvidia-cufft-cu12" },
{ name = "nvidia-cufft-cu12", marker = "sys_platform == 'linux'" },
]
cufile = [
{ name = "nvidia-cufile-cu12" },
{ name = "nvidia-cufile-cu12", marker = "sys_platform == 'linux'" },
]
cupti = [
{ name = "nvidia-cuda-cupti-cu12" },
{ name = "nvidia-cuda-cupti-cu12", marker = "sys_platform == 'linux'" },
]
curand = [
{ name = "nvidia-curand-cu12" },
{ name = "nvidia-curand-cu12", marker = "sys_platform == 'linux'" },
]
cusolver = [
{ name = "nvidia-cusolver-cu12" },
{ name = "nvidia-cusolver-cu12", marker = "sys_platform == 'linux'" },
]
cusparse = [
{ name = "nvidia-cusparse-cu12" },
{ name = "nvidia-cusparse-cu12", marker = "sys_platform == 'linux'" },
]
nvjitlink = [
{ name = "nvidia-nvjitlink-cu12" },
{ name = "nvidia-nvjitlink-cu12", marker = "sys_platform == 'linux'" },
]
nvrtc = [
{ name = "nvidia-cuda-nvrtc-cu12" },
{ name = "nvidia-cuda-nvrtc-cu12", marker = "sys_platform == 'linux'" },
]
nvtx = [
{ name = "nvidia-nvtx-cu12" },
{ name = "nvidia-nvtx-cu12", marker = "sys_platform == 'linux'" },
]
[[package]]
@@ -1145,7 +1145,7 @@ name = "decord"
version = "0.6.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "numpy" },
{ name = "numpy", marker = "(platform_machine != 'arm64' and sys_platform == 'darwin') or (platform_machine == 'AMD64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux')" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/11/79/936af42edf90a7bd4e41a6cac89c913d4b47fa48a26b042d5129a9242ee3/decord-0.6.0-py3-none-manylinux2010_x86_64.whl", hash = "sha256:51997f20be8958e23b7c4061ba45d0efcd86bffd5fe81c695d0befee0d442976", size = 13602299, upload-time = "2021-06-14T21:30:55.486Z" },
@@ -1283,10 +1283,10 @@ resolution-markers = [
"python_full_version == '3.14.*' and sys_platform == 'win32'",
]
dependencies = [
{ name = "absl-py" },
{ name = "attrs" },
{ name = "numpy" },
{ name = "wrapt" },
{ name = "absl-py", marker = "python_full_version >= '3.14'" },
{ name = "attrs", marker = "python_full_version >= '3.14'" },
{ name = "numpy", marker = "python_full_version >= '3.14'" },
{ name = "wrapt", marker = "python_full_version >= '3.14'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/a6/83/ce29720ccf934c6cfa9b9c95ebbe96558386e66886626066632b5e44afed/dm_tree-0.1.9.tar.gz", hash = "sha256:a4c7db3d3935a5a2d5e4b383fc26c6b0cd6f78c6d4605d3e7b518800ecd5342b", size = 35623, upload-time = "2025-01-30T20:45:37.13Z" }
wheels = [
@@ -1324,10 +1324,10 @@ resolution-markers = [
"python_full_version < '3.13' and sys_platform == 'win32'",
]
dependencies = [
{ name = "absl-py" },
{ name = "attrs" },
{ name = "numpy" },
{ name = "wrapt" },
{ name = "absl-py", marker = "python_full_version < '3.14'" },
{ name = "attrs", marker = "python_full_version < '3.14'" },
{ name = "numpy", marker = "python_full_version < '3.14'" },
{ name = "wrapt", marker = "python_full_version < '3.14'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/5a/66/a3ec619d22b6baffa5ab853e8dc6ec9d0c837127948af59bb15b988d7312/dm_tree-0.1.10.tar.gz", hash = "sha256:22f37b599e01cc3402a17f79c257a802aebd8d326de05b54657650845956208a", size = 35748, upload-time = "2026-03-31T17:35:39.03Z" }
wheels = [
@@ -1912,7 +1912,7 @@ name = "h5py"
version = "3.16.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "numpy" },
{ name = "numpy", marker = "sys_platform == 'linux'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/db/33/acd0ce6863b6c0d7735007df01815403f5589a21ff8c2e1ee2587a38f548/h5py-3.16.0.tar.gz", hash = "sha256:a0dbaad796840ccaa67a4c144a0d0c8080073c34c76d5a6941d6818678ef2738", size = 446526, upload-time = "2026-03-06T13:49:08.07Z" }
wheels = [
@@ -1956,23 +1956,23 @@ name = "hf-libero"
version = "0.1.4"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "bddl" },
{ name = "cloudpickle" },
{ name = "easydict" },
{ name = "einops" },
{ name = "future" },
{ name = "gymnasium" },
{ name = "hf-egl-probe" },
{ name = "hydra-core" },
{ name = "matplotlib" },
{ name = "mujoco" },
{ name = "numpy" },
{ name = "opencv-python" },
{ name = "robomimic" },
{ name = "robosuite" },
{ name = "thop" },
{ name = "transformers" },
{ name = "wandb" },
{ name = "bddl", marker = "sys_platform == 'linux'" },
{ name = "cloudpickle", marker = "sys_platform == 'linux'" },
{ name = "easydict", marker = "sys_platform == 'linux'" },
{ name = "einops", marker = "sys_platform == 'linux'" },
{ name = "future", marker = "sys_platform == 'linux'" },
{ name = "gymnasium", marker = "sys_platform == 'linux'" },
{ name = "hf-egl-probe", marker = "sys_platform == 'linux'" },
{ name = "hydra-core", marker = "sys_platform == 'linux'" },
{ name = "matplotlib", marker = "sys_platform == 'linux'" },
{ name = "mujoco", marker = "sys_platform == 'linux'" },
{ name = "numpy", marker = "sys_platform == 'linux'" },
{ name = "opencv-python", marker = "sys_platform == 'linux'" },
{ name = "robomimic", marker = "sys_platform == 'linux'" },
{ name = "robosuite", marker = "sys_platform == 'linux'" },
{ name = "thop", marker = "sys_platform == 'linux'" },
{ name = "transformers", marker = "sys_platform == 'linux'" },
{ name = "wandb", marker = "sys_platform == 'linux'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/af/aa/4e9eb8715e0bff9cb6553db563a35d253393097d446f82bd53575e8b253d/hf_libero-0.1.4.tar.gz", hash = "sha256:c058d67ad5a2b589529c14d614282ef4cca3a7763dafa134f58a6c9039657e34", size = 2961319, upload-time = "2026-06-10T09:56:13.994Z" }
wheels = [
@@ -2123,9 +2123,9 @@ name = "hydra-core"
version = "1.3.4"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "antlr4-python3-runtime" },
{ name = "omegaconf" },
{ name = "packaging" },
{ name = "antlr4-python3-runtime", marker = "sys_platform == 'linux'" },
{ name = "omegaconf", marker = "sys_platform == 'linux'" },
{ name = "packaging", marker = "sys_platform == 'linux'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/10/dd/220f0e91743136725352497e98540772a01fc7c3ab96ff16c3c74424e984/hydra_core-1.3.4.tar.gz", hash = "sha256:ad0f7b05a0242255a8984d5a4ed2f6847f7b783ed727368a2c0155ec52d6c34c", size = 3263348, upload-time = "2026-07-04T16:25:38.891Z" }
wheels = [
@@ -2678,11 +2678,11 @@ name = "jupytext"
version = "1.19.5"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "markdown-it-py" },
{ name = "mdit-py-plugins" },
{ name = "nbformat" },
{ name = "packaging" },
{ name = "pyyaml" },
{ name = "markdown-it-py", marker = "sys_platform == 'linux'" },
{ name = "mdit-py-plugins", marker = "sys_platform == 'linux'" },
{ name = "nbformat", marker = "sys_platform == 'linux'" },
{ name = "packaging", marker = "sys_platform == 'linux'" },
{ name = "pyyaml", marker = "sys_platform == 'linux'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/a1/ca/473f8ebb101553fb2ea6ab1d34324d6677844c968947ac050c759d539f2c/jupytext-1.19.5.tar.gz", hash = "sha256:605026446d605aa54fd7f7fc69df6ae51c7a46053d4cebf05afdc64d66de3df0", size = 4600916, upload-time = "2026-07-21T22:00:29.198Z" }
wheels = [
@@ -2903,6 +2903,7 @@ all = [
{ name = "ruff" },
{ name = "scikit-image" },
{ name = "scipy" },
{ name = "sentencepiece" },
{ name = "teleop" },
{ name = "timm" },
{ name = "torchcodec", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'AMD64' and sys_platform == 'linux') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'arm64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'win32'" },
@@ -3151,6 +3152,7 @@ phone = [
]
pi = [
{ name = "scipy" },
{ name = "sentencepiece" },
{ name = "transformers" },
]
placo-dep = [
@@ -3211,6 +3213,9 @@ sarm = [
scipy-dep = [
{ name = "scipy" },
]
sentencepiece-dep = [
{ name = "sentencepiece" },
]
smolvla = [
{ name = "accelerate" },
{ name = "num2words" },
@@ -3411,6 +3416,7 @@ requires-dist = [
{ name = "lerobot", extras = ["scipy-dep"], marker = "extra == 'phone'" },
{ name = "lerobot", extras = ["scipy-dep"], marker = "extra == 'pi'" },
{ name = "lerobot", extras = ["scipy-dep"], marker = "extra == 'wallx'" },
{ name = "lerobot", extras = ["sentencepiece-dep"], marker = "extra == 'pi'" },
{ name = "lerobot", extras = ["smolvla"], marker = "extra == 'all'" },
{ name = "lerobot", extras = ["test"], marker = "extra == 'all'" },
{ name = "lerobot", extras = ["timm-dep"], marker = "extra == 'groot'" },
@@ -3486,6 +3492,7 @@ requires-dist = [
{ name = "scikit-image", marker = "extra == 'video-benchmark'", specifier = ">=0.23.2,<0.26.0" },
{ name = "scipy", marker = "extra == 'all'", specifier = ">=1.14.0,<2.0.0" },
{ name = "scipy", marker = "extra == 'scipy-dep'", specifier = ">=1.14.0,<2.0.0" },
{ name = "sentencepiece", marker = "extra == 'sentencepiece-dep'", specifier = ">=0.2.0,<0.3.0" },
{ name = "setuptools", specifier = ">=71.0.0,<81.0.0" },
{ name = "teleop", marker = "extra == 'phone'", specifier = ">=0.1.0,<0.2.0" },
{ name = "termcolor", specifier = ">=2.4.0,<4.0.0" },
@@ -3502,7 +3509,7 @@ requires-dist = [
{ name = "transformers", marker = "extra == 'transformers-dep'", specifier = ">=5.4.0,<5.6.0" },
{ name = "wandb", marker = "extra == 'training'", specifier = ">=0.24.0,<0.28.0" },
]
provides-extras = ["dataset", "training", "hardware", "viz", "core-scripts", "evaluation", "dataset-viz", "av-dep", "pygame-dep", "placo-dep", "transformers-dep", "grpcio-dep", "accelerate-dep", "can-dep", "peft-dep", "scipy-dep", "diffusers-dep", "qwen-vl-utils-dep", "matplotlib-dep", "pyserial-dep", "deepdiff-dep", "pynput-dep", "pyzmq-dep", "motorbridge-dep", "motorbridge-smart-servo-dep", "timm-dep", "feetech", "dynamixel", "damiao", "robstride", "openarms", "gamepad", "hopejr", "lekiwi", "unitree-g1", "reachy2", "rebot", "kinematics", "intelrealsense", "phone", "diffusion", "wallx", "pi", "molmoact2", "smolvla", "multi-task-dit", "groot", "sarm", "robometer", "topreward", "xvla", "eo1", "fastwam", "evo1", "hilserl", "vla-jepa", "lingbot-va", "async", "peft", "annotations", "dev", "notebook", "test", "video-benchmark", "aloha", "pusht", "libero", "metaworld", "all"]
provides-extras = ["dataset", "training", "hardware", "viz", "core-scripts", "evaluation", "dataset-viz", "av-dep", "pygame-dep", "placo-dep", "transformers-dep", "sentencepiece-dep", "grpcio-dep", "accelerate-dep", "can-dep", "peft-dep", "scipy-dep", "diffusers-dep", "qwen-vl-utils-dep", "matplotlib-dep", "pyserial-dep", "deepdiff-dep", "pynput-dep", "pyzmq-dep", "motorbridge-dep", "motorbridge-smart-servo-dep", "timm-dep", "feetech", "dynamixel", "damiao", "robstride", "openarms", "gamepad", "hopejr", "lekiwi", "unitree-g1", "reachy2", "rebot", "kinematics", "intelrealsense", "phone", "diffusion", "wallx", "pi", "molmoact2", "smolvla", "multi-task-dit", "groot", "sarm", "robometer", "topreward", "xvla", "eo1", "fastwam", "evo1", "hilserl", "vla-jepa", "lingbot-va", "async", "peft", "annotations", "dev", "notebook", "test", "video-benchmark", "aloha", "pusht", "libero", "metaworld", "all"]
[[package]]
name = "librt"
@@ -3817,7 +3824,7 @@ name = "mdit-py-plugins"
version = "0.6.1"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "markdown-it-py" },
{ name = "markdown-it-py", marker = "sys_platform == 'linux'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/59/fc/f8d0863f8862f25602c0404d75568e89fb6b4109804645e5cdfb1be5cf56/mdit_py_plugins-0.6.1.tar.gz", hash = "sha256:a2bca0f039f39dbd35fb74ae1b5f998608c437463371f0ff7f49a19a17a114d0", size = 56114, upload-time = "2026-05-13T09:03:38.91Z" }
wheels = [
@@ -4296,8 +4303,8 @@ name = "numba"
version = "0.66.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "llvmlite" },
{ name = "numpy" },
{ name = "llvmlite", marker = "sys_platform == 'linux'" },
{ name = "numpy", marker = "sys_platform == 'linux'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/ae/a0/570e3dc53e5602b49108f62a13e529f1eec8bfc7ef37d49c825924dcf546/numba-0.66.0.tar.gz", hash = "sha256:b900e63a0e26c05ea9a6d5a3a5a0a177cb64c5011887bf43edb8c3ed2c38d363", size = 2806181, upload-time = "2026-07-01T23:12:46.36Z" }
wheels = [
@@ -4390,7 +4397,7 @@ name = "nvidia-cudnn-cu12"
version = "9.19.0.56"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "nvidia-cublas-cu12" },
{ name = "nvidia-cublas-cu12", marker = "sys_platform == 'linux'" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/09/b8/277c51962ee46fa3e5b203ac5f76107c650f781d6891e681e28e6f3e9fe6/nvidia_cudnn_cu12-9.19.0.56-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:08caaf27fe556aca82a3ee3b5aa49a77e7de0cfcb7ff4e5c29da426387a8267e", size = 656910700, upload-time = "2026-02-03T20:40:25.508Z" },
@@ -4402,7 +4409,7 @@ name = "nvidia-cufft-cu12"
version = "11.3.3.83"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "nvidia-nvjitlink-cu12" },
{ name = "nvidia-nvjitlink-cu12", marker = "sys_platform == 'linux'" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/60/bc/7771846d3a0272026c416fbb7e5f4c1f146d6d80704534d0b187dd6f4800/nvidia_cufft_cu12-11.3.3.83-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:848ef7224d6305cdb2a4df928759dca7b1201874787083b6e7550dd6765ce69a", size = 193109211, upload-time = "2025-03-07T01:44:56.873Z" },
@@ -4432,9 +4439,9 @@ name = "nvidia-cusolver-cu12"
version = "11.7.3.90"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "nvidia-cublas-cu12" },
{ name = "nvidia-cusparse-cu12" },
{ name = "nvidia-nvjitlink-cu12" },
{ name = "nvidia-cublas-cu12", marker = "sys_platform == 'linux'" },
{ name = "nvidia-cusparse-cu12", marker = "sys_platform == 'linux'" },
{ name = "nvidia-nvjitlink-cu12", marker = "sys_platform == 'linux'" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/c8/32/f7cd6ce8a7690544d084ea21c26e910a97e077c9b7f07bf5de623ee19981/nvidia_cusolver_cu12-11.7.3.90-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:db9ed69dbef9715071232caa9b69c52ac7de3a95773c2db65bdba85916e4e5c0", size = 267229841, upload-time = "2025-03-07T01:46:54.356Z" },
@@ -4446,7 +4453,7 @@ name = "nvidia-cusparse-cu12"
version = "12.5.8.93"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "nvidia-nvjitlink-cu12" },
{ name = "nvidia-nvjitlink-cu12", marker = "sys_platform == 'linux'" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/bc/f7/cd777c4109681367721b00a106f491e0d0d15cfa1fd59672ce580ce42a97/nvidia_cusparse_cu12-12.5.8.93-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:9b6c161cb130be1a07a27ea6923df8141f3c295852f4b260c65f18f3e0a091dc", size = 288117129, upload-time = "2025-03-07T01:47:40.407Z" },
@@ -4503,8 +4510,8 @@ name = "omegaconf"
version = "2.3.1"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "antlr4-python3-runtime" },
{ name = "pyyaml" },
{ name = "antlr4-python3-runtime", marker = "sys_platform == 'linux'" },
{ name = "pyyaml", marker = "sys_platform == 'linux'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/ce/3d/e4b57b8d9008c6ebe0d5eff901f91d5700cf7bdb8c8863df817463a7fd5e/omegaconf-2.3.1.tar.gz", hash = "sha256:e5e7de64aeebeddaf8e6d3f7a783b32ac2a01c0fbd9c878012caecb891a1f42a", size = 3298472, upload-time = "2026-06-11T05:05:12.885Z" }
wheels = [
@@ -4743,7 +4750,7 @@ name = "pexpect"
version = "4.9.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "ptyprocess" },
{ name = "ptyprocess", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/42/92/cc564bf6381ff43ce1f4d06852fc19a2f11d180f23dc32d9588bee2f149d/pexpect-4.9.0.tar.gz", hash = "sha256:ee7d41123f3c9911050ea2c2dac107568dc43b2d3b0c7557a33212c398ead30f", size = 166450, upload-time = "2023-11-25T09:07:26.339Z" }
wheels = [
@@ -5317,10 +5324,10 @@ name = "pyobjc-framework-applicationservices"
version = "12.2.1"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "pyobjc-core" },
{ name = "pyobjc-framework-cocoa" },
{ name = "pyobjc-framework-coretext" },
{ name = "pyobjc-framework-quartz" },
{ name = "pyobjc-core", marker = "sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32'" },
{ name = "pyobjc-framework-cocoa", marker = "sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32'" },
{ name = "pyobjc-framework-coretext", marker = "sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32'" },
{ name = "pyobjc-framework-quartz", marker = "sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/5e/4d/0ebdd8144aba94b8fe9828ccee5616a4bf53d1f8bc51cff55f3cce86d695/pyobjc_framework_applicationservices-12.2.1.tar.gz", hash = "sha256:048ea663c9ac75c44a15dc7d5b8d78cbb4c97bf1c76e83835e8d5498e184001f", size = 109342, upload-time = "2026-06-19T16:19:46.149Z" }
wheels = [
@@ -5338,7 +5345,7 @@ name = "pyobjc-framework-cocoa"
version = "12.2.1"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "pyobjc-core" },
{ name = "pyobjc-core", marker = "sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/51/34/fbe38a204643aa4e1b91391cdce07a34da565a69171ebcad08de7438a556/pyobjc_framework_cocoa-12.2.1.tar.gz", hash = "sha256:b94b37fe5730e5ae1fb0052912cd174e6ec329b0bfba4a012ae5db1014b5864b", size = 3125751, upload-time = "2026-06-19T16:20:05.159Z" }
wheels = [
@@ -5356,9 +5363,9 @@ name = "pyobjc-framework-coretext"
version = "12.2.1"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "pyobjc-core" },
{ name = "pyobjc-framework-cocoa" },
{ name = "pyobjc-framework-quartz" },
{ name = "pyobjc-core", marker = "sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32'" },
{ name = "pyobjc-framework-cocoa", marker = "sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32'" },
{ name = "pyobjc-framework-quartz", marker = "sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/5a/9c/4c7f452059dc1d3845b8e627b9113c247a997b9b07518e848c2ab7ff3149/pyobjc_framework_coretext-12.2.1.tar.gz", hash = "sha256:af740e784d7c592c34025ec7165f4f6c1a69b5a2d9075f06e41e4f77c212aed2", size = 97349, upload-time = "2026-06-19T16:20:22.508Z" }
wheels = [
@@ -5376,8 +5383,8 @@ name = "pyobjc-framework-quartz"
version = "12.2.1"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "pyobjc-core" },
{ name = "pyobjc-framework-cocoa" },
{ name = "pyobjc-core", marker = "sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32'" },
{ name = "pyobjc-framework-cocoa", marker = "sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/3b/f6/2a8b84dbf1fe7c04dd96ea73d991678d4e09a909f51971ecc51629bb2ab4/pyobjc_framework_quartz-12.2.1.tar.gz", hash = "sha256:b3b8b6f71e66147f8ff9e6213864cc8527e3a0b1ee90835b93ce221f4802d9b0", size = 3215521, upload-time = "2026-06-19T16:21:30.199Z" }
wheels = [
@@ -5952,18 +5959,18 @@ name = "robomimic"
version = "0.2.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "egl-probe" },
{ name = "h5py" },
{ name = "imageio" },
{ name = "imageio-ffmpeg" },
{ name = "numpy" },
{ name = "psutil" },
{ name = "tensorboard" },
{ name = "tensorboardx" },
{ name = "termcolor" },
{ name = "torch", version = "2.11.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" } },
{ name = "torchvision", version = "0.26.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" } },
{ name = "tqdm" },
{ name = "egl-probe", marker = "sys_platform == 'linux'" },
{ name = "h5py", marker = "sys_platform == 'linux'" },
{ name = "imageio", marker = "sys_platform == 'linux'" },
{ name = "imageio-ffmpeg", marker = "sys_platform == 'linux'" },
{ name = "numpy", marker = "sys_platform == 'linux'" },
{ name = "psutil", marker = "sys_platform == 'linux'" },
{ name = "tensorboard", marker = "sys_platform == 'linux'" },
{ name = "tensorboardx", marker = "sys_platform == 'linux'" },
{ name = "termcolor", marker = "sys_platform == 'linux'" },
{ name = "torch", version = "2.11.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "sys_platform == 'linux'" },
{ name = "torchvision", version = "0.26.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "sys_platform == 'linux'" },
{ name = "tqdm", marker = "sys_platform == 'linux'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/3d/c3/44b1d1ea4bcb4bbed43d19e09505f4142714451ded74020d4f679cdc89fb/robomimic-0.2.0.tar.gz", hash = "sha256:ee3bb5cf9c3e1feead6b57b43c5db738fd0a8e0c015fdf6419808af8fffdc463", size = 192919, upload-time = "2021-12-17T19:00:33.279Z" }
@@ -5972,12 +5979,12 @@ name = "robosuite"
version = "1.4.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "mujoco" },
{ name = "numba" },
{ name = "numpy" },
{ name = "opencv-python" },
{ name = "pillow" },
{ name = "scipy" },
{ name = "mujoco", marker = "sys_platform == 'linux'" },
{ name = "numba", marker = "sys_platform == 'linux'" },
{ name = "numpy", marker = "sys_platform == 'linux'" },
{ name = "opencv-python", marker = "sys_platform == 'linux'" },
{ name = "pillow", marker = "sys_platform == 'linux'" },
{ name = "scipy", marker = "sys_platform == 'linux'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/25/a1/9dd07a9a5e09c6aa032faf531da985808b34437cbf6c8f358fe8f7c47118/robosuite-1.4.0.tar.gz", hash = "sha256:a8a6233d7458dbd91bf00a86cab15aa1c178bd9d1b28d515db2cf3d152cb48e6", size = 192182147, upload-time = "2022-12-01T07:31:55.791Z" }
wheels = [
@@ -6218,6 +6225,49 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/1c/78/504fdd027da3b84ff1aecd9f6957e65f35134534ccc6da8628eb71e76d3f/send2trash-2.1.0-py3-none-any.whl", hash = "sha256:0da2f112e6d6bb22de6aa6daa7e144831a4febf2a87261451c4ad849fe9a873c", size = 17610, upload-time = "2026-01-14T06:27:35.218Z" },
]
[[package]]
name = "sentencepiece"
version = "0.2.2"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/cc/33/ea3cb3839607eb175da835244a798f797f478c5ddf0e8ecdf57ea85a4c70/sentencepiece-0.2.2.tar.gz", hash = "sha256:3d2b5e824b5622038dc7b490897efe05ebbbb9e7350fc142f3ecc8789ef9bdf6", size = 8218435, upload-time = "2026-07-12T08:39:34.701Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/b8/13/7a562289c8d5b49ebdf3f9c1e8ab67cf14a8743b1d90c8f406bfdec36b72/sentencepiece-0.2.2-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:1edb10e520e4bddf74d85b0f5ae74cc2d60c2b448885080bfb618bc2b3a49f6b", size = 2188384, upload-time = "2026-07-12T08:38:28.486Z" },
{ url = "https://files.pythonhosted.org/packages/85/d1/912f14fd5eae168aba726ffb6a9a2dc1c71fe7676c53da6f5c442b886d4a/sentencepiece-0.2.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:f7c06c751c19d923435a54bff4f7e66e728fad160e8da28254f133abc9725820", size = 1441553, upload-time = "2026-07-12T08:38:30.552Z" },
{ url = "https://files.pythonhosted.org/packages/bd/44/caa9cab5f261a019e2808bc5046152775dc57352ba9cbae7525e9e7a1ed4/sentencepiece-0.2.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:38111ed1f79268f399c505028023d5eaaf0ab4e5eafceb709468b0d3323e7838", size = 1347176, upload-time = "2026-07-12T08:38:32.211Z" },
{ url = "https://files.pythonhosted.org/packages/19/90/cd798935668cff71d309d8ff10385844ecf216b1fe454f1993ed8bf2cb91/sentencepiece-0.2.2-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:cbce24284f51f71d10a42b7b9c964dcb9048b28f1c8e5db40bcbcb6f428cba6a", size = 1325200, upload-time = "2026-07-12T08:38:33.689Z" },
{ url = "https://files.pythonhosted.org/packages/b6/2d/37e3da037318a70066ded0d51bc2a7f35491ae6338dd993d5eb1503fc3b5/sentencepiece-0.2.2-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c8a168b040bc61681293f79a949b5d911c8e25086f4260285b8d97ab5f1195da", size = 1397736, upload-time = "2026-07-12T08:38:35.771Z" },
{ url = "https://files.pythonhosted.org/packages/8d/11/753fca2e6b109be3ab7867abf357dfe48677fe726ae5a5363d0b54ca9450/sentencepiece-0.2.2-cp312-cp312-win_amd64.whl", hash = "sha256:7c6e7bf684dc12145bfa685d3060beaea55139134ba848289bee514ed42e7383", size = 1248030, upload-time = "2026-07-12T08:38:37.604Z" },
{ url = "https://files.pythonhosted.org/packages/e2/0a/70efbe861ca182d7d4b6e1a20f58e043400848fa9f2915229f082e221648/sentencepiece-0.2.2-cp312-cp312-win_arm64.whl", hash = "sha256:76ff5814db72e7462dece042d7593cdf102b8ec82c2b1cc201a2add34ee3050d", size = 1187325, upload-time = "2026-07-12T08:38:39.348Z" },
{ url = "https://files.pythonhosted.org/packages/b9/a3/b3b05095c174d6e80d37d5ddc2f57c2c56237333e7bbd6079cf3243c2a8a/sentencepiece-0.2.2-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:77c3ce990b23441e5ecfa5bce181fd6f408b564aeb6d7e1d1e7de9c5612501c8", size = 2188346, upload-time = "2026-07-12T08:38:41.089Z" },
{ url = "https://files.pythonhosted.org/packages/ca/f3/72ebc4acb10a06bcf7503fbc6091c8f5db68300f6aac4356c09e6c76e0e1/sentencepiece-0.2.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:fd523c4992041faa5c2b3cde62253d11a96c30d73a34afe48a486e8e2254cd1c", size = 1441434, upload-time = "2026-07-12T08:38:42.56Z" },
{ url = "https://files.pythonhosted.org/packages/34/db/f9ea1a6844b4fa5dfe2312095cd866a1f724cd0905054ab9d5991778ba50/sentencepiece-0.2.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:201a8e0f55501a76e08dbf2c54bc45f4642b379271e89c667d517bfbc2191f2a", size = 1347267, upload-time = "2026-07-12T08:38:44.389Z" },
{ url = "https://files.pythonhosted.org/packages/32/4f/31c1073314ad94466bca37d29581761d70110237ee3d46b0efece59a8c1e/sentencepiece-0.2.2-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8eed98514bffe5ecac37f493f91869c351fbb05629328bfdbc08502c6c094dc0", size = 1324980, upload-time = "2026-07-12T08:38:46.304Z" },
{ url = "https://files.pythonhosted.org/packages/59/b4/a0356fa04d6a14337a6e0e443556785a0422c53ec58baae6b9568120eb0f/sentencepiece-0.2.2-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:64b656f025355cf8c51abe9fbe3848540756c6d7ca5e6791b1afa664bc24c7cb", size = 1397593, upload-time = "2026-07-12T08:38:48.302Z" },
{ url = "https://files.pythonhosted.org/packages/09/fa/d2d6369257fd2f0de616b1c7110b73fab409ef61b14f1b9e0010ed325914/sentencepiece-0.2.2-cp313-cp313-win_amd64.whl", hash = "sha256:74f0ee601047c0c12a783088b51be4e6214a62ecd9e02278c477433cd16e0ed9", size = 1247987, upload-time = "2026-07-12T08:38:50.15Z" },
{ url = "https://files.pythonhosted.org/packages/17/ee/2bb594da6fd95e32f29057f1aa7fa996701b8980090923c2d8711fdc0a24/sentencepiece-0.2.2-cp313-cp313-win_arm64.whl", hash = "sha256:b23fe17779834d3c27aaf2edac9486d04cca1a7deb8f5facda35150ac6263a91", size = 1187250, upload-time = "2026-07-12T08:38:52.246Z" },
{ url = "https://files.pythonhosted.org/packages/58/9c/dfc82846460e7a712310f5613f23d8b553cabb4e2e648663c11d8382af56/sentencepiece-0.2.2-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:72b7825b331b1b7e7c45be2e674b3e3c65af608fa376bad2d851b20aaf0cdc78", size = 2223080, upload-time = "2026-07-12T08:38:54.391Z" },
{ url = "https://files.pythonhosted.org/packages/8d/4e/3ff12cebe6d31662d9ceeabfb282de20bd0d6098fa282b4a3b8305abc7e8/sentencepiece-0.2.2-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:d795c4ac689a57f9d4ba2288126ec7901d389ad5827d2f8b8533c883974fe563", size = 1458511, upload-time = "2026-07-12T08:38:56.811Z" },
{ url = "https://files.pythonhosted.org/packages/59/5a/16d51d05360be4cee3ebfe4837c184054c4eed16cabaeb3b039524e9a000/sentencepiece-0.2.2-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:3ab3f1ae98970b5590e2209341522718900ba19bcc2c207ffaa6bd417ad960c5", size = 1361138, upload-time = "2026-07-12T08:38:58.808Z" },
{ url = "https://files.pythonhosted.org/packages/0f/af/c30ee2a9f99d51db9844acaa8fa0b611a97c2fa7116646fa43db3300b187/sentencepiece-0.2.2-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3ec27c152a1f1b24bc9168b55a5880f3c16e2334e697da6f55a1046a22405a3d", size = 1328625, upload-time = "2026-07-12T08:39:00.849Z" },
{ url = "https://files.pythonhosted.org/packages/3e/1a/4c6b39d03f5ba8439509adbd5a23c9538088a3cb679e7a47b911e8442bc6/sentencepiece-0.2.2-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:59d6588712101ccfcae9b03692be3aaae1514c2078666d7b05f15ba3a702e41b", size = 1398595, upload-time = "2026-07-12T08:39:02.86Z" },
{ url = "https://files.pythonhosted.org/packages/0f/bc/9eedddcec1fd57bc70200fa3ebf792d18fa63527a5369581cd416c81f97f/sentencepiece-0.2.2-cp313-cp313t-win_amd64.whl", hash = "sha256:89625fb43765cccaa1443b9adb61f283e5fe4cb1536728205d06bada730caa53", size = 1259346, upload-time = "2026-07-12T08:39:04.559Z" },
{ url = "https://files.pythonhosted.org/packages/41/15/7e74c8533848866ff560b29f7d8719921b76c4ec7149592d6d28e0deee75/sentencepiece-0.2.2-cp313-cp313t-win_arm64.whl", hash = "sha256:4f0603267cd15b92b68c2c0e852a441507614b70dc7773659baa6b8c214a91fd", size = 1196596, upload-time = "2026-07-12T08:39:06.454Z" },
{ url = "https://files.pythonhosted.org/packages/0b/7e/f5df63edb6bcb46c1343cfa5d9192d73a4eb61af2e800d9402efff387523/sentencepiece-0.2.2-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:c62bd361cec1f5b556eb8210264ecfff37486cd990c3386cc00310f26c54090a", size = 2190240, upload-time = "2026-07-12T08:39:08.178Z" },
{ url = "https://files.pythonhosted.org/packages/52/0a/095d183b453b2a2e20b016829029c58eca90adc1c9911113e5d26fff45ed/sentencepiece-0.2.2-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:46ba07b543add034de0ff47ac5f907e9a06682f91d85121a972764628933be6b", size = 1442220, upload-time = "2026-07-12T08:39:09.91Z" },
{ url = "https://files.pythonhosted.org/packages/d1/18/823954c9c90e74eba09fb96752dc37a5555df00d69866cb9406d1725dc7e/sentencepiece-0.2.2-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:79bac5a251f23a7341e28fda9ce0d5319edf45328239ce037c0682936f137906", size = 1348056, upload-time = "2026-07-12T08:39:11.744Z" },
{ url = "https://files.pythonhosted.org/packages/10/ca/1b6c251321901cbf8a2d2e48b8b70eb82a449011b766af52a228d0a90b6b/sentencepiece-0.2.2-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1402d8ee36f0d851cea8eee4dbb85fea14643b7503cf4d00d102eec0fe3ca719", size = 1325463, upload-time = "2026-07-12T08:39:13.413Z" },
{ url = "https://files.pythonhosted.org/packages/24/b3/718847349da7b25c8220ed86d85b89080af94740b2d87a59198104ae5c51/sentencepiece-0.2.2-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8d44b20234905ff022b7d535f79d1f823ad7670c9851cc4f03cdc34787cdb3ab", size = 1398138, upload-time = "2026-07-12T08:39:15.564Z" },
{ url = "https://files.pythonhosted.org/packages/33/fe/4906f12c458274edd96387e4baaad7c6f064a2b7c11a1cc2401c8a7bd483/sentencepiece-0.2.2-cp314-cp314-win_amd64.whl", hash = "sha256:63250cfab8b80a1ef82a614eb2b3cadfec2c405f870cedc139d08e2f063eb708", size = 1356144, upload-time = "2026-07-12T08:39:17.313Z" },
{ url = "https://files.pythonhosted.org/packages/d3/eb/22f89b6542aba400b0007cf0b1697cc3f99be8fb682fdb4c05eec450e33f/sentencepiece-0.2.2-cp314-cp314-win_arm64.whl", hash = "sha256:65d84ec36888de4a848eee5f910e67fbc79b064685ef1e10a502e14520ead9c9", size = 1294351, upload-time = "2026-07-12T08:39:18.967Z" },
{ url = "https://files.pythonhosted.org/packages/84/c4/7afe8c2315b76e46818851a057e50a378a0382aa00b970a1fa444181b6f6/sentencepiece-0.2.2-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:d254c98ca6387655400b3959c33c83efd807f5edeb608e3aca45800ceaa77151", size = 2223281, upload-time = "2026-07-12T08:39:20.978Z" },
{ url = "https://files.pythonhosted.org/packages/98/42/fb678e472c554ef086be6375d20060ca610a2c4218854d4c091001fc6f91/sentencepiece-0.2.2-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:3fd9ce2ab4460c713cfdeb4aca693ca6732a11538e05fb332d5af42e3d7fde25", size = 1458779, upload-time = "2026-07-12T08:39:22.812Z" },
{ url = "https://files.pythonhosted.org/packages/78/52/ffe402b13bce1889228a98dc6cd86ae8afac1112362236be3468be784441/sentencepiece-0.2.2-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:7fc14c1585139fa6b68775e616a6b90cf622ebf219f9558c0aeaf5d253ee6c9b", size = 1361736, upload-time = "2026-07-12T08:39:24.602Z" },
{ url = "https://files.pythonhosted.org/packages/78/4a/2288f60e7283583ec0a0f16e72f9c8e68557d7e7a4b585d2cda4f9f47e64/sentencepiece-0.2.2-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:df88b0c34f2fa909d322f7b06b1398e1e81af4b2f42a7b8e3556f928b25d1811", size = 1328155, upload-time = "2026-07-12T08:39:26.422Z" },
{ url = "https://files.pythonhosted.org/packages/26/31/5dd6882ebe899f741a5cfe40ff56c6efc06bc26ee287abdb723b671f409c/sentencepiece-0.2.2-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3f5851441ab1ef8634963a5100b733a8bbeefe623e0c5c005b1f1f3880e574cf", size = 1398307, upload-time = "2026-07-12T08:39:28.637Z" },
{ url = "https://files.pythonhosted.org/packages/da/05/7d7780fa63f4b8c1821953b916e25f89ae8f14d4da6ba91e10f6d06dc2b4/sentencepiece-0.2.2-cp314-cp314t-win_amd64.whl", hash = "sha256:046b15ea22d8042e2e173561d464ec3b64a9c2081324df70ebce7bf7ebb3e497", size = 1367133, upload-time = "2026-07-12T08:39:30.546Z" },
{ url = "https://files.pythonhosted.org/packages/49/a1/70007fef3f818c688de4a730f98024a671599ab67f20270f8efb03d69dcc/sentencepiece-0.2.2-cp314-cp314t-win_arm64.whl", hash = "sha256:fa9f5ef0e2a82233dd0b8b32ea3f5710e0c44afbc07ed3620219f32601e56090", size = 1302760, upload-time = "2026-07-12T08:39:32.457Z" },
]
[[package]]
name = "sentry-sdk"
version = "2.66.1"
@@ -6398,16 +6448,16 @@ name = "tensorboard"
version = "2.20.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "absl-py" },
{ name = "grpcio" },
{ name = "markdown" },
{ name = "numpy" },
{ name = "packaging" },
{ name = "pillow" },
{ name = "protobuf" },
{ name = "setuptools" },
{ name = "tensorboard-data-server" },
{ name = "werkzeug" },
{ name = "absl-py", marker = "sys_platform == 'linux'" },
{ name = "grpcio", marker = "sys_platform == 'linux'" },
{ name = "markdown", marker = "sys_platform == 'linux'" },
{ name = "numpy", marker = "sys_platform == 'linux'" },
{ name = "packaging", marker = "sys_platform == 'linux'" },
{ name = "pillow", marker = "sys_platform == 'linux'" },
{ name = "protobuf", marker = "sys_platform == 'linux'" },
{ name = "setuptools", marker = "sys_platform == 'linux'" },
{ name = "tensorboard-data-server", marker = "sys_platform == 'linux'" },
{ name = "werkzeug", marker = "sys_platform == 'linux'" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/9c/d9/a5db55f88f258ac669a92858b70a714bbbd5acd993820b41ec4a96a4d77f/tensorboard-2.20.0-py3-none-any.whl", hash = "sha256:9dc9f978cb84c0723acf9a345d96c184f0293d18f166bb8d59ee098e6cfaaba6", size = 5525680, upload-time = "2025-07-17T19:20:49.638Z" },
@@ -6427,9 +6477,9 @@ name = "tensorboardx"
version = "2.6.5"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "numpy" },
{ name = "packaging" },
{ name = "protobuf" },
{ name = "numpy", marker = "sys_platform == 'linux'" },
{ name = "packaging", marker = "sys_platform == 'linux'" },
{ name = "protobuf", marker = "sys_platform == 'linux'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/48/a9/fc520ea91ab1f3ba51cbf3fe24f2b6364ed3b49046969e0868d46d6da372/tensorboardx-2.6.5.tar.gz", hash = "sha256:ca176db3997ee8c07d2eb77381225956a3fd1c10c91beafab1f17069adc47017", size = 4770195, upload-time = "2026-04-03T15:40:23.803Z" }
wheels = [
@@ -6464,7 +6514,7 @@ name = "thop"
version = "0.1.1.post2209072238"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "torch", version = "2.11.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" } },
{ name = "torch", version = "2.11.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "sys_platform == 'linux'" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/bb/0f/72beeab4ff5221dc47127c80f8834b4bcd0cb36f6ba91c0b1d04a1233403/thop-0.1.1.post2209072238-py3-none-any.whl", hash = "sha256:01473c225231927d2ad718351f78ebf7cffe6af3bed464c4f1ba1ef0f7cdda27", size = 15443, upload-time = "2022-09-07T14:38:37.211Z" },
@@ -6570,13 +6620,13 @@ resolution-markers = [
"python_full_version < '3.13' and sys_platform == 'win32'",
]
dependencies = [
{ name = "filelock" },
{ name = "fsspec" },
{ name = "jinja2" },
{ name = "networkx" },
{ name = "setuptools" },
{ name = "sympy" },
{ name = "typing-extensions" },
{ name = "filelock", marker = "sys_platform != 'linux'" },
{ name = "fsspec", marker = "sys_platform != 'linux'" },
{ name = "jinja2", marker = "sys_platform != 'linux'" },
{ name = "networkx", marker = "sys_platform != 'linux'" },
{ name = "setuptools", marker = "sys_platform != 'linux'" },
{ name = "sympy", marker = "sys_platform != 'linux'" },
{ name = "typing-extensions", marker = "sys_platform != 'linux'" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/6f/8b/69e3008d78e5cee2b30183340cc425081b78afc5eff3d080daab0adda9aa/torch-2.11.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:4b5866312ee6e52ea625cd211dcb97d6a2cdc1131a5f15cc0d87eec948f6dd34", size = 80606338, upload-time = "2026-03-23T18:11:34.781Z" },
@@ -6610,20 +6660,20 @@ resolution-markers = [
"python_full_version < '3.13' and platform_machine != 'AMD64' and platform_machine != 'aarch64' and platform_machine != 'arm64' and platform_machine != 'x86_64' and sys_platform == 'linux'",
]
dependencies = [
{ name = "cuda-bindings" },
{ name = "cuda-toolkit", extra = ["cublas", "cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"] },
{ name = "filelock" },
{ name = "fsspec" },
{ name = "jinja2" },
{ name = "networkx" },
{ name = "nvidia-cudnn-cu12" },
{ name = "nvidia-cusparselt-cu12" },
{ name = "nvidia-nccl-cu12" },
{ name = "nvidia-nvshmem-cu12" },
{ name = "setuptools" },
{ name = "sympy" },
{ name = "triton" },
{ name = "typing-extensions" },
{ name = "cuda-bindings", marker = "sys_platform == 'linux'" },
{ name = "cuda-toolkit", extra = ["cublas", "cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], marker = "sys_platform == 'linux'" },
{ name = "filelock", marker = "sys_platform == 'linux'" },
{ name = "fsspec", marker = "sys_platform == 'linux'" },
{ name = "jinja2", marker = "sys_platform == 'linux'" },
{ name = "networkx", marker = "sys_platform == 'linux'" },
{ name = "nvidia-cudnn-cu12", marker = "sys_platform == 'linux'" },
{ name = "nvidia-cusparselt-cu12", marker = "sys_platform == 'linux'" },
{ name = "nvidia-nccl-cu12", marker = "sys_platform == 'linux'" },
{ name = "nvidia-nvshmem-cu12", marker = "sys_platform == 'linux'" },
{ name = "setuptools", marker = "sys_platform == 'linux'" },
{ name = "sympy", marker = "sys_platform == 'linux'" },
{ name = "triton", marker = "sys_platform == 'linux'" },
{ name = "typing-extensions", marker = "sys_platform == 'linux'" },
]
wheels = [
{ url = "https://download-r2.pytorch.org/whl/cu128/torch-2.11.0%2Bcu128-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:9c8f38efee365cb9d334de8a83ce52fc7e5fc9e5a7b0853285efa1b69e00b0f2", upload-time = "2026-04-27T17:41:30Z" },
@@ -6694,9 +6744,9 @@ resolution-markers = [
"python_full_version < '3.13' and sys_platform == 'win32'",
]
dependencies = [
{ name = "numpy" },
{ name = "pillow" },
{ name = "torch", version = "2.11.0", source = { registry = "https://pypi.org/simple" } },
{ name = "numpy", marker = "sys_platform != 'linux'" },
{ name = "pillow", marker = "sys_platform != 'linux'" },
{ name = "torch", version = "2.11.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux'" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/ae/e7/56b47cc3b132aea90ccce22bcb8975dec688b002150012acc842846039d0/torchvision-0.26.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:c409e1c3fdebec7a3834465086dbda8bf7680eff79abf7fd2f10c6b59520a7a4", size = 1863502, upload-time = "2026-03-23T18:12:57.326Z" },
@@ -6730,9 +6780,9 @@ resolution-markers = [
"python_full_version < '3.13' and platform_machine != 'AMD64' and platform_machine != 'aarch64' and platform_machine != 'arm64' and platform_machine != 'x86_64' and sys_platform == 'linux'",
]
dependencies = [
{ name = "numpy" },
{ name = "pillow" },
{ name = "torch", version = "2.11.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" } },
{ name = "numpy", marker = "sys_platform == 'linux'" },
{ name = "pillow", marker = "sys_platform == 'linux'" },
{ name = "torch", version = "2.11.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "sys_platform == 'linux'" },
]
wheels = [
{ url = "https://download-r2.pytorch.org/whl/cu128/torchvision-0.26.0%2Bcu128-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:63e35234aed13b6edda37056f417b5c281249669db631e706811917af36b21d7", upload-time = "2026-04-09T23:21:35Z" },
@@ -7222,7 +7272,7 @@ name = "werkzeug"
version = "3.1.8"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "markupsafe" },
{ name = "markupsafe", marker = "sys_platform == 'linux'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/dd/b2/381be8cfdee792dd117872481b6e378f85c957dd7c5bca38897b08f765fd/werkzeug-3.1.8.tar.gz", hash = "sha256:9bad61a4268dac112f1c5cd4630a56ede601b6ed420300677a869083d70a4c44", size = 875852, upload-time = "2026-04-02T18:49:14.268Z" }
wheels = [