mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-29 12:39:41 +00:00
Compare commits
10 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Martino Russi
|
51ca2be48a | ||
|
Steven Palma
|
a9879e69ed | ||
|
Steven Palma
|
9d82bb9871 | ||
|
Steven Palma
|
c5371d0691 | ||
|
Steven Palma
|
b2c062c0f4 | ||
|
Maxime Ellerbach
|
051b13573e | ||
|
Pepijn
|
7de2e4c1ef | ||
|
Nikodem Bartnik
|
8db50611c2 | ||
|
Maxime Ellerbach
|
92f96f33b3 | ||
|
Steven Palma
|
d4b3ca569c |
@@ -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), [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) |
|
||||
| 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) |
|
||||
|
||||
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
|
||||
|
||||
|
||||
@@ -63,8 +63,6 @@
|
||||
title: π₀-FAST (Pi0Fast)
|
||||
- local: pi05
|
||||
title: π₀.₅ (Pi05)
|
||||
- local: pi052
|
||||
title: π₀.₅ with language supervision (Pi052)
|
||||
- local: molmoact2
|
||||
title: MolmoAct2
|
||||
- local: vla_jepa
|
||||
|
||||
@@ -150,14 +150,14 @@ class MyPolicy(PreTrainedPolicy):
|
||||
|
||||
The methods called by the train/eval loops:
|
||||
|
||||
| Method | Used by | What it does |
|
||||
| ----------------------------------------------------------------- | ----------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
|
||||
| `reset() -> None` | `lerobot-eval` | Clear per-episode state at the start of each episode. |
|
||||
| `select_action(batch, **kwargs) -> Tensor` | `lerobot-eval` | Return the next action `(B, action_dim)`. Called every step. |
|
||||
| `predict_action_chunk(batch, **kwargs) -> Tensor` | the policy itself | Return an action chunk `(B, chunk_size, action_dim)`. Currently abstract on the base class — raise `NotImplementedError` if your policy doesn't chunk. |
|
||||
| `forward(batch, reduction="mean") -> tuple[Tensor, dict \| None]` | `lerobot-train` | Return `(loss, output_dict)`. Accept `reduction="none"` if you want to support per-sample weighting. |
|
||||
| `get_optim_params() -> dict` | the optimizer | Return `self.parameters()` for simple policies; return a named parameter dict for [multi-optimizer policies](https://github.com/huggingface/lerobot/blob/ecd38c50d7d15b4184cf42649ff1185ee2e11eeb/src/lerobot/policies/sac/modeling_sac.py#L61-L73). |
|
||||
| `update() -> None` _(optional)_ | `lerobot-train` | Called after each optimizer step _if defined_. Use for EMA, target nets, replay buffers (TDMPC uses this). |
|
||||
| Method | Used by | What it does |
|
||||
| ----------------------------------------------------------------- | ----------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
|
||||
| `reset() -> None` | `lerobot-eval` | Clear per-episode state at the start of each episode. |
|
||||
| `select_action(batch, **kwargs) -> Tensor` | `lerobot-eval` | Return the next action `(B, action_dim)`. Called every step. |
|
||||
| `predict_action_chunk(batch, **kwargs) -> Tensor` | the policy itself | Return an action chunk `(B, chunk_size, action_dim)`. Currently abstract on the base class — raise `NotImplementedError` if your policy doesn't chunk. |
|
||||
| `forward(batch, reduction="mean") -> tuple[Tensor, dict \| None]` | `lerobot-train` | Return `(loss, output_dict)`. Accept `reduction="none"` if you want to support per-sample weighting. |
|
||||
| `get_optim_params() -> dict` | the optimizer | Return `self.parameters()` for simple policies; return a named parameter dict for multi-optimizer policies (see `get_optim_params` in [`modeling_act.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/act/modeling_act.py) for a per-group learning-rate example). |
|
||||
| `update() -> None` _(optional)_ | `lerobot-train` | Called after each optimizer step _if defined_. Use for EMA, target nets, replay buffers (TDMPC uses this). |
|
||||
|
||||
Batches are flat dictionaries keyed by the constants in [`lerobot.utils.constants`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/utils/constants.py): `OBS_STATE` (`observation.state.<motor>`), `OBS_IMAGES` (`observation.images.<camera>`), `OBS_LANGUAGE`, `ACTION`, etc. Reuse the constants — don't invent new prefixes.
|
||||
|
||||
@@ -189,162 +189,6 @@ def make_my_policy_pre_post_processors(
|
||||
|
||||
---
|
||||
|
||||
## Adding high- and low-level language control
|
||||
|
||||
The policy API above is sufficient for training and standard evaluation. To use a language-conditioned policy with interactive `lerobot-rollout`, also register a runtime adapter. The adapter keeps policy-specific prompting and tokenization out of the generic control loop.
|
||||
|
||||
The runtime supports two policy shapes:
|
||||
|
||||
| Policy shape | Behavior | Adapter |
|
||||
| ---------------- | ----------------------------------------------------------------------- | ---------------------------------------------- |
|
||||
| Low-level / flat | The operator's task or subtask directly conditions action prediction. | Reuse `DirectTaskPolicyAdapter`. |
|
||||
| High + low level | The policy generates subtasks or memory, then conditions actions on it. | Subclass `BaseLanguageAdapter`, as PI052 does. |
|
||||
|
||||
During a rollout, `RuntimeState` stores the high-level task and the active language context:
|
||||
|
||||
```text
|
||||
task ──> adapter.generate_text("subtask", ...) ──> state.language_context["subtask"]
|
||||
│
|
||||
observation ──> processors ──> adapter.select_action() ─┴─> action chunk ──> robot
|
||||
```
|
||||
|
||||
The generic runtime handles generation frequency, pause/resume, prompt replacement, action queues, and dispatch. The adapter only translates between that runtime contract and your policy.
|
||||
|
||||
### Low-level policies
|
||||
|
||||
If your policy already consumes the live task through its normal preprocessor and implements `predict_action_chunk`, register the shared direct adapter. PI0.5 and MolmoAct2 use this path:
|
||||
|
||||
```python
|
||||
# src/lerobot/runtime/registry.py
|
||||
_ADAPTERS = {
|
||||
# ...
|
||||
"my_policy": "lerobot.runtime.adapter:DirectTaskPolicyAdapter",
|
||||
}
|
||||
```
|
||||
|
||||
Run it with direct-subtask mode so the operator supplies the instruction used by the action policy:
|
||||
|
||||
```bash
|
||||
lerobot-rollout \
|
||||
--language \
|
||||
--policy.path=user/my_policy_checkpoint \
|
||||
--robot.type=so101_follower \
|
||||
--robot.port=/dev/ttyACM0 \
|
||||
--direct_subtask
|
||||
```
|
||||
|
||||
The rollout context builds the observation batch with the current instruction before `DirectTaskPolicyAdapter` calls `policy.predict_action_chunk(observation)`. No text-generation method is required.
|
||||
|
||||
### Hierarchical policies
|
||||
|
||||
For a policy that generates language and actions, subclass [`BaseLanguageAdapter`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/runtime/adapter.py) and implement two methods:
|
||||
|
||||
- `generate_text(kind, observation, state, user_text=None) -> str` generates a `subtask`, `memory`, or interjection response.
|
||||
- `select_action(observation, state)` builds the low-level prompt from the active context and returns an action chunk.
|
||||
|
||||
This abbreviated adapter follows [`PI052PolicyAdapter`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/pi052/inference/pi052_adapter.py):
|
||||
|
||||
```python
|
||||
# inference/my_policy_adapter.py
|
||||
from typing import Any
|
||||
|
||||
from lerobot.runtime import RuntimeState
|
||||
from lerobot.runtime.adapter import BaseLanguageAdapter
|
||||
from lerobot.utils.constants import (
|
||||
OBS_LANGUAGE_ATTENTION_MASK,
|
||||
OBS_LANGUAGE_TOKENS,
|
||||
)
|
||||
|
||||
|
||||
class MyPolicyAdapter(BaseLanguageAdapter):
|
||||
def select_action(self, observation: dict[str, Any], state: RuntimeState):
|
||||
instruction = state.language_context.get("subtask") or state.task or ""
|
||||
tokens, attention_mask = tokenize_instruction(instruction)
|
||||
|
||||
batch = dict(observation)
|
||||
batch[OBS_LANGUAGE_TOKENS] = tokens
|
||||
batch[OBS_LANGUAGE_ATTENTION_MASK] = attention_mask
|
||||
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)
|
||||
batch, tokenizer = tokenize_messages(messages, observation)
|
||||
return self.policy.select_message(
|
||||
batch,
|
||||
tokenizer=tokenizer,
|
||||
min_new_tokens=self.gen.min_new_tokens,
|
||||
temperature=self.gen.temperature,
|
||||
top_p=self.gen.top_p,
|
||||
)
|
||||
|
||||
def build_messages(
|
||||
self, kind: str, state: RuntimeState, user_text: str | None
|
||||
) -> list[dict[str, str]]:
|
||||
if kind == "subtask":
|
||||
return [{"role": "user", "content": state.task or ""}]
|
||||
if kind == "memory":
|
||||
return [
|
||||
{"role": "user", "content": state.task or ""},
|
||||
{
|
||||
"role": "user",
|
||||
"content": f"Completed subtask: {state.extra.get('prior_subtask', '')}",
|
||||
},
|
||||
]
|
||||
if kind == "interjection":
|
||||
return [
|
||||
{"role": "user", "content": state.task or ""},
|
||||
{"role": "user", "content": user_text or ""},
|
||||
]
|
||||
raise ValueError(f"Unsupported text kind: {kind}")
|
||||
```
|
||||
|
||||
`tokenize_instruction` and `tokenize_messages` are policy-specific helpers. They must reproduce the prompt format used during training; PI052, for example, adds the discretized robot state to its low-level subtask prompt and uses the same PaliGemma formatting for `select_message`.
|
||||
|
||||
`BaseLanguageAdapter` provides the default hierarchy: regenerate a subtask at action-chunk boundaries, update memory when the subtask changes, and handle user interjections. Override `_regenerate_context` only if your policy uses a different hierarchy.
|
||||
|
||||
Register the adapter with a lazy import so importing LeRobot does not load the model or its optional dependencies:
|
||||
|
||||
```python
|
||||
# src/lerobot/runtime/registry.py
|
||||
_ADAPTERS = {
|
||||
# ...
|
||||
"my_policy": "lerobot.policies.my_policy.inference.my_policy_adapter:MyPolicyAdapter",
|
||||
}
|
||||
```
|
||||
|
||||
The key must match the policy's registered type. Once registered, the same checkpoint works through the shared entry point:
|
||||
|
||||
```bash
|
||||
lerobot-rollout \
|
||||
--language \
|
||||
--policy.path=user/my_hierarchical_checkpoint \
|
||||
--robot.type=so101_follower \
|
||||
--robot.port=/dev/ttyACM0 \
|
||||
--task="put the cup in the sink"
|
||||
```
|
||||
|
||||
For RoboCasa-compatible policies, replace the robot arguments with `--sim --sim.task=<task>`. Without `--direct_subtask`, the adapter generates the low-level subtask; with it, the operator bypasses high-level generation and supplies each subtask.
|
||||
|
||||
### Keep training and deployment aligned
|
||||
|
||||
The adapter is intentionally small, but its prompts are part of the model contract:
|
||||
|
||||
- Use the same tokenizer, role formatting, special tokens, image ordering, and state encoding as training.
|
||||
- Condition `select_action` on `state.language_context["subtask"]`, falling back to `state.task` for direct or not-yet-generated prompts.
|
||||
- Return a full action chunk from `select_action`; the runtime handles control-rate dispatch.
|
||||
- Keep optional model dependencies inside lazy imports.
|
||||
- Test adapter selection, generated-message routing, action-batch construction, and direct-subtask behavior with a lightweight fake policy.
|
||||
|
||||
PI052 is the complete in-tree reference: its [processor](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/pi052/processor_pi052.py) renders the training recipe, its [policy](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/pi052/modeling_pi052.py) exposes text and action generation, and its [adapter](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/pi052/inference/pi052_adapter.py) reconstructs those same prompts at deployment.
|
||||
|
||||
---
|
||||
|
||||
## Path A: Out-of-tree plugin
|
||||
|
||||
The fastest way to ship a policy: package it as a standalone Python distribution and install it alongside LeRobot. No PR required, you own the release cycle, and you can publish to PyPI under your own namespace.
|
||||
@@ -451,12 +295,10 @@ The file names are load-bearing: the factory does lazy imports by name, and the
|
||||
|
||||
### Wiring
|
||||
|
||||
Four places need to know about your policy. All by name.
|
||||
Two places need to know about your policy. All by name.
|
||||
|
||||
1. **`policies/__init__.py`** — re-export `MyPolicyConfig` and add it to `__all__`. **Don't** re-export the modeling class; it loads lazily through the factory (so `import lerobot` stays fast).
|
||||
2. **`factory.py:get_policy_class`** — add a branch returning `MyPolicy` from a lazy import.
|
||||
3. **`factory.py:make_policy_config`** and **`factory.py:make_pre_post_processors`** — same idea, two more branches.
|
||||
4. **`templates/lerobot_modelcard_template.md` and the root `README.md`** — the template is what `push_model_to_hub` renders into the model card of every checkpoint trained with your policy: add a one-line description of your policy in the `model_name` branches, map it in `policy_docs` so cards link to your MDX guide, and optionally add an architecture image to `diagrams`. Then add your policy to the models table in the root `README.md`, under the right category, linking to your doc page.
|
||||
1. **`policies/__init__.py`** — re-export `MyPolicyConfig` and add it to `__all__`. This import is what registers your policy: `@PreTrainedConfig.register_subclass("my_policy")` runs, and from then on the factory resolves everything by convention. **Don't** re-export the modeling class; it loads lazily through the factory (so `import lerobot` stays fast).
|
||||
2. **`templates/lerobot_modelcard_template.md` and the root `README.md`** — the template is what `push_model_to_hub` renders into the model card of every checkpoint trained with your policy: add a one-line description of your policy in the `model_name` branches, map it in `policy_docs` so cards link to your MDX guide, and optionally add an architecture image to `diagrams`. Then add your policy to the models table in the root `README.md`, under the right category, linking to your doc page.
|
||||
|
||||
Mirror an existing policy that's structurally similar to yours; the diff is small.
|
||||
|
||||
@@ -488,6 +330,10 @@ This way:
|
||||
|
||||
Add a matching extra to [`pyproject.toml`](https://github.com/huggingface/lerobot/blob/main/pyproject.toml) `[project.optional-dependencies]` and include it in the `all` extra so `pip install 'lerobot[all]'` keeps installing everything.
|
||||
|
||||
### Avoid copying a modeling file — subclass it
|
||||
|
||||
If your policy needs to modify a backbone that already exists in `transformers` (custom conditioning, extra inputs, a swapped sub-module), **do not vendor a copy of its `modeling_*.py`**. Instead, subclass the smallest upstream unit and override only what changes. [`pi_gemma.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/pi_gemma.py) is the canonical reference: it injects AdaRMS conditioning into PaliGemma/Gemma in ~370 lines by subclassing `GemmaModel`/`PaliGemmaModel` and overriding the decoder-layer forward, instead of forking the ~2,000-line modeling file. Model surgery on a _loaded_ native model is also fine (layer truncation, tokenizer expansion, hidden-state capture — see `evo1/internvl3_embedder.py`, `eo1/modeling_eo1.py`, `groot/groot_n1_7.py` for working examples). Reviewers will ask for this pattern when a PR arrives with a copied modeling file; the only accepted exception is a model that does not exist in `transformers` at all.
|
||||
|
||||
### Benchmarks and a published checkpoint
|
||||
|
||||
A new policy is much easier to review — and far more useful — when it ships with a working checkpoint and at least one number you can reproduce.
|
||||
@@ -523,7 +369,7 @@ If your policy is real-robot-only and no sim benchmark applies, swap the sim eva
|
||||
The general expectations are in [`CONTRIBUTING.md`](https://github.com/huggingface/lerobot/blob/main/CONTRIBUTING.md) and the [PR template](https://github.com/huggingface/lerobot/blob/main/.github/PULL_REQUEST_TEMPLATE.md). On top of those, reviewers will look for:
|
||||
|
||||
- [ ] `MyPolicy` and `MyPolicyConfig` cover the surface above; `__init_subclass__` accepts the class.
|
||||
- [ ] `factory.py` and `policies/__init__.py` are wired (lazy imports for modeling).
|
||||
- [ ] `policies/__init__.py` re-exports the config (this registers the policy; the factory resolves modeling/processor by naming convention).
|
||||
- [ ] `make_my_policy_pre_post_processors` follows the naming convention.
|
||||
- [ ] Optional deps live behind a `[project.optional-dependencies]` extra and the `TYPE_CHECKING + require_package` guard.
|
||||
- [ ] `tests/policies/` updated; backward-compat artifact committed & policy-specific tests.
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
# Policy Deployment (lerobot-rollout)
|
||||
|
||||
`lerobot-rollout` is the single CLI for deploying trained policies on real robots or in an interactive simulator. It supports multiple execution strategies and inference backends, from quick evaluation to continuous recording, language-driven control, and human-in-the-loop data collection.
|
||||
`lerobot-rollout` is the single CLI for deploying trained policies on real robots. It supports multiple execution strategies and inference backends, from quick evaluation to continuous recording and human-in-the-loop data collection.
|
||||
|
||||
## Quick Start
|
||||
|
||||
@@ -197,52 +197,6 @@ Teleop is optional — if omitted the robot holds its position during the reset
|
||||
|
||||
---
|
||||
|
||||
## Interactive language control
|
||||
|
||||
Language-conditioned policies can expose a high-level text head in addition to
|
||||
their action head. Add `--language` to open-prompt one of these policies on a
|
||||
real robot. Language-only flags such as `--direct_subtask` select this mode
|
||||
automatically.
|
||||
|
||||
MolmoAct2 has no high-level planner, so use direct-subtask mode and type each
|
||||
next low-level instruction yourself:
|
||||
|
||||
```bash
|
||||
lerobot-rollout \
|
||||
--policy.path=lerobot/MolmoAct2-SO100_101-LeRobot \
|
||||
--policy.device=cuda \
|
||||
--robot.type=so101_follower \
|
||||
--robot.port=/dev/ttyACM1 \
|
||||
--robot.cameras='{"cam0":{"type":"opencv","index_or_path":"/dev/video0","width":640,"height":480,"fps":30,"fourcc":"MJPG","backend":200},"cam1":{"type":"opencv","index_or_path":"/dev/video2","width":640,"height":480,"fps":30,"fourcc":"MJPG","backend":200}}' \
|
||||
--direct_subtask \
|
||||
--robot.max_relative_target='{"shoulder_pan":5,"shoulder_lift":5,"elbow_flex":5,"wrist_flex":5,"wrist_roll":5,"gripper":5}'
|
||||
```
|
||||
|
||||
The robot starts paused. Type a subtask, then use `/resume` and `/pause` to
|
||||
control action dispatch. Check the workspace and motion limits before resuming.
|
||||
Without `--direct_subtask`, a policy such as PI052 generates its active subtask
|
||||
from the high-level `--task` itself.
|
||||
|
||||
RoboCasa uses the same runtime and processor path. `--sim` selects it
|
||||
automatically, so no robot configuration is needed:
|
||||
|
||||
```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 simulator view. Add
|
||||
`--sim.direct_subtask` to bypass the language planner and make each typed prompt
|
||||
the action policy's current subtask.
|
||||
|
||||
---
|
||||
|
||||
## Inference Backends
|
||||
|
||||
Select a backend with `--inference.type=<name>`. All strategies work with both backends.
|
||||
|
||||
@@ -141,17 +141,6 @@ sample["target_message_indices"]
|
||||
|
||||
The renderer does not apply a tokenizer chat template. Policy processors decide how to serialize the messages for their backbone, which keeps the same dataset usable across SmolVLA, Pi0.5, and any future VLM that expects OpenAI-style chat messages.
|
||||
|
||||
## Blends
|
||||
|
||||
Blend recipes select one weighted sub-recipe deterministically from the sample index.
|
||||
`recipes/subtask_mem.yaml` trains the compact core blend — high-level subtask prediction, low-level execution, and memory. `recipes/subtask_mem_vqa_speech.yaml` is the fuller variant that also adds VQA and spoken interjection responses.
|
||||
|
||||
A message recipe with a supervised assistant turn on the `low_level` stream trains
|
||||
the π0.5 paper's joint sequence instead of a blend: the target span gets text CE
|
||||
while also conditioning the action losses in the same forward.
|
||||
`recipes/subtask_joint.yaml` is the provided example; pair it with
|
||||
`--policy.joint_subtask_conditioning=true` at inference.
|
||||
|
||||
## Graceful absence
|
||||
|
||||
If both language columns are missing, `None`, or empty, `RenderMessagesStep` is a no-op.
|
||||
|
||||
@@ -1,274 +0,0 @@
|
||||
# π₀.₅ 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.
|
||||
+9
-15
@@ -109,21 +109,15 @@ 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.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.
|
||||
| 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` |
|
||||
|
||||
## Inference
|
||||
|
||||
|
||||
@@ -46,8 +46,11 @@ CMD = (
|
||||
"apt-get update -qq && apt-get install -y -qq git ffmpeg && "
|
||||
"pip install --no-deps "
|
||||
"'lerobot @ git+https://github.com/huggingface/lerobot.git@main' && "
|
||||
# Pins mirror pyproject.toml — unpinned installs pull av 18 / datasets 5 /
|
||||
# draccus 0.11, which break lerobot at import time.
|
||||
"pip install --upgrade-strategy only-if-needed "
|
||||
"datasets pyarrow av jsonlines draccus gymnasium torchcodec mergedeep pyyaml-include toml typing-inspect "
|
||||
"'datasets>=4.7.0,<5.0.0' 'pyarrow>=21.0.0,<30.0.0' 'av>=15.0.0,<16.0.0' 'draccus==0.10.0' "
|
||||
"'pandas>=2.0.0,<3.0.0' jsonlines gymnasium torchcodec mergedeep pyyaml-include toml typing-inspect "
|
||||
"openai && "
|
||||
"export VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0 && "
|
||||
"export VLLM_VIDEO_BACKEND=pyav && "
|
||||
|
||||
+1
-4
@@ -150,7 +150,6 @@ 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"]
|
||||
@@ -213,7 +212,7 @@ wallx = [
|
||||
"torchdiffeq>=0.2.4,<0.3.0",
|
||||
"lerobot[qwen-vl-utils-dep]",
|
||||
]
|
||||
pi = ["lerobot[transformers-dep]", "lerobot[scipy-dep]", "lerobot[sentencepiece-dep]"]
|
||||
pi = ["lerobot[transformers-dep]", "lerobot[scipy-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]"]
|
||||
@@ -414,8 +413,6 @@ ignore = [
|
||||
"__init__.py" = ["F401", "F403", "E402"]
|
||||
# E402: conditional-import guards (TYPE_CHECKING / is_package_available) must precede the imports they protect
|
||||
"src/lerobot/scripts/convert_dataset_v21_to_v30.py" = ["E402"]
|
||||
"src/lerobot/policies/wall_x/**" = ["N801", "N812", "SIM102", "SIM108", "SIM210", "SIM211", "B006", "B007", "SIM118"] # Supprese these as they are coming from original Qwen2_5_vl code TODO(pepijn): refactor original
|
||||
|
||||
[tool.ruff.lint.isort]
|
||||
combine-as-imports = true
|
||||
known-first-party = ["lerobot"]
|
||||
|
||||
@@ -1,227 +0,0 @@
|
||||
#!/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()
|
||||
@@ -33,8 +33,6 @@ class DatasetConfig:
|
||||
# looked up under $HF_LEROBOT_HOME/repo_id and Hub downloads use a revision-safe cache under $HF_LEROBOT_HOME/hub.
|
||||
root: str | None = None
|
||||
episodes: list[int] | None = None
|
||||
# Episode indices to drop (e.g. corrupt or heterogeneous ones). Applied on top of `episodes`.
|
||||
exclude_episodes: list[int] | None = None
|
||||
image_transforms: ImageTransformsConfig = field(default_factory=ImageTransformsConfig)
|
||||
revision: str | None = None
|
||||
use_imagenet_stats: bool = True
|
||||
@@ -64,10 +62,6 @@ class DatasetConfig:
|
||||
if len(self.episodes) != len(set(self.episodes)):
|
||||
duplicates = sorted({ep for ep in self.episodes if self.episodes.count(ep) > 1})
|
||||
raise ValueError(f"Episode indices contain duplicates: {duplicates}")
|
||||
if self.exclude_episodes is not None and any(ep < 0 for ep in self.exclude_episodes):
|
||||
raise ValueError(
|
||||
f"exclude_episodes must be non-negative, got: {[ep for ep in self.exclude_episodes if ep < 0]}"
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
|
||||
@@ -205,24 +205,30 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC): # type: igno
|
||||
f"{CONFIG_NAME} not found on the HuggingFace Hub in {model_id}"
|
||||
) from e
|
||||
|
||||
# HACK: Parse the original config to get the config subclass, so that we can
|
||||
# apply cli overrides.
|
||||
# This is very ugly, ideally we'd like to be able to do that natively with draccus
|
||||
# something like --policy.path (in addition to --policy.type)
|
||||
with draccus.config_type("json"):
|
||||
orig_config = draccus.parse(cls, config_file, args=[])
|
||||
|
||||
if config_file is None:
|
||||
raise FileNotFoundError(f"{CONFIG_NAME} not found in {model_id}")
|
||||
|
||||
with open(config_file) as f:
|
||||
config = json.load(f)
|
||||
|
||||
config.pop("type")
|
||||
# Resolve the concrete config subclass from the serialized "type" tag, then parse
|
||||
# the config (with CLI overrides) directly for that class. The "type" key is
|
||||
# stripped because draccus only consumes it when parsing the registry base class.
|
||||
policy_type = config.pop("type", None)
|
||||
if policy_type is None:
|
||||
raise ValueError(f"Missing 'type' field in {CONFIG_NAME} of {model_id}")
|
||||
try:
|
||||
config_cls = cls.get_choice_class(policy_type)
|
||||
except Exception as e:
|
||||
raise ValueError(
|
||||
f"Policy type '{policy_type}' (from {CONFIG_NAME} of {model_id}) is not registered. "
|
||||
f"Available policy types: {cls.get_known_choices()}"
|
||||
) from e
|
||||
|
||||
with tempfile.NamedTemporaryFile("w+", delete=False, suffix=".json") as f:
|
||||
json.dump(config, f)
|
||||
config_file = f.name
|
||||
|
||||
cli_overrides = policy_kwargs.pop("cli_overrides", [])
|
||||
with draccus.config_type("json"):
|
||||
return draccus.parse(orig_config.__class__, config_file, args=cli_overrides)
|
||||
return draccus.parse(config_cls, config_file, args=cli_overrides)
|
||||
|
||||
@@ -147,7 +147,7 @@ class TrainingRecipe:
|
||||
return cls.from_dict(data)
|
||||
|
||||
def _validate_message_recipe(self) -> None:
|
||||
"""Validate bindings and require text or low-level action supervision."""
|
||||
"""Ensure every templated binding is known and at least one turn is a target."""
|
||||
assert self.messages is not None
|
||||
known_bindings = set(DEFAULT_BINDINGS) | set(self.bindings or {}) | {"task"}
|
||||
|
||||
@@ -156,14 +156,8 @@ class TrainingRecipe:
|
||||
if missing:
|
||||
raise ValueError(f"MessageTurn references unknown binding(s): {sorted(missing)}")
|
||||
|
||||
has_target = any(turn.target for turn in self.messages)
|
||||
has_low_level = any(turn.stream == "low_level" for turn in self.messages)
|
||||
if not (has_target or has_low_level):
|
||||
raise ValueError(
|
||||
"Message recipes must contain at least one supervised turn — "
|
||||
"either ``target: true`` (text CE) or ``stream: low_level`` "
|
||||
"(flow/action loss)."
|
||||
)
|
||||
if not any(turn.target for turn in self.messages):
|
||||
raise ValueError("Message recipes must contain at least one target turn.")
|
||||
|
||||
def _validate_blend_recipe(self) -> None:
|
||||
"""Ensure each blend component is a non-empty, weighted message recipe."""
|
||||
|
||||
@@ -1,16 +0,0 @@
|
||||
# Predicts subtasks from tasks and trains subtask-conditioned action flow without memory or plans.
|
||||
# Requires `subtask` annotations; samples with missing `if_present` bindings do not render.
|
||||
|
||||
blend:
|
||||
|
||||
high_level_subtask:
|
||||
weight: 0.30
|
||||
messages:
|
||||
- {role: user, content: "${task}", stream: high_level}
|
||||
- {role: assistant, content: "${subtask}", stream: high_level, target: true, if_present: subtask}
|
||||
|
||||
low_level_execution:
|
||||
weight: 0.70
|
||||
messages:
|
||||
# The low-level stream trains action flow on the generated or annotated subtask.
|
||||
- {role: user, content: "${subtask}", stream: low_level, if_present: subtask}
|
||||
@@ -1,13 +0,0 @@
|
||||
# Paper-style joint sequence (pi0.5 §IV-B): one sample supervises the subtask
|
||||
# text with CE and, because the assistant turn is part of the prefix, conditions
|
||||
# the FAST and flow action losses on the same annotated subtask in one forward.
|
||||
# The supervised span is attended causally; the action losses see task + subtask.
|
||||
#
|
||||
# Pair with `--policy.joint_subtask_conditioning=true` at inference so the flow
|
||||
# prefix reproduces this layout (task turn with state + causal generated subtask).
|
||||
# Samples without a `subtask` annotation fall back to a plain task-prompt
|
||||
# low-level sample via `if_present`.
|
||||
|
||||
messages:
|
||||
- {role: user, content: "${task}", stream: low_level}
|
||||
- {role: assistant, content: "${subtask}", stream: low_level, target: true, if_present: subtask}
|
||||
@@ -1,30 +0,0 @@
|
||||
# Trains subtask prediction, subtask-conditioned action flow, and memory updates without plans.
|
||||
# Requires `subtask` and `memory`; missing `if_present` bindings skip the affected sub-recipe.
|
||||
|
||||
blend:
|
||||
|
||||
high_level_subtask:
|
||||
weight: 0.25
|
||||
messages:
|
||||
- {role: user, content: "${task}", stream: high_level}
|
||||
- {role: assistant, content: "${subtask}", stream: high_level, target: true, if_present: subtask}
|
||||
|
||||
low_level_execution:
|
||||
weight: 0.60
|
||||
messages:
|
||||
# The low-level stream trains action flow on the generated or annotated subtask.
|
||||
- {role: user, content: "${subtask}", stream: low_level, if_present: subtask}
|
||||
|
||||
memory_update:
|
||||
# `active_at` densifies sparse boundaries while preserving the prior-memory/subtask mapping.
|
||||
# Inference controls update timing through `subtask_change` events.
|
||||
weight: 0.15
|
||||
bindings:
|
||||
prior_memory: "nth_prev(style=memory, offset=1)"
|
||||
current_memory: "active_at(t, style=memory)"
|
||||
completed_subtask: "nth_prev(style=subtask, offset=1)"
|
||||
messages:
|
||||
- {role: user, content: "${task}", stream: high_level}
|
||||
- {role: assistant, content: "Previous memory: ${prior_memory}", stream: high_level, if_present: prior_memory}
|
||||
- {role: user, content: "Completed subtask: ${completed_subtask}", stream: high_level, if_present: completed_subtask}
|
||||
- {role: assistant, content: "${current_memory}", stream: high_level, target: true, if_present: current_memory}
|
||||
@@ -1,70 +0,0 @@
|
||||
# Adds memory, spoken interjection responses, and camera-grounded VQA to subtask/action training.
|
||||
# Missing optional annotations skip only their sub-recipe; `say` tool calls tokenize as `<say>...</say>`.
|
||||
|
||||
blend:
|
||||
|
||||
high_level_subtask:
|
||||
weight: 0.25
|
||||
messages:
|
||||
- {role: user, content: "${task}", stream: high_level}
|
||||
- {role: assistant, content: "${subtask}", stream: high_level, target: true, if_present: subtask}
|
||||
|
||||
low_level_execution:
|
||||
weight: 0.40
|
||||
messages:
|
||||
# The low-level stream trains action flow on the generated or annotated subtask.
|
||||
- {role: user, content: "${subtask}", stream: low_level, if_present: subtask}
|
||||
|
||||
memory_update:
|
||||
# `active_at` densifies sparse boundaries while preserving the prior-memory/subtask mapping.
|
||||
# Inference controls update timing through `subtask_change` events.
|
||||
weight: 0.10
|
||||
bindings:
|
||||
prior_memory: "nth_prev(style=memory, offset=1)"
|
||||
current_memory: "active_at(t, style=memory)"
|
||||
completed_subtask: "nth_prev(style=subtask, offset=1)"
|
||||
messages:
|
||||
- {role: user, content: "${task}", stream: high_level}
|
||||
- {role: assistant, content: "Previous memory: ${prior_memory}", stream: high_level, if_present: prior_memory}
|
||||
- {role: user, content: "Completed subtask: ${completed_subtask}", stream: high_level, if_present: completed_subtask}
|
||||
- {role: assistant, content: "${current_memory}", stream: high_level, target: true, if_present: current_memory}
|
||||
|
||||
user_interjection_response:
|
||||
weight: 0.10
|
||||
bindings:
|
||||
interjection: "emitted_at(t, style=interjection)"
|
||||
speech: "emitted_at(t, role=assistant, tool_name=say)"
|
||||
messages:
|
||||
- {role: user, content: "${task}", stream: high_level}
|
||||
- {role: user, content: "${interjection}", stream: high_level, if_present: interjection}
|
||||
# The assistant target is a `say` tool call flattened to a `<say>...</say>` marker.
|
||||
- {role: assistant, stream: high_level, target: true, if_present: speech, tool_calls_from: speech}
|
||||
|
||||
# Each camera uses a separate VQA sub-recipe for view-specific binding.
|
||||
ask_vqa_top:
|
||||
weight: 0.075
|
||||
bindings:
|
||||
vqa_query: "emitted_at(t, style=vqa, role=user, camera=observation.images.front)"
|
||||
vqa: "emitted_at(t, style=vqa, role=assistant, camera=observation.images.front)"
|
||||
messages:
|
||||
- role: user
|
||||
stream: high_level
|
||||
if_present: vqa_query
|
||||
content:
|
||||
- {type: image, feature: observation.images.front}
|
||||
- {type: text, text: "${vqa_query}"}
|
||||
- {role: assistant, content: "${vqa}", stream: high_level, target: true, if_present: vqa}
|
||||
|
||||
ask_vqa_wrist:
|
||||
weight: 0.075
|
||||
bindings:
|
||||
vqa_query: "emitted_at(t, style=vqa, role=user, camera=observation.images.wrist)"
|
||||
vqa: "emitted_at(t, style=vqa, role=assistant, camera=observation.images.wrist)"
|
||||
messages:
|
||||
- role: user
|
||||
stream: high_level
|
||||
if_present: vqa_query
|
||||
content:
|
||||
- {type: image, feature: observation.images.wrist}
|
||||
- {type: text, text: "${vqa_query}"}
|
||||
- {role: assistant, content: "${vqa}", stream: high_level, target: true, if_present: vqa}
|
||||
@@ -163,40 +163,10 @@ class DatasetReader:
|
||||
def _load_hf_dataset(self) -> datasets.Dataset:
|
||||
"""hf_dataset contains all the observations, states, actions, rewards, etc."""
|
||||
features = get_hf_features_from_features(self._meta.features)
|
||||
# Annotated datasets may have language columns absent from metadata.
|
||||
# Extend the schema before the strict Parquet cast.
|
||||
features = self._extend_features_with_language_columns(features)
|
||||
hf_dataset = load_nested_dataset(self.root / "data", features=features, episodes=self.episodes)
|
||||
hf_dataset.set_transform(hf_transform_to_torch)
|
||||
return hf_dataset
|
||||
|
||||
def _extend_features_with_language_columns(self, features: datasets.Features) -> datasets.Features:
|
||||
"""Register language columns found in Parquet but missing from metadata."""
|
||||
# Leave empty datasets to fail through the normal loading path.
|
||||
try:
|
||||
sample = next((self.root / "data").glob("*/*.parquet"))
|
||||
except StopIteration:
|
||||
return features
|
||||
|
||||
from pyarrow import parquet as _pq # noqa: PLC0415
|
||||
|
||||
schema_names = set(_pq.read_schema(sample).names)
|
||||
from .language import ( # noqa: PLC0415
|
||||
LANGUAGE_EVENTS,
|
||||
LANGUAGE_PERSISTENT,
|
||||
language_events_column_feature,
|
||||
language_persistent_column_feature,
|
||||
)
|
||||
|
||||
extra: dict[str, object] = {}
|
||||
if LANGUAGE_PERSISTENT in schema_names and LANGUAGE_PERSISTENT not in features:
|
||||
extra[LANGUAGE_PERSISTENT] = language_persistent_column_feature()
|
||||
if LANGUAGE_EVENTS in schema_names and LANGUAGE_EVENTS not in features:
|
||||
extra[LANGUAGE_EVENTS] = language_events_column_feature()
|
||||
if not extra:
|
||||
return features
|
||||
return datasets.Features({**features, **extra})
|
||||
|
||||
def _check_cached_episodes_sufficient(self) -> bool:
|
||||
"""Check if the cached dataset contains all requested episodes and their video files."""
|
||||
if self.hf_dataset is None or len(self.hf_dataset) == 0:
|
||||
|
||||
@@ -66,17 +66,6 @@ def resolve_delta_timestamps(
|
||||
return delta_timestamps
|
||||
|
||||
|
||||
def _resolve_episodes(
|
||||
episodes: list[int] | None, exclude_episodes: list[int] | None, total_episodes: int
|
||||
) -> list[int] | None:
|
||||
"""Apply an episode exclusion list on top of an optional allowlist."""
|
||||
if not exclude_episodes:
|
||||
return episodes
|
||||
base = episodes if episodes is not None else list(range(total_episodes))
|
||||
excluded = set(exclude_episodes)
|
||||
return [episode for episode in base if episode not in excluded]
|
||||
|
||||
|
||||
def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDataset:
|
||||
"""Handles the logic of setting up delta timestamps and image transforms before creating a dataset.
|
||||
|
||||
@@ -98,14 +87,11 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
|
||||
cfg.dataset.repo_id, root=cfg.dataset.root, revision=cfg.dataset.revision
|
||||
)
|
||||
delta_timestamps = resolve_delta_timestamps(cfg.trainable_config, ds_meta)
|
||||
episodes = _resolve_episodes(
|
||||
cfg.dataset.episodes, cfg.dataset.exclude_episodes, ds_meta.total_episodes
|
||||
)
|
||||
if not cfg.dataset.streaming:
|
||||
dataset = LeRobotDataset(
|
||||
cfg.dataset.repo_id,
|
||||
root=cfg.dataset.root,
|
||||
episodes=episodes,
|
||||
episodes=cfg.dataset.episodes,
|
||||
delta_timestamps=delta_timestamps,
|
||||
image_transforms=image_transforms,
|
||||
revision=cfg.dataset.revision,
|
||||
@@ -118,7 +104,7 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
|
||||
dataset = StreamingLeRobotDataset(
|
||||
cfg.dataset.repo_id,
|
||||
root=cfg.dataset.root,
|
||||
episodes=episodes,
|
||||
episodes=cfg.dataset.episodes,
|
||||
delta_timestamps=delta_timestamps,
|
||||
image_transforms=image_transforms,
|
||||
revision=cfg.dataset.revision,
|
||||
|
||||
@@ -162,28 +162,14 @@ def render_sample(
|
||||
task: str | None = None,
|
||||
dataset_ctx: Any | None = None,
|
||||
) -> RenderedMessages | None:
|
||||
"""Resolve one sample's bindings and render its message recipe.
|
||||
"""Render the chat-style messages for a single dataset sample.
|
||||
|
||||
Returns ``None`` when no text or low-level action supervision applies.
|
||||
Resolves the recipe's bindings against ``persistent`` and ``events`` rows
|
||||
at frame timestamp ``t``, then expands the recipe's message templates.
|
||||
Returns ``None`` if the resolved sample contains no target message.
|
||||
"""
|
||||
persistent_rows = _normalize_rows(persistent or [])
|
||||
event_rows = _normalize_rows(events or [])
|
||||
|
||||
# Route sparse VQA frames to a matching view-specific component before weighted selection.
|
||||
# This avoids dropping annotated frames or selecting VQA without annotations.
|
||||
if recipe.blend is not None:
|
||||
vqa_rendered = _render_vqa_if_present(
|
||||
recipe,
|
||||
persistent=persistent_rows,
|
||||
events=event_rows,
|
||||
t=t,
|
||||
sample_idx=sample_idx,
|
||||
task=task,
|
||||
dataset_ctx=dataset_ctx,
|
||||
)
|
||||
if vqa_rendered is not None:
|
||||
return vqa_rendered
|
||||
|
||||
selected_recipe = _select_recipe(recipe, sample_idx)
|
||||
bindings = _resolve_bindings(
|
||||
selected_recipe,
|
||||
@@ -197,55 +183,6 @@ def render_sample(
|
||||
return _render_message_recipe(selected_recipe, bindings)
|
||||
|
||||
|
||||
def _render_vqa_if_present(
|
||||
recipe: TrainingRecipe,
|
||||
*,
|
||||
persistent: Sequence[LanguageRow],
|
||||
events: Sequence[LanguageRow],
|
||||
t: float,
|
||||
sample_idx: int,
|
||||
task: str | None,
|
||||
dataset_ctx: Any | None,
|
||||
) -> RenderedMessages | None:
|
||||
"""Render a matching VQA component, or return ``None`` for normal selection.
|
||||
|
||||
Multiple matching views are selected deterministically by relative weight.
|
||||
"""
|
||||
assert recipe.blend is not None
|
||||
renderable: list[tuple[float, RenderedMessages]] = []
|
||||
for name, component in recipe.blend.items():
|
||||
if not name.startswith("ask_vqa"):
|
||||
continue
|
||||
bindings = _resolve_bindings(
|
||||
component,
|
||||
persistent=persistent,
|
||||
events=events,
|
||||
t=t,
|
||||
sample_idx=sample_idx,
|
||||
task=task,
|
||||
dataset_ctx=dataset_ctx,
|
||||
)
|
||||
rendered = _render_message_recipe(component, bindings)
|
||||
if rendered is not None:
|
||||
renderable.append((float(component.weight or 0.0), rendered))
|
||||
|
||||
if not renderable:
|
||||
return None
|
||||
if len(renderable) == 1:
|
||||
return renderable[0][1]
|
||||
|
||||
# Choose among matching cameras by relative weight, or uniformly when all weights are zero.
|
||||
total = sum(w for w, _ in renderable) or float(len(renderable))
|
||||
digest = hashlib.blake2b(f"vqa:{sample_idx}".encode(), digest_size=8).digest()
|
||||
draw = int.from_bytes(digest, "big") / 2**64 * total
|
||||
cumulative = 0.0
|
||||
for w, rendered in renderable:
|
||||
cumulative += w or (total / len(renderable))
|
||||
if draw < cumulative:
|
||||
return rendered
|
||||
return renderable[-1][1]
|
||||
|
||||
|
||||
def _select_recipe(recipe: TrainingRecipe, sample_idx: int) -> TrainingRecipe:
|
||||
"""Pick a deterministic blend component for ``sample_idx`` (or return ``recipe``)."""
|
||||
if recipe.blend is None:
|
||||
@@ -409,9 +346,7 @@ def _render_message_recipe(
|
||||
if turn.target:
|
||||
target_indices.append(message_idx)
|
||||
|
||||
# Keep samples with either text targets or low-level action supervision.
|
||||
has_low_level = any(stream == "low_level" for stream in streams)
|
||||
if not target_indices and not has_low_level:
|
||||
if not target_indices:
|
||||
return None
|
||||
|
||||
rendered = {
|
||||
@@ -468,12 +403,14 @@ def _validate_rendered(rendered: RenderedMessages) -> None:
|
||||
|
||||
if len(streams) != len(messages):
|
||||
raise ValueError("message_streams must be aligned with messages.")
|
||||
# Require text or low-level action supervision.
|
||||
if not target_indices and not any(s == "low_level" for s in streams):
|
||||
raise ValueError("Rendered samples must contain a target message or a low_level-stream message.")
|
||||
if not target_indices:
|
||||
raise ValueError("Rendered samples must contain at least one target message.")
|
||||
for idx in target_indices:
|
||||
if idx < 0 or idx >= len(messages):
|
||||
raise ValueError(f"Target message index {idx} is out of bounds.")
|
||||
# ``stream`` is enforced non-None at MessageTurn construction time
|
||||
# (see ``MessageTurn.__post_init__``), so a missing stream here would
|
||||
# mean the dataclass invariant was bypassed; no need to re-check.
|
||||
|
||||
|
||||
def _nth_relative(
|
||||
|
||||
@@ -556,13 +556,7 @@ class RoboCasaEnv(EnvConfig):
|
||||
kwargs["split"] = self.split
|
||||
return kwargs
|
||||
|
||||
def create_envs(
|
||||
self,
|
||||
n_envs: int,
|
||||
use_async_envs: bool = False,
|
||||
terminate_on_success: bool = True,
|
||||
horizon: int | None = None,
|
||||
):
|
||||
def create_envs(self, n_envs: int, use_async_envs: bool = False):
|
||||
from .robocasa import create_robocasa_envs
|
||||
|
||||
if self.task is None:
|
||||
@@ -576,8 +570,6 @@ class RoboCasaEnv(EnvConfig):
|
||||
env_cls=env_cls,
|
||||
episode_length=self.episode_length,
|
||||
obj_registries=tuple(self.obj_registries),
|
||||
terminate_on_success=terminate_on_success,
|
||||
horizon=horizon,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -33,8 +33,8 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
# Dimensions for the flat action/state vectors used by the LeRobot wrapper.
|
||||
# These correspond to the PandaOmron robot in RoboCasa365.
|
||||
OBS_STATE_DIM = 16 # ee_pos_rel(3) + ee_quat_rel(4) + base_pos(3) + base_quat(4) + gripper_qpos(2)
|
||||
ACTION_DIM = 12 # ee_pos(3) + ee_rot(3) + gripper(1) + base_motion(4) + control_mode(1)
|
||||
OBS_STATE_DIM = 16 # base_pos(3) + base_quat(4) + ee_pos_rel(3) + ee_quat_rel(4) + gripper_qpos(2)
|
||||
ACTION_DIM = 12 # base_motion(4) + control_mode(1) + ee_pos(3) + ee_rot(3) + gripper(1)
|
||||
ACTION_LOW = -1.0
|
||||
ACTION_HIGH = 1.0
|
||||
|
||||
@@ -101,15 +101,14 @@ def _resolve_tasks(task: str) -> tuple[list[str], str | None]:
|
||||
def convert_action(flat_action: np.ndarray) -> dict[str, Any]:
|
||||
"""Split a flat (12,) action vector into a RoboCasa action dict.
|
||||
|
||||
Layout (openpi / robocasa.utils.env_utils.convert_action order):
|
||||
ee_pos(3) + ee_rot(3) + gripper(1) + base_motion(4) + control_mode(1)
|
||||
Layout: base_motion(4) + control_mode(1) + ee_pos(3) + ee_rot(3) + gripper(1)
|
||||
"""
|
||||
return {
|
||||
"action.end_effector_position": flat_action[0:3],
|
||||
"action.end_effector_rotation": flat_action[3:6],
|
||||
"action.gripper_close": flat_action[6:7],
|
||||
"action.base_motion": flat_action[7:11],
|
||||
"action.control_mode": flat_action[11:12],
|
||||
"action.base_motion": flat_action[0:4],
|
||||
"action.control_mode": flat_action[4:5],
|
||||
"action.end_effector_position": flat_action[5:8],
|
||||
"action.end_effector_rotation": flat_action[8:11],
|
||||
"action.gripper_close": flat_action[11:12],
|
||||
}
|
||||
|
||||
|
||||
@@ -137,16 +136,9 @@ class RoboCasaEnv(gym.Env):
|
||||
episode_length: int | None = None,
|
||||
obj_registries: Sequence[str] = DEFAULT_OBJ_REGISTRIES,
|
||||
episode_index: int = 0,
|
||||
terminate_on_success: bool = True,
|
||||
horizon: int | None = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.task = task
|
||||
# When False, a task-success does NOT end/reset the episode — used by the
|
||||
# interactive sim so one kitchen persists across sequential prompts.
|
||||
self.terminate_on_success = terminate_on_success
|
||||
# Underlying robosuite horizon (steps before truncation). None -> default.
|
||||
self.horizon = horizon
|
||||
self.obs_type = obs_type
|
||||
self.render_mode = render_mode
|
||||
self.observation_width = observation_width
|
||||
@@ -218,16 +210,12 @@ class RoboCasaEnv(gym.Env):
|
||||
# (only None/"all"/"pretrain"/"target" are valid). Always pass a
|
||||
# valid value so we don't hit that default. Extra kwargs are
|
||||
# forwarded to the underlying kitchen env via create_env/robosuite.make.
|
||||
extra_kwargs: dict[str, Any] = {}
|
||||
if self.horizon is not None:
|
||||
extra_kwargs["horizon"] = int(self.horizon)
|
||||
self._env = RoboCasaGymEnv(
|
||||
env_name=self.task,
|
||||
camera_widths=self.observation_width,
|
||||
camera_heights=self.observation_height,
|
||||
split=self.split if self.split is not None else "all",
|
||||
obj_registries=self.obj_registries,
|
||||
**extra_kwargs,
|
||||
)
|
||||
|
||||
ep_meta = self._env.env.get_ep_meta()
|
||||
@@ -242,14 +230,12 @@ class RoboCasaEnv(gym.Env):
|
||||
return {"pixels": images}
|
||||
|
||||
# `state.*` keys come from PandaOmronKeyConverter inside the wrapper.
|
||||
# openpi state order: ee first, then base, then gripper (matches the
|
||||
# openpi robocasa pipeline / examples/robocasa/main.py state layout).
|
||||
agent_pos = np.concatenate(
|
||||
[
|
||||
raw_obs.get("state.end_effector_position_relative", np.zeros(3)),
|
||||
raw_obs.get("state.end_effector_rotation_relative", np.zeros(4)),
|
||||
raw_obs.get("state.base_position", np.zeros(3)),
|
||||
raw_obs.get("state.base_rotation", np.zeros(4)),
|
||||
raw_obs.get("state.end_effector_position_relative", np.zeros(3)),
|
||||
raw_obs.get("state.end_effector_rotation_relative", np.zeros(4)),
|
||||
raw_obs.get("state.gripper_qpos", np.zeros(2)),
|
||||
],
|
||||
axis=-1,
|
||||
@@ -294,7 +280,7 @@ class RoboCasaEnv(gym.Env):
|
||||
raw_obs, reward, done, truncated, info = self._env.step(action_dict)
|
||||
|
||||
is_success = bool(info.get("success", False))
|
||||
terminated = done or (is_success and self.terminate_on_success)
|
||||
terminated = done or is_success
|
||||
info.update({"task": self.task, "done": done, "is_success": is_success})
|
||||
|
||||
observation = self._format_raw_obs(raw_obs)
|
||||
@@ -327,8 +313,6 @@ def _make_env_fns(
|
||||
split: str | None,
|
||||
episode_length: int | None,
|
||||
obj_registries: Sequence[str],
|
||||
terminate_on_success: bool = True,
|
||||
horizon: int | None = None,
|
||||
) -> list[Callable[[], RoboCasaEnv]]:
|
||||
"""Build n_envs factory callables for a single task.
|
||||
|
||||
@@ -351,8 +335,6 @@ def _make_env_fns(
|
||||
episode_length=episode_length,
|
||||
obj_registries=obj_registries,
|
||||
episode_index=episode_index,
|
||||
terminate_on_success=terminate_on_success,
|
||||
horizon=horizon,
|
||||
)
|
||||
|
||||
return [partial(_make_env, i) for i in range(n_envs)]
|
||||
@@ -366,8 +348,6 @@ def create_robocasa_envs(
|
||||
env_cls: Callable[[Sequence[Callable[[], Any]]], Any] | None = None,
|
||||
episode_length: int | None = None,
|
||||
obj_registries: Sequence[str] = DEFAULT_OBJ_REGISTRIES,
|
||||
terminate_on_success: bool = True,
|
||||
horizon: int | None = None,
|
||||
) -> dict[str, dict[int, Any]]:
|
||||
"""Create vectorized RoboCasa365 environments with a consistent return shape.
|
||||
|
||||
@@ -429,8 +409,6 @@ def create_robocasa_envs(
|
||||
split=split,
|
||||
episode_length=episode_length,
|
||||
obj_registries=obj_registries,
|
||||
terminate_on_success=terminate_on_success,
|
||||
horizon=horizon,
|
||||
)
|
||||
|
||||
if is_async:
|
||||
|
||||
@@ -104,8 +104,6 @@ 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)
|
||||
|
||||
@@ -28,11 +28,11 @@ 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
|
||||
from .utils import make_robot_action, prepare_observation_for_inference
|
||||
from .vla_jepa.configuration_vla_jepa import VLAJEPAConfig as VLAJEPAConfig
|
||||
from .vqbet.configuration_vqbet import VQBeTConfig as VQBeTConfig
|
||||
from .wall_x.configuration_wall_x import WallXConfig as WallXConfig
|
||||
from .xvla.configuration_xvla import XVLAConfig as XVLAConfig
|
||||
@@ -56,9 +56,9 @@ __all__ = [
|
||||
"PI0Config",
|
||||
"PI0FastConfig",
|
||||
"PI05Config",
|
||||
"PI052Config",
|
||||
"SmolVLAConfig",
|
||||
"TDMPCConfig",
|
||||
"VLAJEPAConfig",
|
||||
"VQBeTConfig",
|
||||
"WallXConfig",
|
||||
"XVLAConfig",
|
||||
|
||||
@@ -18,17 +18,10 @@ from typing import Any
|
||||
import torch
|
||||
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
RenameObservationsProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
make_default_pre_post_processors,
|
||||
)
|
||||
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
|
||||
|
||||
from .configuration_act import ACTConfig
|
||||
|
||||
@@ -54,34 +47,4 @@ def make_act_pre_post_processors(
|
||||
tuple[PolicyProcessorPipeline[dict[str, Any], dict[str, Any]], PolicyProcessorPipeline[PolicyAction, PolicyAction]]: A tuple containing the
|
||||
pre-processor pipeline and the post-processor pipeline.
|
||||
"""
|
||||
|
||||
input_steps = [
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
device=config.device,
|
||||
),
|
||||
]
|
||||
output_steps = [
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
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,
|
||||
),
|
||||
)
|
||||
return make_default_pre_post_processors(config, dataset_stats, normalizer_device=config.device)
|
||||
|
||||
@@ -0,0 +1,161 @@
|
||||
#!/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.
|
||||
|
||||
"""Flow-matching sampling primitives shared across policies.
|
||||
|
||||
Canonical versions of the beta-distributed timestep sampler and the forward-Euler
|
||||
denoising loop (with its real-time-chunking hook) that the openpi-derived policies
|
||||
(pi0, pi05, smolvla, eo1) historically each carried a copy of. All functions are
|
||||
stateless; adopting them does not affect checkpoints.
|
||||
"""
|
||||
|
||||
from collections.abc import Callable
|
||||
from functools import lru_cache
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from lerobot.policies.rtc.modeling_rtc import RTCProcessor
|
||||
|
||||
|
||||
@lru_cache(maxsize=None)
|
||||
def _beta_distribution(alpha: float, beta: float) -> "torch.distributions.Beta":
|
||||
# Beta sampling uses _sample_dirichlet which isn't implemented for MPS, so build on CPU.
|
||||
# Cached (groot convention) so the distribution object is constructed once per (alpha, beta).
|
||||
alpha_t = torch.tensor(alpha, dtype=torch.float32)
|
||||
beta_t = torch.tensor(beta, dtype=torch.float32)
|
||||
return torch.distributions.Beta(alpha_t, beta_t)
|
||||
|
||||
|
||||
def sample_beta(alpha: float, beta: float, bsize: int, device) -> Tensor: # see openpi (exact copy)
|
||||
return _beta_distribution(alpha, beta).sample((bsize,)).to(device)
|
||||
|
||||
|
||||
def sample_noise(shape, device, *, distribution: str = "normal") -> Tensor:
|
||||
"""Float32 flow-matching noise sample.
|
||||
|
||||
``distribution="normal"`` (default, openpi: pi0/pi05/eo1/smolvla/groot/wall_x) draws
|
||||
standard-normal noise. ``distribution="uniform"`` (evo1) draws uniformly from
|
||||
``[-1, 1)`` via ``rand * 2 - 1``.
|
||||
"""
|
||||
if distribution == "normal":
|
||||
return torch.normal(
|
||||
mean=0.0,
|
||||
std=1.0,
|
||||
size=shape,
|
||||
dtype=torch.float32,
|
||||
device=device,
|
||||
)
|
||||
if distribution == "uniform":
|
||||
return torch.rand(shape, dtype=torch.float32, device=device) * 2 - 1
|
||||
raise ValueError(f"Unknown noise distribution: {distribution!r} (expected 'normal' or 'uniform')")
|
||||
|
||||
|
||||
def sample_time_beta(
|
||||
bsize: int,
|
||||
device,
|
||||
*,
|
||||
alpha: float,
|
||||
beta: float,
|
||||
scale: float = 1.0,
|
||||
offset: float = 0.0,
|
||||
complement: bool = False,
|
||||
clamp_min: float | None = None,
|
||||
clamp_max: float | None = None,
|
||||
) -> Tensor:
|
||||
"""Beta-distributed flow-matching timesteps.
|
||||
|
||||
Computes ``t = f(Beta(alpha, beta)) * scale + offset`` where ``f`` is the identity by
|
||||
default or ``1 - x`` when ``complement=True``, then optionally clamps to
|
||||
``[clamp_min, clamp_max]``. This covers the known per-policy conventions:
|
||||
|
||||
* openpi backward (pi0/pi05/eo1/smolvla): ``scale=0.999, offset=0.001``.
|
||||
* forward (groot/wall_x): ``complement=True, scale=0.999`` giving ``(1 - beta) * 0.999``.
|
||||
* evo1: ``alpha=beta=2, clamp_min=0.02, clamp_max=0.98``.
|
||||
"""
|
||||
time_beta = sample_beta(alpha, beta, bsize, device)
|
||||
if complement:
|
||||
time_beta = 1.0 - time_beta
|
||||
time = time_beta * scale + offset
|
||||
if clamp_min is not None or clamp_max is not None:
|
||||
time = time.clamp(min=clamp_min, max=clamp_max)
|
||||
return time.to(dtype=torch.float32, device=device)
|
||||
|
||||
|
||||
def euler_integrate(
|
||||
denoise_fn: Callable[[Tensor, Tensor], Tensor],
|
||||
noise: Tensor,
|
||||
num_steps: int,
|
||||
*,
|
||||
rtc_processor: "RTCProcessor | None" = None,
|
||||
rtc_enabled: bool = False,
|
||||
inference_delay: int | None = None,
|
||||
prev_chunk_left_over: Tensor | None = None,
|
||||
execution_horizon: int | None = None,
|
||||
) -> Tensor:
|
||||
"""Forward-Euler integration of a velocity field from t=1 (noise) to t=0 (actions).
|
||||
|
||||
This is the openpi sampling loop: ``dt = -1/num_steps``, ``time = 1.0 + step*dt``,
|
||||
``x_t <- x_t + dt * v_t``, with the optional real-time-chunking (RTC) guidance hook
|
||||
wrapping the velocity computation and debug tracking after each step.
|
||||
|
||||
Args:
|
||||
denoise_fn: Computes the velocity ``v_t`` from ``(x_t, time_tensor)`` where
|
||||
``time_tensor`` is a float32 tensor of shape ``(batch_size,)``. The returned
|
||||
velocity must have the same shape and dtype as ``x_t``.
|
||||
noise: Initial sample ``x_1`` of shape ``(batch_size, ...)``.
|
||||
num_steps: Number of Euler steps.
|
||||
rtc_processor: Optional RTC processor. Debug tracking fires whenever it is set and
|
||||
has debugging enabled, even if RTC guidance itself is disabled (this mirrors
|
||||
the historical per-policy loops).
|
||||
rtc_enabled: Whether to route the velocity computation through
|
||||
``rtc_processor.denoise_step`` (requires ``rtc_processor``).
|
||||
inference_delay: RTC guidance parameter, forwarded verbatim.
|
||||
prev_chunk_left_over: RTC guidance parameter, forwarded verbatim.
|
||||
execution_horizon: RTC guidance parameter, forwarded verbatim.
|
||||
"""
|
||||
bsize = noise.shape[0]
|
||||
device = noise.device
|
||||
|
||||
dt = -1.0 / num_steps
|
||||
x_t = noise
|
||||
for step in range(num_steps):
|
||||
time = 1.0 + step * dt
|
||||
time_tensor = torch.tensor(time, dtype=torch.float32, device=device).expand(bsize)
|
||||
|
||||
def denoise_step_partial_call(input_x_t, current_timestep=time_tensor):
|
||||
return denoise_fn(input_x_t, current_timestep)
|
||||
|
||||
if rtc_enabled:
|
||||
v_t = 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 rtc_processor is not None and rtc_processor.is_debug_enabled():
|
||||
rtc_processor.track(time=time, x_t=x_t, v_t=v_t)
|
||||
|
||||
return x_t
|
||||
@@ -0,0 +1,243 @@
|
||||
#!/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.
|
||||
|
||||
"""Helpers shared by the openpi-derived VLA policies (pi0, pi05, pi0_fast, smolvla, eo1, xvla).
|
||||
|
||||
These are the canonical versions of functions that historically were copy-pasted per
|
||||
policy. They are pure (no parameters, no module state), so importing them from here
|
||||
instead of a policy-local copy has no effect on checkpoints.
|
||||
"""
|
||||
|
||||
import math
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F # noqa: N812
|
||||
from torch import Tensor
|
||||
|
||||
from lerobot.utils.constants import OPENPI_ATTENTION_MASK_VALUE
|
||||
from lerobot.utils.device_utils import get_safe_dtype
|
||||
from lerobot.utils.import_utils import _transformers_available, require_package
|
||||
|
||||
if TYPE_CHECKING or _transformers_available:
|
||||
from transformers import DynamicCache
|
||||
else:
|
||||
DynamicCache = None
|
||||
|
||||
|
||||
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 make_att_2d_masks(pad_masks: Tensor, att_masks: Tensor) -> Tensor: # see openpi (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 prepare_attention_masks_4d(att_2d_masks: Tensor, dtype: torch.dtype | None = None) -> Tensor:
|
||||
"""Expand boolean 2D attention masks to the additive 4D layout expected by transformers.
|
||||
|
||||
Valid positions become 0.0 and masked positions the large negative openpi constant.
|
||||
"""
|
||||
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 clone_past_key_values(past_key_values):
|
||||
"""Clone the DynamicCache returned by prefix prefill for compiled denoising."""
|
||||
if DynamicCache is None:
|
||||
require_package("transformers", extra="transformers-dep")
|
||||
|
||||
return DynamicCache(
|
||||
tuple(
|
||||
(keys.clone(), values.clone(), sliding_window) for keys, values, sliding_window in past_key_values
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def pad_vector(vector: Tensor, new_dim: int, *, truncate: bool = False) -> Tensor:
|
||||
"""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)
|
||||
|
||||
With ``truncate=False`` (openpi behavior), vectors whose last dimension is already
|
||||
>= new_dim are returned unchanged. With ``truncate=True`` (xVLA behavior), the last
|
||||
dimension is truncated to exactly ``new_dim`` (which may be 0).
|
||||
"""
|
||||
if vector.shape[-1] == new_dim:
|
||||
return vector
|
||||
if not truncate:
|
||||
if vector.shape[-1] >= new_dim:
|
||||
return vector
|
||||
return F.pad(vector, (0, new_dim - vector.shape[-1]))
|
||||
shape = list(vector.shape)
|
||||
current_dim = shape[-1]
|
||||
shape[-1] = new_dim
|
||||
new_vector = vector.new_zeros(*shape)
|
||||
length = min(current_dim, new_dim)
|
||||
new_vector[..., :length] = vector[..., :length]
|
||||
return new_vector
|
||||
|
||||
|
||||
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].
|
||||
|
||||
Padding is centered (openpi convention). For the top-left-padding variant used by
|
||||
smolvla/xvla, see :func:`resize_with_pad`.
|
||||
|
||||
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
|
||||
|
||||
|
||||
def resize_with_pad(img: torch.Tensor, height: int, width: int, *, pad_value: float) -> torch.Tensor:
|
||||
"""Resize a (b, c, h, w) image without distortion, padding on the LEFT and TOP.
|
||||
|
||||
This is the smolvla/xvla convention. For the centered-padding openpi variant, see
|
||||
:func:`resize_with_pad_torch`. ``pad_value`` is keyword-only on purpose: callers
|
||||
historically used different values (0, -1) and must state their choice explicitly.
|
||||
"""
|
||||
if img.ndim != 4:
|
||||
raise ValueError(f"(b,c,h,w) expected, but got {img.shape}")
|
||||
|
||||
current_height, current_width = img.shape[2:]
|
||||
if current_height == height and current_width == width:
|
||||
return img
|
||||
|
||||
ratio = max(current_width / width, current_height / height)
|
||||
resized_height = int(current_height / ratio)
|
||||
resized_width = int(current_width / ratio)
|
||||
resized_img = F.interpolate(
|
||||
img, size=(resized_height, resized_width), mode="bilinear", align_corners=False
|
||||
)
|
||||
|
||||
pad_height = max(0, height - resized_height)
|
||||
pad_width = max(0, width - resized_width)
|
||||
padded_img = F.pad(resized_img, (pad_width, 0, pad_height, 0), value=pad_value)
|
||||
return padded_img
|
||||
@@ -19,17 +19,10 @@ from typing import Any
|
||||
import torch
|
||||
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
RenameObservationsProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
make_default_pre_post_processors,
|
||||
)
|
||||
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
|
||||
|
||||
from .configuration_diffusion import DiffusionConfig
|
||||
|
||||
@@ -63,32 +56,4 @@ def make_diffusion_pre_post_processors(
|
||||
Returns:
|
||||
A tuple containing the configured pre-processor and post-processor pipelines.
|
||||
"""
|
||||
|
||||
input_steps = [
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
]
|
||||
output_steps = [
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
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,
|
||||
),
|
||||
)
|
||||
return make_default_pre_post_processors(config, dataset_stats)
|
||||
|
||||
@@ -23,24 +23,16 @@ import torch
|
||||
|
||||
from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
ComplementaryDataProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
ProcessorStep,
|
||||
ProcessorStepRegistry,
|
||||
RenameObservationsProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
make_default_policy_processor_steps,
|
||||
make_policy_processor_pipelines,
|
||||
)
|
||||
from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
|
||||
from lerobot.types import TransitionKey
|
||||
from lerobot.utils.constants import (
|
||||
OBS_STATE,
|
||||
POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
)
|
||||
from lerobot.utils.constants import OBS_STATE
|
||||
from lerobot.utils.import_utils import _transformers_available, require_package
|
||||
|
||||
from .configuration_eo1 import EO1Config
|
||||
@@ -242,14 +234,12 @@ def make_eo1_pre_post_processors(
|
||||
]:
|
||||
"""Build pre/post processor pipelines for EO1."""
|
||||
|
||||
steps = make_default_policy_processor_steps(config, dataset_stats)
|
||||
|
||||
input_steps: list[ProcessorStep] = [
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
steps.rename_observations,
|
||||
steps.add_batch_dim,
|
||||
steps.normalize,
|
||||
EO1ConversationTemplateStep(input_features=config.input_features, chunk_size=config.chunk_size),
|
||||
EO1QwenProcessorStep(
|
||||
processor_name=config.vlm_base,
|
||||
@@ -257,27 +247,12 @@ def make_eo1_pre_post_processors(
|
||||
image_max_pixels=config.image_max_pixels,
|
||||
use_fast_processor=config.use_fast_processor,
|
||||
),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
steps.to_device,
|
||||
]
|
||||
|
||||
output_steps: list[ProcessorStep] = [
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features,
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
steps.unnormalize,
|
||||
steps.to_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,
|
||||
),
|
||||
)
|
||||
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
|
||||
|
||||
+66
-363
@@ -17,6 +17,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import inspect
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Any, TypedDict, Unpack
|
||||
|
||||
@@ -44,26 +45,10 @@ from lerobot.utils.constants import (
|
||||
)
|
||||
from lerobot.utils.feature_utils import dataset_to_policy_features
|
||||
|
||||
from .act.configuration_act import ACTConfig
|
||||
from .diffusion.configuration_diffusion import DiffusionConfig
|
||||
from .eo1.configuration_eo1 import EO1Config
|
||||
from .evo1.configuration_evo1 import Evo1Config
|
||||
from .fastwam.configuration_fastwam import FastWAMConfig
|
||||
from .gaussian_actor.configuration_gaussian_actor import GaussianActorConfig
|
||||
from .groot.configuration_groot import GrootConfig
|
||||
from .lingbot_va.configuration_lingbot_va import LingBotVAConfig
|
||||
from .molmoact2.configuration_molmoact2 import MolmoAct2Config
|
||||
from .multi_task_dit.configuration_multi_task_dit import MultiTaskDiTConfig
|
||||
from .pi0.configuration_pi0 import PI0Config
|
||||
from .pi05.configuration_pi05 import PI05Config
|
||||
from .pretrained import PreTrainedPolicy
|
||||
from .smolvla.configuration_smolvla import SmolVLAConfig
|
||||
from .tdmpc.configuration_tdmpc import TDMPCConfig
|
||||
from .utils import validate_visual_features_consistency
|
||||
from .vla_jepa.configuration_vla_jepa import VLAJEPAConfig
|
||||
from .vqbet.configuration_vqbet import VQBeTConfig
|
||||
from .wall_x.configuration_wall_x import WallXConfig
|
||||
from .xvla.configuration_xvla import XVLAConfig
|
||||
|
||||
|
||||
def _reconnect_relative_absolute_steps(
|
||||
@@ -88,104 +73,23 @@ def get_policy_class(name: str) -> type[PreTrainedPolicy]:
|
||||
"""
|
||||
Retrieves a policy class by its registered name.
|
||||
|
||||
This function uses dynamic imports to avoid loading all policy classes into memory
|
||||
at once, improving startup time and reducing dependencies.
|
||||
Resolution is convention-based: the draccus-registered config class of ``name`` is
|
||||
looked up, its ``configuration_*`` module path is rewritten to ``modeling_*``, and
|
||||
the ``<X>Policy`` class is imported from there. The modeling module is only imported
|
||||
at call time, keeping heavy optional dependencies lazy. This works for both built-in
|
||||
policies and third-party lerobot plugins (anything registered via
|
||||
``@PreTrainedConfig.register_subclass``).
|
||||
|
||||
Args:
|
||||
name: The name of the policy. Supported names are "tdmpc", "diffusion", "act",
|
||||
"multi_task_dit", "vqbet", "pi0", "pi05", "gaussian_actor", "smolvla", "wall_x",
|
||||
"molmoact2", "eo1", "evo1".
|
||||
name: The registered name of the policy (e.g. "act", "diffusion", "pi0").
|
||||
Returns:
|
||||
The policy class corresponding to the given name.
|
||||
|
||||
Raises:
|
||||
NotImplementedError: If the policy name is not recognized.
|
||||
ValueError: If the policy name is not registered.
|
||||
ImportError: If the policy's optional dependencies are not installed.
|
||||
"""
|
||||
if name == "tdmpc":
|
||||
from .tdmpc.modeling_tdmpc import TDMPCPolicy
|
||||
|
||||
return TDMPCPolicy
|
||||
elif name == "diffusion":
|
||||
from .diffusion.modeling_diffusion import DiffusionPolicy
|
||||
|
||||
return DiffusionPolicy
|
||||
elif name == "act":
|
||||
from .act.modeling_act import ACTPolicy
|
||||
|
||||
return ACTPolicy
|
||||
elif name == "multi_task_dit":
|
||||
from .multi_task_dit.modeling_multi_task_dit import MultiTaskDiTPolicy
|
||||
|
||||
return MultiTaskDiTPolicy
|
||||
elif name == "vqbet":
|
||||
from .vqbet.modeling_vqbet import VQBeTPolicy
|
||||
|
||||
return VQBeTPolicy
|
||||
elif name == "pi0":
|
||||
from .pi0.modeling_pi0 import PI0Policy
|
||||
|
||||
return PI0Policy
|
||||
elif name == "pi0_fast":
|
||||
from .pi0_fast.modeling_pi0_fast import PI0FastPolicy
|
||||
|
||||
return PI0FastPolicy
|
||||
elif name == "pi05":
|
||||
from .pi05.modeling_pi05 import PI05Policy
|
||||
|
||||
return PI05Policy
|
||||
elif name == "pi052":
|
||||
from .pi052.modeling_pi052 import PI052Policy
|
||||
|
||||
return PI052Policy
|
||||
elif name == "gaussian_actor":
|
||||
from .gaussian_actor.modeling_gaussian_actor import GaussianActorPolicy
|
||||
|
||||
return GaussianActorPolicy
|
||||
elif name == "smolvla":
|
||||
from .smolvla.modeling_smolvla import SmolVLAPolicy
|
||||
|
||||
return SmolVLAPolicy
|
||||
elif name == "groot":
|
||||
from .groot.modeling_groot import GrootPolicy
|
||||
|
||||
return GrootPolicy
|
||||
elif name == "xvla":
|
||||
from .xvla.modeling_xvla import XVLAPolicy
|
||||
|
||||
return XVLAPolicy
|
||||
elif name == "wall_x":
|
||||
from .wall_x.modeling_wall_x import WallXPolicy
|
||||
|
||||
return WallXPolicy
|
||||
elif name == "eo1":
|
||||
from .eo1.modeling_eo1 import EO1Policy
|
||||
|
||||
return EO1Policy
|
||||
elif name == "molmoact2":
|
||||
from .molmoact2.modeling_molmoact2 import MolmoAct2Policy
|
||||
|
||||
return MolmoAct2Policy
|
||||
elif name == "vla_jepa":
|
||||
from .vla_jepa.modeling_vla_jepa import VLAJEPAPolicy
|
||||
|
||||
return VLAJEPAPolicy
|
||||
elif name == "lingbot_va":
|
||||
from .lingbot_va.modeling_lingbot_va import LingBotVAPolicy
|
||||
|
||||
return LingBotVAPolicy
|
||||
elif name == "fastwam":
|
||||
from .fastwam.modeling_fastwam import FastWAMPolicy
|
||||
|
||||
return FastWAMPolicy
|
||||
elif name == "evo1":
|
||||
from .evo1.modeling_evo1 import Evo1Policy
|
||||
|
||||
return Evo1Policy
|
||||
else:
|
||||
try:
|
||||
return _get_policy_cls_from_policy_name(name=name)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Policy type '{name}' is not available.") from e
|
||||
return _get_policy_cls_from_policy_name(name=name)
|
||||
|
||||
|
||||
def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
|
||||
@@ -196,9 +100,8 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
|
||||
mapping a string identifier to the corresponding config class.
|
||||
|
||||
Args:
|
||||
policy_type: The type of the policy. Supported types include "tdmpc",
|
||||
"multi_task_dit", "diffusion", "act", "vqbet", "pi0", "pi05", "pi052",
|
||||
"gaussian_actor", "smolvla", "wall_x", "molmoact2", "eo1", "evo1".
|
||||
policy_type: The registered type of the policy (any name registered via
|
||||
``@PreTrainedConfig.register_subclass``, e.g. "act", "diffusion", "pi0").
|
||||
**kwargs: Keyword arguments to be passed to the configuration class constructor.
|
||||
|
||||
Returns:
|
||||
@@ -207,52 +110,11 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
|
||||
Raises:
|
||||
ValueError: If the `policy_type` is not recognized.
|
||||
"""
|
||||
if policy_type == "tdmpc":
|
||||
return TDMPCConfig(**kwargs)
|
||||
elif policy_type == "diffusion":
|
||||
return DiffusionConfig(**kwargs)
|
||||
elif policy_type == "act":
|
||||
return ACTConfig(**kwargs)
|
||||
elif policy_type == "multi_task_dit":
|
||||
return MultiTaskDiTConfig(**kwargs)
|
||||
elif policy_type == "vqbet":
|
||||
return VQBeTConfig(**kwargs)
|
||||
elif policy_type == "pi0":
|
||||
return PI0Config(**kwargs)
|
||||
elif policy_type == "pi05":
|
||||
return PI05Config(**kwargs)
|
||||
elif policy_type == "pi052":
|
||||
from .pi052.configuration_pi052 import PI052Config
|
||||
|
||||
return PI052Config(**kwargs)
|
||||
elif policy_type == "gaussian_actor":
|
||||
return GaussianActorConfig(**kwargs)
|
||||
elif policy_type == "smolvla":
|
||||
return SmolVLAConfig(**kwargs)
|
||||
elif policy_type == "groot":
|
||||
return GrootConfig(**kwargs)
|
||||
elif policy_type == "xvla":
|
||||
return XVLAConfig(**kwargs)
|
||||
elif policy_type == "wall_x":
|
||||
return WallXConfig(**kwargs)
|
||||
elif policy_type == "eo1":
|
||||
return EO1Config(**kwargs)
|
||||
elif policy_type == "molmoact2":
|
||||
return MolmoAct2Config(**kwargs)
|
||||
elif policy_type == "vla_jepa":
|
||||
return VLAJEPAConfig(**kwargs)
|
||||
elif policy_type == "lingbot_va":
|
||||
return LingBotVAConfig(**kwargs)
|
||||
elif policy_type == "fastwam":
|
||||
return FastWAMConfig(**kwargs)
|
||||
elif policy_type == "evo1":
|
||||
return Evo1Config(**kwargs)
|
||||
else:
|
||||
try:
|
||||
config_cls = PreTrainedConfig.get_choice_class(policy_type)
|
||||
return config_cls(**kwargs)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Policy type '{policy_type}' is not available.") from e
|
||||
try:
|
||||
config_cls = PreTrainedConfig.get_choice_class(policy_type)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Policy type '{policy_type}' is not available.") from e
|
||||
return config_cls(**kwargs)
|
||||
|
||||
|
||||
class ProcessorConfigKwargs(TypedDict, total=False):
|
||||
@@ -275,12 +137,6 @@ 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 repo used for optional processor fitting; omit it to use universal tokenizers.
|
||||
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
|
||||
|
||||
|
||||
@@ -312,13 +168,9 @@ def make_pre_post_processors(
|
||||
A tuple containing the input (pre-processor) and output (post-processor) pipelines.
|
||||
|
||||
Raises:
|
||||
NotImplementedError: If a processor factory is not implemented for the given
|
||||
policy configuration type.
|
||||
ValueError: If no processor factory exists for the given policy configuration type.
|
||||
"""
|
||||
if pretrained_path:
|
||||
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
|
||||
|
||||
@@ -368,194 +220,13 @@ def make_pre_post_processors(
|
||||
)
|
||||
return preprocessor, postprocessor
|
||||
|
||||
# Create a new processor based on policy type
|
||||
if isinstance(policy_cfg, TDMPCConfig):
|
||||
from .tdmpc.processor_tdmpc import make_tdmpc_pre_post_processors
|
||||
|
||||
processors = make_tdmpc_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, DiffusionConfig):
|
||||
from .diffusion.processor_diffusion import make_diffusion_pre_post_processors
|
||||
|
||||
processors = make_diffusion_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, ACTConfig):
|
||||
from .act.processor_act import make_act_pre_post_processors
|
||||
|
||||
processors = make_act_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, MultiTaskDiTConfig):
|
||||
from .multi_task_dit.processor_multi_task_dit import (
|
||||
make_multi_task_dit_pre_post_processors,
|
||||
)
|
||||
|
||||
processors = make_multi_task_dit_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, VQBeTConfig):
|
||||
from .vqbet.processor_vqbet import make_vqbet_pre_post_processors
|
||||
|
||||
processors = make_vqbet_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, PI0Config):
|
||||
from .pi0.processor_pi0 import make_pi0_pre_post_processors
|
||||
|
||||
processors = make_pi0_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif policy_cfg.type == "pi0_fast":
|
||||
from .pi0_fast.processor_pi0_fast import make_pi0_fast_pre_post_processors
|
||||
|
||||
processors = make_pi0_fast_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
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"),
|
||||
)
|
||||
|
||||
elif policy_cfg.type == "pi052":
|
||||
# PI052 must precede PI05 because its config subclasses PI05Config.
|
||||
from .pi052.processor_pi052 import make_pi052_pre_post_processors
|
||||
|
||||
processors = make_pi052_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
# Without a dataset repo, FAST auto-fit falls back to the universal tokenizer.
|
||||
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"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, PI05Config):
|
||||
from .pi05.processor_pi05 import make_pi05_pre_post_processors
|
||||
|
||||
processors = make_pi05_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, GaussianActorConfig):
|
||||
from .gaussian_actor.processor_gaussian_actor import make_gaussian_actor_pre_post_processors
|
||||
|
||||
processors = make_gaussian_actor_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, SmolVLAConfig):
|
||||
from .smolvla.processor_smolvla import make_smolvla_pre_post_processors
|
||||
|
||||
processors = make_smolvla_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, GrootConfig):
|
||||
from .groot.processor_groot import make_groot_pre_post_processors
|
||||
|
||||
processors = make_groot_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
dataset_meta=kwargs.get("dataset_meta"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, XVLAConfig):
|
||||
from .xvla.processor_xvla import (
|
||||
make_xvla_pre_post_processors,
|
||||
)
|
||||
|
||||
processors = make_xvla_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, WallXConfig):
|
||||
from .wall_x.processor_wall_x import make_wall_x_pre_post_processors
|
||||
|
||||
processors = make_wall_x_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, EO1Config):
|
||||
from .eo1.processor_eo1 import make_eo1_pre_post_processors
|
||||
|
||||
processors = make_eo1_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
elif isinstance(policy_cfg, Evo1Config):
|
||||
from .evo1.processor_evo1 import make_evo1_pre_post_processors
|
||||
|
||||
processors = make_evo1_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, MolmoAct2Config):
|
||||
from .molmoact2.processor_molmoact2 import make_molmoact2_pre_post_processors
|
||||
|
||||
processors = make_molmoact2_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
dataset_meta=kwargs.get("dataset_meta"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, VLAJEPAConfig):
|
||||
from .vla_jepa.processor_vla_jepa import make_vla_jepa_pre_post_processors
|
||||
|
||||
processors = make_vla_jepa_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, LingBotVAConfig):
|
||||
from .lingbot_va.processor_lingbot_va import make_lingbot_va_pre_post_processors
|
||||
|
||||
processors = make_lingbot_va_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, FastWAMConfig):
|
||||
from .fastwam.processor_fastwam import make_fastwam_pre_post_processors
|
||||
|
||||
processors = make_fastwam_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
else:
|
||||
try:
|
||||
processors = _make_processors_from_policy_config(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Processor for policy type '{policy_cfg.type}' is not implemented.") from e
|
||||
|
||||
return processors
|
||||
# Create new processors from the policy config, resolving the per-policy factory
|
||||
# function by naming convention (lazy import keeps optional dependencies optional).
|
||||
return _make_processors_from_policy_config(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
dataset_meta=kwargs.get("dataset_meta"),
|
||||
)
|
||||
|
||||
|
||||
def make_policy(
|
||||
@@ -699,10 +370,12 @@ def make_policy(
|
||||
return policy
|
||||
|
||||
|
||||
def _get_policy_cls_from_policy_name(name: str) -> type[PreTrainedConfig]:
|
||||
def _get_policy_cls_from_policy_name(name: str) -> type[PreTrainedPolicy]:
|
||||
"""Get policy class from its registered name using dynamic imports.
|
||||
|
||||
This is used as a helper function to import policies from 3rd party lerobot plugins.
|
||||
Works for built-in policies and 3rd party lerobot plugins alike: the config class
|
||||
registered under ``name`` is resolved via the draccus ChoiceRegistry, and the policy
|
||||
class is imported from the sibling ``modeling_*`` module by naming convention.
|
||||
|
||||
Args:
|
||||
name: The name of the policy.
|
||||
@@ -728,22 +401,39 @@ def _get_policy_cls_from_policy_name(name: str) -> type[PreTrainedConfig]:
|
||||
"configuration_", "modeling_"
|
||||
) # e.g., configuration_diffusion -> modeling_diffusion
|
||||
|
||||
module = importlib.import_module(module_path)
|
||||
policy_cls = getattr(module, cls_name)
|
||||
try:
|
||||
module = importlib.import_module(module_path)
|
||||
except ModuleNotFoundError as e:
|
||||
if e.name == module_path:
|
||||
# The modeling_* module itself does not exist for this policy type. A missing
|
||||
# optional dependency inside an existing module propagates unchanged instead,
|
||||
# so its actionable install hint stays visible.
|
||||
raise ValueError(f"Policy class for '{name}' is not implemented.") from e
|
||||
raise
|
||||
policy_cls = getattr(module, cls_name, None)
|
||||
if policy_cls is None:
|
||||
raise ValueError(
|
||||
f"Policy class '{cls_name}' not found in '{module_path}'. "
|
||||
f"Policies must expose '<Name>Policy' in the sibling 'modeling_*' module by naming convention."
|
||||
)
|
||||
return policy_cls
|
||||
|
||||
|
||||
def _make_processors_from_policy_config(
|
||||
config: PreTrainedConfig,
|
||||
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
|
||||
dataset_meta: Any | None = None,
|
||||
) -> tuple[Any, Any]:
|
||||
"""Create pre- and post-processors from a policy configuration using dynamic imports.
|
||||
|
||||
This is used as a helper function to import processor factories from 3rd party lerobot plugins.
|
||||
Resolves ``make_{type}_pre_post_processors`` from the policy's ``processor_*`` module
|
||||
by naming convention. Works for built-in policies and 3rd party lerobot plugins.
|
||||
|
||||
Args:
|
||||
config: The policy configuration object.
|
||||
dataset_stats: Dataset statistics for normalization.
|
||||
dataset_meta: Dataset metadata, forwarded only to factories that declare a
|
||||
``dataset_meta`` parameter (e.g. groot, molmoact2).
|
||||
Returns:
|
||||
A tuple containing the input (pre-processor) and output (post-processor) pipelines.
|
||||
"""
|
||||
@@ -756,6 +446,19 @@ def _make_processors_from_policy_config(
|
||||
logging.debug(
|
||||
f"Instantiating pre/post processors using function '{function_name}' from module '{module_path}'"
|
||||
)
|
||||
module = importlib.import_module(module_path)
|
||||
function = getattr(module, function_name)
|
||||
return function(config, dataset_stats=dataset_stats)
|
||||
try:
|
||||
module = importlib.import_module(module_path)
|
||||
except ModuleNotFoundError as e:
|
||||
if e.name == module_path:
|
||||
# The processor_* module itself does not exist for this policy type. A missing
|
||||
# optional dependency inside an existing module propagates unchanged instead,
|
||||
# so its actionable install hint stays visible.
|
||||
raise ValueError(f"Processor for policy type '{policy_type}' is not implemented.") from e
|
||||
raise
|
||||
function = getattr(module, function_name, None)
|
||||
if function is None:
|
||||
raise ValueError(f"Processor for policy type '{policy_type}' is not implemented.")
|
||||
call_kwargs: dict[str, Any] = {"dataset_stats": dataset_stats}
|
||||
if "dataset_meta" in inspect.signature(function).parameters:
|
||||
call_kwargs["dataset_meta"] = dataset_meta
|
||||
return function(config, **call_kwargs)
|
||||
|
||||
@@ -22,20 +22,11 @@ import torch
|
||||
from lerobot.configs import PipelineFeatureType, PolicyFeature
|
||||
from lerobot.processor import (
|
||||
ActionProcessorStep,
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
ProcessorStepRegistry,
|
||||
RenameObservationsProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
)
|
||||
from lerobot.utils.constants import (
|
||||
POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
make_default_policy_processor_steps,
|
||||
make_policy_processor_pipelines,
|
||||
)
|
||||
|
||||
from .configuration_fastwam import FastWAMConfig
|
||||
@@ -105,38 +96,20 @@ def make_fastwam_pre_post_processors(
|
||||
# anyway) and unsafe across fine-tuning: its `resize_size` would be inherited from the base
|
||||
# checkpoint's camera geometry, not this dataset's, making the concatenation N_cameras x too wide.
|
||||
|
||||
steps = make_default_policy_processor_steps(config, normalization_stats, normalizer_device=config.device)
|
||||
|
||||
input_steps = [
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=normalization_stats,
|
||||
device=config.device,
|
||||
),
|
||||
steps.rename_observations,
|
||||
steps.add_batch_dim,
|
||||
steps.to_device,
|
||||
steps.normalize,
|
||||
]
|
||||
output_steps = [
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features,
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=normalization_stats,
|
||||
),
|
||||
steps.unnormalize,
|
||||
]
|
||||
if config.toggle_action_dimensions:
|
||||
output_steps.append(
|
||||
FastWAMActionToggleProcessorStep(toggle_dimensions=config.toggle_action_dimensions)
|
||||
)
|
||||
output_steps.append(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,
|
||||
),
|
||||
)
|
||||
output_steps.append(steps.to_cpu)
|
||||
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
|
||||
|
||||
@@ -20,17 +20,10 @@ from typing import Any
|
||||
import torch
|
||||
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
RenameObservationsProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
make_default_pre_post_processors,
|
||||
)
|
||||
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
|
||||
|
||||
from .configuration_gaussian_actor import GaussianActorConfig
|
||||
|
||||
@@ -62,33 +55,4 @@ def make_gaussian_actor_pre_post_processors(
|
||||
Returns:
|
||||
A tuple containing the configured pre-processor and post-processor pipelines.
|
||||
"""
|
||||
|
||||
# Add remaining processors
|
||||
input_steps = [
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
]
|
||||
output_steps = [
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
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,
|
||||
),
|
||||
)
|
||||
return make_default_pre_post_processors(config, dataset_stats)
|
||||
|
||||
@@ -25,19 +25,12 @@ import torch
|
||||
|
||||
from lerobot.configs.types import FeatureType, NormalizationMode
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
ProcessorStep,
|
||||
RenameObservationsProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
)
|
||||
from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
|
||||
from lerobot.utils.constants import (
|
||||
POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
make_default_policy_processor_steps,
|
||||
make_policy_processor_pipelines,
|
||||
)
|
||||
|
||||
from .configuration_lingbot_va import LingBotVAConfig
|
||||
@@ -52,15 +45,13 @@ def make_lingbot_va_pre_post_processors(
|
||||
]:
|
||||
"""Build the pre/post processor pipelines for LingBot-VA."""
|
||||
|
||||
steps = make_default_policy_processor_steps(config, dataset_stats)
|
||||
|
||||
input_steps: list[ProcessorStep] = [
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
steps.rename_observations,
|
||||
steps.add_batch_dim,
|
||||
steps.normalize,
|
||||
steps.to_device,
|
||||
]
|
||||
|
||||
# Unnormalize actions from [-1, 1] to physical units (QUANTILES) using q01/q99 restored from the checkpoint.
|
||||
@@ -70,18 +61,7 @@ def make_lingbot_va_pre_post_processors(
|
||||
norm_map={FeatureType.ACTION: NormalizationMode.QUANTILES},
|
||||
stats=dataset_stats,
|
||||
),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
steps.to_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,
|
||||
),
|
||||
)
|
||||
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
|
||||
|
||||
@@ -19,18 +19,12 @@ from typing import Any
|
||||
import torch
|
||||
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
RenameObservationsProcessorStep,
|
||||
TokenizerProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
make_default_policy_processor_steps,
|
||||
make_policy_processor_pipelines,
|
||||
)
|
||||
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
|
||||
|
||||
from .configuration_multi_task_dit import MultiTaskDiTConfig
|
||||
|
||||
@@ -66,9 +60,11 @@ def make_multi_task_dit_pre_post_processors(
|
||||
A tuple containing the configured pre-processor and post-processor pipelines.
|
||||
"""
|
||||
|
||||
steps = make_default_policy_processor_steps(config, dataset_stats, normalizer_device=config.device)
|
||||
|
||||
input_steps = [
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
steps.rename_observations,
|
||||
steps.add_batch_dim,
|
||||
TokenizerProcessorStep(
|
||||
tokenizer_name=config.text_encoder_name,
|
||||
padding=config.tokenizer_padding,
|
||||
@@ -76,32 +72,12 @@ def make_multi_task_dit_pre_post_processors(
|
||||
max_length=config.tokenizer_max_length,
|
||||
truncation=config.tokenizer_truncation,
|
||||
),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
device=config.device,
|
||||
),
|
||||
steps.to_device,
|
||||
steps.normalize,
|
||||
]
|
||||
output_steps = [
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features,
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
steps.unnormalize,
|
||||
steps.to_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,
|
||||
),
|
||||
)
|
||||
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
|
||||
|
||||
@@ -21,22 +21,16 @@ import torch
|
||||
from lerobot.configs import PipelineFeatureType, PolicyFeature
|
||||
from lerobot.processor import (
|
||||
AbsoluteActionsProcessorStep,
|
||||
AddBatchDimensionProcessorStep,
|
||||
ComplementaryDataProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
ProcessorStep,
|
||||
ProcessorStepRegistry,
|
||||
RelativeActionsProcessorStep,
|
||||
RenameObservationsProcessorStep,
|
||||
TokenizerProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
make_default_policy_processor_steps,
|
||||
make_policy_processor_pipelines,
|
||||
)
|
||||
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
|
||||
|
||||
from .configuration_pi0 import PI0Config
|
||||
|
||||
@@ -136,10 +130,12 @@ def make_pi0_pre_post_processors(
|
||||
action_names=getattr(config, "action_feature_names", None),
|
||||
)
|
||||
|
||||
steps = make_default_policy_processor_steps(config, dataset_stats)
|
||||
|
||||
# OpenPI order: raw → relative → normalize → model → unnormalize → absolute
|
||||
input_steps: list[ProcessorStep] = [
|
||||
RenameObservationsProcessorStep(rename_map={}), # To mimic the same processor as pretrained one
|
||||
AddBatchDimensionProcessorStep(),
|
||||
steps.rename_observations, # To mimic the same processor as pretrained one
|
||||
steps.add_batch_dim,
|
||||
Pi0NewLineProcessor(), # Add newlines before tokenization for PaliGemma
|
||||
TokenizerProcessorStep(
|
||||
tokenizer_name="google/paligemma-3b-pt-224",
|
||||
@@ -147,32 +143,15 @@ def make_pi0_pre_post_processors(
|
||||
padding_side="right",
|
||||
padding="max_length",
|
||||
),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
steps.to_device,
|
||||
relative_step,
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
steps.normalize,
|
||||
]
|
||||
|
||||
output_steps: list[ProcessorStep] = [
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
steps.unnormalize,
|
||||
AbsoluteActionsProcessorStep(enabled=config.use_relative_actions, relative_step=relative_step),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
steps.to_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,
|
||||
),
|
||||
)
|
||||
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
|
||||
|
||||
@@ -15,7 +15,6 @@
|
||||
# limitations under the License.
|
||||
|
||||
import builtins
|
||||
import json
|
||||
import logging
|
||||
import math
|
||||
from collections import deque
|
||||
@@ -24,7 +23,6 @@ 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
|
||||
@@ -34,7 +32,6 @@ 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,
|
||||
@@ -50,7 +47,6 @@ else:
|
||||
_gated_residual = None
|
||||
layernorm_forward = None
|
||||
PaliGemmaForConditionalGenerationWithPiGemma = None
|
||||
cached_file = None
|
||||
from lerobot.configs import PreTrainedConfig
|
||||
from lerobot.utils.constants import (
|
||||
ACTION,
|
||||
@@ -70,84 +66,6 @@ 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:
|
||||
@@ -645,12 +563,6 @@ 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
|
||||
@@ -694,11 +606,7 @@ 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
|
||||
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.paligemma.model.vision_tower.gradient_checkpointing = True
|
||||
self.paligemma_with_expert.gemma_expert.model.gradient_checkpointing = True
|
||||
logging.info("Enabled gradient checkpointing for PI05Pytorch model")
|
||||
|
||||
@@ -706,11 +614,7 @@ 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
|
||||
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.paligemma.model.vision_tower.gradient_checkpointing = False
|
||||
self.paligemma_with_expert.gemma_expert.model.gradient_checkpointing = False
|
||||
logging.info("Disabled gradient checkpointing for PI05Pytorch model")
|
||||
|
||||
@@ -725,13 +629,10 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch`
|
||||
)
|
||||
return func(*args, **kwargs)
|
||||
|
||||
def _prepare_attention_masks_4d(self, att_2d_masks, dtype=None):
|
||||
def _prepare_attention_masks_4d(self, att_2d_masks):
|
||||
"""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
|
||||
return torch.where(att_2d_masks_4d, 0.0, OPENPI_ATTENTION_MASK_VALUE)
|
||||
|
||||
def sample_noise(self, shape, device):
|
||||
return torch.normal(
|
||||
@@ -757,16 +658,13 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch`
|
||||
pad_masks = []
|
||||
att_masks = []
|
||||
|
||||
if self.checkpoint_vision_embeddings:
|
||||
# Process images
|
||||
for img, img_mask in zip(images, img_masks, strict=True):
|
||||
|
||||
def embed_image(img):
|
||||
return self._apply_checkpoint(self.paligemma_with_expert.embed_image, img)
|
||||
def image_embed_func(img):
|
||||
return self.paligemma_with_expert.embed_image(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):
|
||||
img_emb = self._apply_checkpoint(image_embed_func, img)
|
||||
bsize, num_img_embs = img_emb.shape[:2]
|
||||
|
||||
embs.append(img_emb)
|
||||
@@ -836,14 +734,8 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch`
|
||||
|
||||
embs = torch.cat(embs, dim=1)
|
||||
pad_masks = torch.cat(pad_masks, dim=1)
|
||||
if self.use_on_device_suffix_mask:
|
||||
n = len(att_masks)
|
||||
att_masks = torch.zeros(n, dtype=embs.dtype, device=embs.device)
|
||||
att_masks[0] = 1
|
||||
att_masks = att_masks[None, :].expand(bsize, n)
|
||||
else:
|
||||
att_masks = torch.tensor(att_masks, dtype=embs.dtype, device=embs.device)
|
||||
att_masks = att_masks[None, :].expand(bsize, len(att_masks))
|
||||
att_masks = torch.tensor(att_masks, dtype=embs.dtype, device=embs.device)
|
||||
att_masks = att_masks[None, :].expand(bsize, len(att_masks))
|
||||
|
||||
return embs, pad_masks, att_masks, adarms_cond
|
||||
|
||||
@@ -927,8 +819,7 @@ 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
|
||||
|
||||
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)
|
||||
prefix_att_2d_masks_4d = self._prepare_attention_masks_4d(prefix_att_2d_masks)
|
||||
self.paligemma_with_expert.paligemma.model.language_model.config._attn_implementation = "eager" # noqa: SLF001
|
||||
|
||||
_, past_key_values = self.paligemma_with_expert.forward(
|
||||
@@ -941,19 +832,10 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch`
|
||||
|
||||
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)
|
||||
time_tensor = torch.tensor(time, dtype=torch.float32, device=device).expand(bsize)
|
||||
|
||||
def denoise_step_partial_call(input_x_t, current_timestep=time_tensor):
|
||||
return self.denoise_step(
|
||||
@@ -1031,9 +913,6 @@ class PI05Policy(PreTrainedPolicy):
|
||||
|
||||
config_class = PI05Config
|
||||
name = "pi05"
|
||||
model_class = PI05Pytorch
|
||||
eval_after_pretrained_load = False
|
||||
show_openpi_disclaimer = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@@ -1051,7 +930,7 @@ class PI05Policy(PreTrainedPolicy):
|
||||
|
||||
# Initialize the core PI05 model
|
||||
self.init_rtc_processor()
|
||||
self.model = self.model_class(config, rtc_processor=self.rtc_processor)
|
||||
self.model = PI05Pytorch(config, rtc_processor=self.rtc_processor)
|
||||
|
||||
# Enable gradient checkpointing if requested
|
||||
if config.gradient_checkpointing:
|
||||
@@ -1077,16 +956,16 @@ class PI05Policy(PreTrainedPolicy):
|
||||
strict: bool = True,
|
||||
**kwargs,
|
||||
) -> T:
|
||||
"""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"
|
||||
)
|
||||
"""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"
|
||||
)
|
||||
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,
|
||||
@@ -1100,34 +979,84 @@ class PI05Policy(PreTrainedPolicy):
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# Initialize model without loading weights
|
||||
# Check if dataset_stats were provided in kwargs
|
||||
model = cls(config, **kwargs)
|
||||
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
|
||||
# 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}")
|
||||
|
||||
return model
|
||||
|
||||
def _fix_pytorch_state_dict_keys(
|
||||
self, state_dict, model_config
|
||||
@@ -1299,16 +1228,12 @@ class PI05Policy(PreTrainedPolicy):
|
||||
|
||||
# Action queue logic for n_action_steps > 1
|
||||
if len(self._action_queue) == 0:
|
||||
action_batch = self._prepare_action_batch(batch)
|
||||
actions = self.predict_action_chunk(action_batch)[:, : self.config.n_action_steps]
|
||||
actions = self.predict_action_chunk(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."""
|
||||
|
||||
@@ -24,26 +24,17 @@ import torch
|
||||
from lerobot.configs import PipelineFeatureType, PolicyFeature
|
||||
from lerobot.processor import (
|
||||
AbsoluteActionsProcessorStep,
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
ProcessorStep,
|
||||
ProcessorStepRegistry,
|
||||
RelativeActionsProcessorStep,
|
||||
RenameObservationsProcessorStep,
|
||||
TokenizerProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
make_default_policy_processor_steps,
|
||||
make_policy_processor_pipelines,
|
||||
)
|
||||
from lerobot.types import EnvTransition, TransitionKey
|
||||
from lerobot.utils.constants import (
|
||||
OBS_STATE,
|
||||
POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
)
|
||||
from lerobot.utils.constants import OBS_STATE
|
||||
|
||||
from .configuration_pi05 import PI05Config
|
||||
|
||||
@@ -135,18 +126,16 @@ def make_pi05_pre_post_processors(
|
||||
action_names=getattr(config, "action_feature_names", None),
|
||||
)
|
||||
|
||||
steps = make_default_policy_processor_steps(config, dataset_stats)
|
||||
|
||||
# OpenPI order: raw → relative → normalize → model → unnormalize → absolute
|
||||
input_steps: list[ProcessorStep] = [
|
||||
RenameObservationsProcessorStep(rename_map={}), # To mimic the same processor as pretrained one
|
||||
AddBatchDimensionProcessorStep(),
|
||||
steps.rename_observations, # To mimic the same processor as pretrained one
|
||||
steps.add_batch_dim,
|
||||
relative_step,
|
||||
# NOTE: NormalizerProcessorStep MUST come before Pi05PrepareStateTokenizerProcessorStep
|
||||
# because the tokenizer step expects normalized state in [-1, 1] range for discretization
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
steps.normalize,
|
||||
Pi05PrepareStateTokenizerProcessorStep(max_state_dim=config.max_state_dim),
|
||||
TokenizerProcessorStep(
|
||||
tokenizer_name="google/paligemma-3b-pt-224",
|
||||
@@ -154,26 +143,13 @@ def make_pi05_pre_post_processors(
|
||||
padding_side="right",
|
||||
padding="max_length",
|
||||
),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
steps.to_device,
|
||||
]
|
||||
|
||||
output_steps: list[ProcessorStep] = [
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
steps.unnormalize,
|
||||
AbsoluteActionsProcessorStep(enabled=config.use_relative_actions, relative_step=relative_step),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
steps.to_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,
|
||||
),
|
||||
)
|
||||
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
|
||||
|
||||
@@ -1,19 +0,0 @@
|
||||
# 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"]
|
||||
@@ -1,195 +0,0 @@
|
||||
# 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")
|
||||
@@ -1,512 +0,0 @@
|
||||
# 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
|
||||
]
|
||||
|
||||
return fit_fast_tokenizer(
|
||||
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_samples=config.fast_tokenizer_validation_samples,
|
||||
max_reconstruction_rmse=config.fast_tokenizer_max_reconstruction_rmse,
|
||||
max_dim_rmse=config.fast_tokenizer_max_dim_rmse,
|
||||
)
|
||||
@@ -1,263 +0,0 @@
|
||||
# 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
|
||||
@@ -1,19 +0,0 @@
|
||||
# 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"]
|
||||
@@ -1,254 +0,0 @@
|
||||
# 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 ""
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,164 +0,0 @@
|
||||
# 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)
|
||||
@@ -1,521 +0,0 @@
|
||||
# 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
|
||||
@@ -61,9 +61,6 @@ 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
|
||||
|
||||
@@ -25,26 +25,17 @@ from lerobot.configs import PipelineFeatureType, PolicyFeature
|
||||
from lerobot.processor import (
|
||||
AbsoluteActionsProcessorStep,
|
||||
ActionTokenizerProcessorStep,
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
ProcessorStep,
|
||||
ProcessorStepRegistry,
|
||||
RelativeActionsProcessorStep,
|
||||
RenameObservationsProcessorStep,
|
||||
TokenizerProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
make_default_policy_processor_steps,
|
||||
make_policy_processor_pipelines,
|
||||
)
|
||||
from lerobot.types import EnvTransition, TransitionKey
|
||||
from lerobot.utils.constants import (
|
||||
OBS_STATE,
|
||||
POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
)
|
||||
from lerobot.utils.constants import OBS_STATE
|
||||
|
||||
from .configuration_pi0_fast import PI0FastConfig
|
||||
|
||||
@@ -101,11 +92,6 @@ 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],
|
||||
@@ -140,6 +126,8 @@ def make_pi0_fast_pre_post_processors(
|
||||
action_names=getattr(config, "action_feature_names", None),
|
||||
)
|
||||
|
||||
steps = make_default_policy_processor_steps(config, dataset_stats)
|
||||
|
||||
# Pi0Fast order: relative → normalize → tokenize → model → unnormalize → absolute
|
||||
# This matches pi0/pi0.5: RelativeActionsProcessorStep runs first on raw absolute actions,
|
||||
# caching the raw state. NormalizerProcessorStep then normalizes the raw relative actions,
|
||||
@@ -148,27 +136,11 @@ 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] = [
|
||||
RenameObservationsProcessorStep(rename_map={}), # To mimic the same processor as pretrained one
|
||||
AddBatchDimensionProcessorStep(),
|
||||
steps.rename_observations, # To mimic the same processor as pretrained one
|
||||
steps.add_batch_dim,
|
||||
relative_step,
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
steps.normalize,
|
||||
Pi0FastPrepareStateAndLanguageTokenizerProcessorStep(max_state_dim=config.max_state_dim),
|
||||
TokenizerProcessorStep(
|
||||
tokenizer_name=config.text_tokenizer_name,
|
||||
@@ -177,31 +149,18 @@ def make_pi0_fast_pre_post_processors(
|
||||
padding="max_length",
|
||||
),
|
||||
ActionTokenizerProcessorStep(
|
||||
action_tokenizer_name=action_tokenizer_path,
|
||||
action_tokenizer_name=config.action_tokenizer_name,
|
||||
max_action_tokens=config.max_action_tokens,
|
||||
fast_skip_tokens=config.fast_skip_tokens,
|
||||
paligemma_tokenizer_name=config.text_tokenizer_name,
|
||||
),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
steps.to_device,
|
||||
]
|
||||
|
||||
output_steps: list[ProcessorStep] = [
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
steps.unnormalize,
|
||||
AbsoluteActionsProcessorStep(enabled=config.use_relative_actions, relative_step=relative_step),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
steps.to_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,
|
||||
),
|
||||
)
|
||||
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
|
||||
|
||||
@@ -14,27 +14,18 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Literal
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
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,
|
||||
@@ -58,8 +49,6 @@ else:
|
||||
GradientCheckpointingLayer = None
|
||||
BaseModelOutputWithPast = None
|
||||
create_causal_mask = None
|
||||
CONFIG_MAPPING = None
|
||||
modeling_gemma = None
|
||||
|
||||
|
||||
def _gated_residual(
|
||||
@@ -132,10 +121,7 @@ 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)
|
||||
# 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:
|
||||
if len(x.shape) == 3:
|
||||
modulation = modulation.unsqueeze(1)
|
||||
scale, shift, gate = modulation.chunk(3, dim=-1)
|
||||
normed = normed * (1 + scale.float()) + shift.float()
|
||||
@@ -289,8 +275,6 @@ 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)
|
||||
@@ -383,374 +367,3 @@ __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
|
||||
|
||||
@@ -23,8 +23,6 @@ from pathlib import Path
|
||||
from tempfile import TemporaryDirectory
|
||||
from typing import TYPE_CHECKING, TypedDict, TypeVar, Unpack
|
||||
|
||||
import packaging
|
||||
import safetensors
|
||||
from huggingface_hub import HfApi, ModelCard, ModelCardData, hf_hub_download, save_torch_state_dict
|
||||
from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
|
||||
from huggingface_hub.errors import HfHubHTTPError
|
||||
@@ -34,6 +32,7 @@ from torch import Tensor, nn
|
||||
from lerobot.__version__ import __version__
|
||||
from lerobot.configs import PreTrainedConfig
|
||||
from lerobot.configs.train import TrainPipelineConfig
|
||||
from lerobot.utils.device_utils import resolve_safetensors_device
|
||||
from lerobot.utils.hub import HubMixin
|
||||
|
||||
from .utils import log_model_loading_keys
|
||||
@@ -221,26 +220,10 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
|
||||
|
||||
@classmethod
|
||||
def _load_as_safetensor(cls, model: T, model_file: str, map_location: str, strict: bool) -> T:
|
||||
# Create base kwargs
|
||||
kwargs = {"strict": strict}
|
||||
|
||||
# Add device parameter for newer versions that support it
|
||||
if packaging.version.parse(safetensors.__version__) >= packaging.version.parse("0.4.3"):
|
||||
kwargs["device"] = map_location
|
||||
|
||||
# Load the model with appropriate kwargs
|
||||
missing_keys, unexpected_keys = load_model_as_safetensor(model, model_file, **kwargs)
|
||||
missing_keys, unexpected_keys = load_model_as_safetensor(
|
||||
model, model_file, strict=strict, device=resolve_safetensors_device(map_location)
|
||||
)
|
||||
log_model_loading_keys(missing_keys, unexpected_keys)
|
||||
|
||||
# For older versions, manually move to device if needed
|
||||
if "device" not in kwargs and map_location != "cpu":
|
||||
logging.warning(
|
||||
"Loading model weights on other devices than 'cpu' is not supported natively in your version of safetensors."
|
||||
" This means that the model is loaded on 'cpu' first and then copied to the device."
|
||||
" This leads to a slower loading time."
|
||||
" Please update safetensors to version 0.4.3 or above for improved performance."
|
||||
)
|
||||
model.to(map_location)
|
||||
return model
|
||||
|
||||
@abc.abstractmethod
|
||||
|
||||
@@ -19,19 +19,13 @@ from typing import Any
|
||||
import torch
|
||||
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NewLineTaskProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
RenameObservationsProcessorStep,
|
||||
TokenizerProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
make_default_policy_processor_steps,
|
||||
make_policy_processor_pipelines,
|
||||
)
|
||||
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
|
||||
|
||||
from .configuration_smolvla import SmolVLAConfig
|
||||
|
||||
@@ -66,9 +60,11 @@ def make_smolvla_pre_post_processors(
|
||||
A tuple containing the configured pre-processor and post-processor pipelines.
|
||||
"""
|
||||
|
||||
steps = make_default_policy_processor_steps(config, dataset_stats)
|
||||
|
||||
input_steps = [
|
||||
RenameObservationsProcessorStep(rename_map={}), # To mimic the same processor as pretrained one
|
||||
AddBatchDimensionProcessorStep(),
|
||||
steps.rename_observations, # To mimic the same processor as pretrained one
|
||||
steps.add_batch_dim,
|
||||
NewLineTaskProcessorStep(),
|
||||
TokenizerProcessorStep(
|
||||
tokenizer_name=config.vlm_model_name,
|
||||
@@ -76,28 +72,11 @@ def make_smolvla_pre_post_processors(
|
||||
padding_side="right",
|
||||
max_length=config.tokenizer_max_length,
|
||||
),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
steps.to_device,
|
||||
steps.normalize,
|
||||
]
|
||||
output_steps = [
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
steps.unnormalize,
|
||||
steps.to_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,
|
||||
),
|
||||
)
|
||||
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
|
||||
|
||||
@@ -19,17 +19,10 @@ from typing import Any
|
||||
import torch
|
||||
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
RenameObservationsProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
make_default_pre_post_processors,
|
||||
)
|
||||
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
|
||||
|
||||
from .configuration_tdmpc import TDMPCConfig
|
||||
|
||||
@@ -61,32 +54,4 @@ def make_tdmpc_pre_post_processors(
|
||||
Returns:
|
||||
A tuple containing the configured pre-processor and post-processor pipelines.
|
||||
"""
|
||||
|
||||
input_steps = [
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
]
|
||||
output_steps = [
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
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,
|
||||
),
|
||||
)
|
||||
return make_default_pre_post_processors(config, dataset_stats)
|
||||
|
||||
@@ -20,20 +20,16 @@ import torch
|
||||
|
||||
from lerobot.policies.vla_jepa.configuration_vla_jepa import VLAJEPAConfig
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
EnvTransition,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
ProcessorStep,
|
||||
ProcessorStepRegistry,
|
||||
RenameObservationsProcessorStep,
|
||||
TransitionKey,
|
||||
UnnormalizerProcessorStep,
|
||||
make_default_policy_processor_steps,
|
||||
make_policy_processor_pipelines,
|
||||
)
|
||||
from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
|
||||
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
|
||||
|
||||
|
||||
@ProcessorStepRegistry.register(name="vla_jepa_clip_actions")
|
||||
@@ -112,15 +108,12 @@ def make_vla_jepa_pre_post_processors(
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction],
|
||||
]:
|
||||
features = {**config.input_features, **config.output_features}
|
||||
steps = make_default_policy_processor_steps(config, dataset_stats)
|
||||
input_steps = [
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
NormalizerProcessorStep(
|
||||
features=features,
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
steps.rename_observations,
|
||||
steps.add_batch_dim,
|
||||
steps.to_device,
|
||||
steps.normalize,
|
||||
]
|
||||
output_steps: list[ProcessorStep] = []
|
||||
if config.clip_normalized_actions:
|
||||
@@ -129,6 +122,8 @@ def make_vla_jepa_pre_post_processors(
|
||||
output_steps.append(
|
||||
PreSnapGripperProcessorStep(gripper_dim=config.gripper_dim, threshold=config.gripper_threshold)
|
||||
)
|
||||
# NOTE: unlike the default policy unnormalizer (output features only), VLA-JEPA
|
||||
# unnormalizes over BOTH input and output features.
|
||||
output_steps.append(
|
||||
UnnormalizerProcessorStep(
|
||||
features=features,
|
||||
@@ -140,16 +135,5 @@ def make_vla_jepa_pre_post_processors(
|
||||
output_steps.append(
|
||||
BinarizeGripperProcessorStep(gripper_dim=config.gripper_dim, threshold=config.gripper_threshold)
|
||||
)
|
||||
output_steps.append(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,
|
||||
),
|
||||
)
|
||||
output_steps.append(steps.to_cpu)
|
||||
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
|
||||
|
||||
@@ -20,17 +20,10 @@ from typing import Any
|
||||
import torch
|
||||
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
RenameObservationsProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
make_default_pre_post_processors,
|
||||
)
|
||||
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
|
||||
|
||||
from .configuration_vqbet import VQBeTConfig
|
||||
|
||||
@@ -62,32 +55,4 @@ def make_vqbet_pre_post_processors(
|
||||
Returns:
|
||||
A tuple containing the configured pre-processor and post-processor pipelines.
|
||||
"""
|
||||
|
||||
input_steps = [
|
||||
RenameObservationsProcessorStep(rename_map={}), # Let the possibility to the user to rename the keys
|
||||
AddBatchDimensionProcessorStep(),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
]
|
||||
output_steps = [
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
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,
|
||||
),
|
||||
)
|
||||
return make_default_pre_post_processors(config, dataset_stats)
|
||||
|
||||
@@ -58,10 +58,14 @@ class WallXConfig(PreTrainedConfig):
|
||||
# Action prediction mode: "diffusion" or "fast"
|
||||
prediction_mode: str = "diffusion"
|
||||
|
||||
# Attention Implementation, options: "eager", "flash_attention_2", "sdpa"
|
||||
# NOTE: flash-attn==2.7.4.post1 is required for flash_attention_2 implementation
|
||||
# Wall-X's bidirectional action-token islands currently require eager attention.
|
||||
attn_implementation: str = "eager"
|
||||
|
||||
# Vision attention is independent from the text action-token mask. ``auto`` uses
|
||||
# PyTorch's packed variable-length attention when the runtime supports it and
|
||||
# otherwise falls back to the native per-chunk SDPA implementation.
|
||||
vision_attn_implementation: str = "auto"
|
||||
|
||||
# ==================== Optimizer Presets ====================
|
||||
optimizer_lr: float = 2e-5
|
||||
optimizer_betas: tuple[float, float] = (0.9, 0.95)
|
||||
@@ -86,6 +90,18 @@ class WallXConfig(PreTrainedConfig):
|
||||
if self.prediction_mode not in ["diffusion", "fast"]:
|
||||
raise ValueError(f"prediction_mode must be 'diffusion' or 'fast', got {self.prediction_mode}")
|
||||
|
||||
if self.attn_implementation != "eager":
|
||||
raise ValueError(
|
||||
"Wall-X currently supports only attn_implementation='eager' because its "
|
||||
"bidirectional action-token islands require an explicit attention mask."
|
||||
)
|
||||
|
||||
if self.vision_attn_implementation not in {"auto", "sdpa", "varlen"}:
|
||||
raise ValueError(
|
||||
"vision_attn_implementation must be one of 'auto', 'sdpa', or 'varlen', got "
|
||||
f"{self.vision_attn_implementation!r}"
|
||||
)
|
||||
|
||||
# Assign use_fast_tokenizer based on prediction_mode
|
||||
if self.prediction_mode == "fast":
|
||||
self.use_fast_tokenizer = True
|
||||
|
||||
@@ -43,11 +43,14 @@ from typing import TYPE_CHECKING, Any
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from PIL import Image
|
||||
import torch.nn.functional as functional
|
||||
from safetensors import SafetensorError
|
||||
from safetensors.torch import load_file
|
||||
from torch import Tensor
|
||||
from torch.distributions import Beta
|
||||
from torch.nn import CrossEntropyLoss
|
||||
from torchvision.transforms import InterpolationMode
|
||||
from torchvision.transforms.v2 import functional as tv_functional
|
||||
|
||||
from lerobot.utils.constants import ACTION, OBS_STATE
|
||||
from lerobot.utils.import_utils import (
|
||||
@@ -74,17 +77,17 @@ if TYPE_CHECKING or _wallx_deps_available:
|
||||
from qwen_vl_utils.vision_process import smart_resize
|
||||
from torchdiffeq import odeint
|
||||
from transformers import AutoProcessor, BatchFeature
|
||||
from transformers.cache_utils import StaticCache
|
||||
from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
|
||||
Qwen2_5_VisionTransformerPretrainedModel,
|
||||
Qwen2_5_VLForConditionalGeneration,
|
||||
)
|
||||
from transformers.utils import is_torchdynamo_compiling
|
||||
from transformers.utils import cached_file, is_torchdynamo_compiling
|
||||
|
||||
from .qwen_model.configuration_qwen2_5_vl import Qwen2_5_VLConfig
|
||||
from .qwen_model.qwen2_5_vl_moe import (
|
||||
Qwen2_5_VisionTransformerPretrainedModel,
|
||||
from .qwen_model import (
|
||||
Qwen2_5_VLACausalLMOutputWithPast,
|
||||
Qwen2_5_VLConfig,
|
||||
Qwen2_5_VLMoEModel,
|
||||
configure_wall_x_vision_attention,
|
||||
)
|
||||
else:
|
||||
LoraConfig = None
|
||||
@@ -93,13 +96,14 @@ else:
|
||||
odeint = None
|
||||
AutoProcessor = None
|
||||
BatchFeature = None
|
||||
StaticCache = None
|
||||
Qwen2_5_VLForConditionalGeneration = None
|
||||
cached_file = None
|
||||
is_torchdynamo_compiling = None
|
||||
Qwen2_5_VLConfig = None
|
||||
Qwen2_5_VisionTransformerPretrainedModel = None
|
||||
Qwen2_5_VLACausalLMOutputWithPast = None
|
||||
Qwen2_5_VLMoEModel = None
|
||||
configure_wall_x_vision_attention = None
|
||||
|
||||
from .utils import (
|
||||
get_wallx_normal_text,
|
||||
@@ -111,6 +115,75 @@ from .utils import (
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _wall_x_resize_dimensions(height: int, width: int) -> tuple[int, int, int, int]:
|
||||
"""Return the intermediate and final Wall-X resize dimensions as ``(H, W, H, W)``."""
|
||||
if RESOLUTION == -1:
|
||||
intermediate_height, intermediate_width = height, width
|
||||
elif width > height:
|
||||
intermediate_width = RESOLUTION
|
||||
intermediate_height = int(RESOLUTION * height / width)
|
||||
else:
|
||||
intermediate_height = RESOLUTION
|
||||
intermediate_width = int(RESOLUTION * width / height)
|
||||
|
||||
resized_height, resized_width = smart_resize(
|
||||
intermediate_height,
|
||||
intermediate_width,
|
||||
factor=IMAGE_FACTOR,
|
||||
min_pixels=MIN_PIXELS,
|
||||
max_pixels=MAX_PIXELS,
|
||||
)
|
||||
return intermediate_height, intermediate_width, resized_height, resized_width
|
||||
|
||||
|
||||
def _resize_wall_x_image_batch(images: Tensor) -> tuple[Tensor, tuple[int, int, int, int]]:
|
||||
"""Quantize and resize a BCHW camera batch without leaving its current device."""
|
||||
if images.ndim != 4:
|
||||
raise ValueError(f"Wall-X images must be BCHW tensors, got shape {tuple(images.shape)}")
|
||||
|
||||
original_height, original_width = images.shape[-2:]
|
||||
intermediate_height, intermediate_width, resized_height, resized_width = _wall_x_resize_dimensions(
|
||||
original_height, original_width
|
||||
)
|
||||
|
||||
if images.is_floating_point():
|
||||
# Match the previous PIL path, which quantized via `(image * 255).to(torch.uint8)`.
|
||||
images = (images * 255).to(torch.uint8)
|
||||
elif images.dtype != torch.uint8:
|
||||
raise TypeError(f"Wall-X images must be floating point or uint8, got {images.dtype}")
|
||||
|
||||
if images.shape[-2:] != (intermediate_height, intermediate_width):
|
||||
images = tv_functional.resize(
|
||||
images,
|
||||
[intermediate_height, intermediate_width],
|
||||
interpolation=InterpolationMode.BICUBIC,
|
||||
antialias=True,
|
||||
)
|
||||
if images.shape[-2:] != (resized_height, resized_width):
|
||||
images = tv_functional.resize(
|
||||
images,
|
||||
[resized_height, resized_width],
|
||||
interpolation=InterpolationMode.BICUBIC,
|
||||
antialias=True,
|
||||
)
|
||||
|
||||
return images, (original_height, original_width, resized_height, resized_width)
|
||||
|
||||
|
||||
def _prepare_wall_x_image_inputs(
|
||||
batch: dict[str, Any], img_keys: list[str]
|
||||
) -> tuple[list[list[Tensor]], dict[str, tuple[int, int, int, int]]]:
|
||||
"""Resize each camera as a batch, then restore sample-major/camera-minor ordering."""
|
||||
resized_by_key: dict[str, Tensor] = {}
|
||||
dimensions_by_key: dict[str, tuple[int, int, int, int]] = {}
|
||||
for key in img_keys:
|
||||
resized_by_key[key], dimensions_by_key[key] = _resize_wall_x_image_batch(batch[key])
|
||||
|
||||
batch_size = batch[img_keys[0]].shape[0]
|
||||
image_inputs = [[resized_by_key[key][i] for key in img_keys] for i in range(batch_size)]
|
||||
return image_inputs, dimensions_by_key
|
||||
|
||||
|
||||
class SinusoidalPosEmb(nn.Module):
|
||||
"""Sinusoidal positional embedding for diffusion timesteps."""
|
||||
|
||||
@@ -246,7 +319,7 @@ class ActionHead(nn.Module):
|
||||
flow = flow.to(torch.float32)
|
||||
|
||||
action_pred = self.action_proj_back(action_hidden_states)
|
||||
loss = F.mse_loss(action_pred, flow, reduction="none")
|
||||
loss = functional.mse_loss(action_pred, flow, reduction="none")
|
||||
|
||||
if dof_mask is not None:
|
||||
dof_mask = dof_mask.reshape(-1, dof_mask.shape[-1]).to(torch.float32)
|
||||
@@ -254,7 +327,7 @@ class ActionHead(nn.Module):
|
||||
|
||||
return loss
|
||||
|
||||
def proprioception_proj(self, proprioception, dof_mask=None, use_history=False):
|
||||
def proprioception_proj(self, proprioception, dof_mask=None):
|
||||
"""Project proprioceptive data to hidden space."""
|
||||
# Ensure proper device and dtype alignment
|
||||
proprioception = proprioception.to(device=self.propri_proj.weight.device).to(
|
||||
@@ -264,10 +337,7 @@ class ActionHead(nn.Module):
|
||||
if dof_mask is not None:
|
||||
# Concatenate proprioception with DOF mask
|
||||
# TODO: Use variable-based dimension checking for better flexibility
|
||||
if use_history:
|
||||
proprioception = torch.cat([proprioception, dof_mask], dim=-1)
|
||||
else:
|
||||
proprioception = torch.cat([proprioception, dof_mask], dim=-1)
|
||||
proprioception = torch.cat([proprioception, dof_mask], dim=-1)
|
||||
|
||||
proprioception = proprioception.to(device=self.propri_proj.weight.device).to(
|
||||
dtype=self.propri_proj.weight.dtype
|
||||
@@ -281,7 +351,7 @@ class ActionHead(nn.Module):
|
||||
_Qwen2_5_VLForAction_Base = Qwen2_5_VLForConditionalGeneration if _wallx_deps_available else nn.Module
|
||||
|
||||
|
||||
class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base): # noqa: N801
|
||||
"""
|
||||
Qwen2.5 Vision-Language Mixture of Experts model for action processing.
|
||||
|
||||
@@ -305,6 +375,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
config=None,
|
||||
action_tokenizer_path=None,
|
||||
attn_implementation: str = "eager",
|
||||
vision_attn_implementation: str = "auto",
|
||||
cache_dir: str | PathLike | None = None,
|
||||
force_download: bool = False,
|
||||
local_files_only: bool = False,
|
||||
@@ -321,11 +392,14 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
config_path (str, optional): Configuration file path, if None will look for qwen25_config.json in pretrained_model_path
|
||||
action_tokenizer_path (str, optional): Action tokenizer path, if None will load from default config
|
||||
attn_implementation (str, optional): Attention implementation, if None will load from default config
|
||||
vision_attn_implementation (str, optional): Vision attention backend. ``auto`` uses packed
|
||||
variable-length attention when supported and otherwise falls back to SDPA.
|
||||
**kwargs: Additional arguments
|
||||
|
||||
Returns:
|
||||
Qwen2_5_VLMoEForAction: Loaded model instance
|
||||
"""
|
||||
Qwen2_5_VLMoEModel._require_eager_attention(attn_implementation)
|
||||
if config is None:
|
||||
config = cls.config_class.from_pretrained(
|
||||
pretrained_name_or_path,
|
||||
@@ -339,7 +413,15 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
)
|
||||
if attn_implementation is not None:
|
||||
config._attn_implementation = attn_implementation
|
||||
processor = AutoProcessor.from_pretrained(pretrained_name_or_path, use_fast=True)
|
||||
processor = AutoProcessor.from_pretrained(
|
||||
pretrained_name_or_path,
|
||||
cache_dir=cache_dir,
|
||||
force_download=force_download,
|
||||
local_files_only=local_files_only,
|
||||
token=token,
|
||||
revision=revision,
|
||||
use_fast=True,
|
||||
)
|
||||
if action_tokenizer_path is not None:
|
||||
action_tokenizer = AutoProcessor.from_pretrained(action_tokenizer_path, trust_remote_code=True)
|
||||
processor.action_processor = action_tokenizer
|
||||
@@ -351,41 +433,41 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
config.text_config.pad_token_id = processor.tokenizer.pad_token_id
|
||||
|
||||
# Initialize model with configuration and processor
|
||||
model = cls(config, processor=processor, action_tokenizer=action_tokenizer, **kwargs)
|
||||
model = cls(
|
||||
config,
|
||||
processor=processor,
|
||||
action_tokenizer=action_tokenizer,
|
||||
vision_attn_implementation=vision_attn_implementation,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# Resize token embeddings to match processor tokenizer vocabulary size
|
||||
model.resize_token_embeddings(len(processor.tokenizer))
|
||||
|
||||
# Try to load the model.safetensors file
|
||||
print(f"Loading model from: {pretrained_name_or_path}")
|
||||
logger.info("Loading Wall-X model from %s", pretrained_name_or_path)
|
||||
try:
|
||||
from transformers.utils import cached_file
|
||||
|
||||
# Try safetensors first
|
||||
resolved_file = cached_file(
|
||||
pretrained_name_or_path,
|
||||
"model.safetensors",
|
||||
cache_dir=kwargs.get("cache_dir"),
|
||||
force_download=kwargs.get("force_download", False),
|
||||
cache_dir=cache_dir,
|
||||
force_download=force_download,
|
||||
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),
|
||||
token=token,
|
||||
revision=revision,
|
||||
local_files_only=local_files_only,
|
||||
)
|
||||
from safetensors.torch import load_file
|
||||
|
||||
sd = 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
|
||||
except (OSError, SafetensorError) as error:
|
||||
raise OSError(
|
||||
f"Failed to load pretrained Wall-X weights from {pretrained_name_or_path!r}"
|
||||
) from error
|
||||
logger.info("Loaded Wall-X state dict from model.safetensors")
|
||||
|
||||
state_dict = {}
|
||||
# filter normalizer statistic params
|
||||
del_keys = []
|
||||
for key in sd.keys():
|
||||
for key in sd:
|
||||
if "action_preprocessor.normalizer" in key:
|
||||
del_keys.append(key)
|
||||
for key in del_keys:
|
||||
@@ -404,6 +486,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
action_tokenizer=None,
|
||||
action_mapper=None,
|
||||
flow_loss_weight=1.0,
|
||||
vision_attn_implementation: str = "auto",
|
||||
):
|
||||
"""
|
||||
Initialize the Qwen2.5 VLMoE model for action processing.
|
||||
@@ -416,10 +499,16 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
action_mapper: Action mapping utility
|
||||
flow_loss_weight (float): Weight for flow loss computation
|
||||
"""
|
||||
Qwen2_5_VLMoEModel._require_eager_attention(config._attn_implementation)
|
||||
config._attn_implementation = "eager"
|
||||
# Text needs eager attention for action-token islands. Vision has no such
|
||||
# constraint, so keep its portable native fallback on SDPA.
|
||||
config.vision_config._attn_implementation = "sdpa"
|
||||
super().__init__(config)
|
||||
|
||||
# Initialize vision transformer and language model components
|
||||
self.visual = Qwen2_5_VisionTransformerPretrainedModel._from_config(config.vision_config)
|
||||
configure_wall_x_vision_attention(self.visual, vision_attn_implementation)
|
||||
self.model = Qwen2_5_VLMoEModel(config)
|
||||
self.vocab_size = config.vocab_size
|
||||
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
|
||||
@@ -457,7 +546,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
|
||||
params_to_keep_float32 = []
|
||||
|
||||
for name, param in self.named_parameters():
|
||||
for name, _param in self.named_parameters():
|
||||
if "input_layernorm" in name or "post_attention_layernorm" in name or "model.norm" in name:
|
||||
params_to_keep_float32.append(name)
|
||||
if "action_preprocessor" in name:
|
||||
@@ -491,7 +580,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
"action_token_id": action_token_id,
|
||||
}
|
||||
|
||||
def add_lora(self, r=8, lora_alpha=32, target_modules=["q_proj", "v_proj"], lora_dropout=0.1):
|
||||
def add_lora(self, r=8, lora_alpha=32, target_modules=None, lora_dropout=0.1):
|
||||
"""
|
||||
Add LoRA (Low-Rank Adaptation) adapters to the model.
|
||||
|
||||
@@ -501,6 +590,9 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
target_modules (list): List of module names to apply LoRA to
|
||||
lora_dropout (float): Dropout probability for LoRA layers
|
||||
"""
|
||||
if target_modules is None:
|
||||
target_modules = ["q_proj", "v_proj"]
|
||||
|
||||
config = LoraConfig(
|
||||
r=r,
|
||||
lora_alpha=lora_alpha,
|
||||
@@ -795,6 +887,9 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
)
|
||||
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
||||
|
||||
if rope_deltas is not None:
|
||||
self.rope_deltas = rope_deltas
|
||||
|
||||
# Calculate RoPE position IDs if not provided
|
||||
# Note: Cannot calculate rope deltas with 4D attention mask. TODO: Fix this limitation
|
||||
if position_ids is None and (attention_mask is None or attention_mask.ndim == 2):
|
||||
@@ -833,7 +928,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
# Process image embeddings
|
||||
if pixel_values is not None:
|
||||
pixel_values = pixel_values.type(self.visual.dtype)
|
||||
image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw)
|
||||
image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw).pooler_output
|
||||
mask = input_ids == self.config.image_token_id
|
||||
mask_unsqueezed = mask.unsqueeze(-1)
|
||||
mask_expanded = mask_unsqueezed.expand_as(inputs_embeds)
|
||||
@@ -845,7 +940,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
# Process video embeddings
|
||||
if pixel_values_videos is not None:
|
||||
pixel_values_videos = pixel_values_videos.type(self.visual.dtype)
|
||||
video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw)
|
||||
video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw).pooler_output
|
||||
n_video_tokens = (input_ids == self.config.video_token_id).sum().item()
|
||||
n_video_features = video_embeds.shape[0]
|
||||
|
||||
@@ -869,7 +964,6 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
proprioception = self.action_preprocessor.proprioception_proj(
|
||||
proprioception,
|
||||
agent_pos_mask,
|
||||
use_history=proprioception.shape[1] > 1,
|
||||
)
|
||||
mask = input_ids == self.action_token_id_set["propri_token_id"]
|
||||
mask_unsqueezed = mask.unsqueeze(-1)
|
||||
@@ -919,6 +1013,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
output_attentions=output_attentions,
|
||||
output_hidden_states=output_hidden_states,
|
||||
return_dict=return_dict,
|
||||
cache_position=cache_position,
|
||||
)
|
||||
|
||||
hidden_states = outputs[0]
|
||||
@@ -1107,7 +1202,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
# Process image embeddings
|
||||
if pixel_values is not None:
|
||||
pixel_values = pixel_values.type(self.visual.dtype)
|
||||
image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw)
|
||||
image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw).pooler_output
|
||||
n_image_tokens = (input_ids == self.config.image_token_id).sum().item()
|
||||
n_image_features = image_embeds.shape[0]
|
||||
|
||||
@@ -1128,7 +1223,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
# Process video embeddings
|
||||
if pixel_values_videos is not None:
|
||||
pixel_values_videos = pixel_values_videos.type(self.visual.dtype)
|
||||
video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw)
|
||||
video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw).pooler_output
|
||||
n_video_tokens = (input_ids == self.config.video_token_id).sum().item()
|
||||
n_video_features = video_embeds.shape[0]
|
||||
|
||||
@@ -1153,7 +1248,6 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
proprio_embed = self.action_preprocessor.proprioception_proj(
|
||||
proprioception,
|
||||
agent_pos_mask,
|
||||
use_history=proprioception.shape[1] > 1,
|
||||
)
|
||||
proprioception_mask = input_ids == self.action_token_id_set["propri_token_id"]
|
||||
proprio_embed = proprio_embed.to(torch.bfloat16)
|
||||
@@ -1202,25 +1296,37 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
|
||||
# Split input sequence for text and fast modes (not needed for diffusion)
|
||||
if predict_mode == "text" or predict_mode == "fast":
|
||||
# Look for generation prompt tokens: <|im_start|>assistant
|
||||
generation_prompt = "<|im_start|>assistant\n"
|
||||
generation_prompt_ids = torch.tensor(
|
||||
[151644, 77091], device=input_ids.device, dtype=input_ids.dtype
|
||||
)
|
||||
matches = (input_ids[0, :-1] == generation_prompt_ids[0]) & (
|
||||
input_ids[0, 1:] == generation_prompt_ids[1]
|
||||
self.processor.tokenizer.encode(generation_prompt, add_special_tokens=False),
|
||||
device=input_ids.device,
|
||||
dtype=input_ids.dtype,
|
||||
)
|
||||
prompt_length = generation_prompt_ids.numel()
|
||||
if prompt_length == 0:
|
||||
raise ValueError(f"Tokenizer produced no tokens for generation prompt {generation_prompt!r}")
|
||||
if input_ids.shape[1] < prompt_length:
|
||||
matches = torch.empty(0, device=input_ids.device, dtype=torch.bool)
|
||||
else:
|
||||
matches = (
|
||||
input_ids[0]
|
||||
.unfold(dimension=0, size=prompt_length, step=1)
|
||||
.eq(generation_prompt_ids)
|
||||
.all(dim=-1)
|
||||
)
|
||||
|
||||
if matches.any():
|
||||
split_pos = torch.nonzero(matches, as_tuple=True)[0][0].item()
|
||||
prompt_end = split_pos + prompt_length
|
||||
# Extract ground truth output tokens (including newline)
|
||||
gt_output_ids = input_ids[:, split_pos + 3 :]
|
||||
gt_output_ids = input_ids[:, prompt_end:]
|
||||
# Remove output part from input, keeping prompt
|
||||
input_ids = input_ids[:, : split_pos + 3]
|
||||
inputs_embeds = inputs_embeds[:, : split_pos + 3, :]
|
||||
input_ids = input_ids[:, :prompt_end]
|
||||
inputs_embeds = inputs_embeds[:, :prompt_end, :]
|
||||
if attention_mask is not None:
|
||||
attention_mask = attention_mask[:, : split_pos + 3]
|
||||
attention_mask = attention_mask[:, :prompt_end]
|
||||
if labels is not None:
|
||||
labels = labels[:, split_pos + 3 :]
|
||||
labels = labels[:, prompt_end:]
|
||||
else:
|
||||
raise ValueError(
|
||||
"input_ids does not contain the generation prompt tokens <|im_start|>assistant"
|
||||
@@ -1255,7 +1361,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
use_cache=True,
|
||||
pad_token_id=self.processor.tokenizer.pad_token_id,
|
||||
temperature=(1.0 if not re_generate else 0.7), # Higher temperature for regeneration
|
||||
do_sample=(False if not re_generate else True), # Enable sampling for regeneration
|
||||
do_sample=re_generate, # Enable sampling for regeneration
|
||||
)
|
||||
|
||||
# Decode generated and ground truth text
|
||||
@@ -1524,27 +1630,6 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
|
||||
else:
|
||||
model_inputs = {"input_ids": input_ids, "inputs_embeds": None}
|
||||
|
||||
# Prepare 4D causal attention mask for static cache
|
||||
if isinstance(past_key_values, StaticCache) and attention_mask.ndim == 2:
|
||||
if model_inputs["inputs_embeds"] is not None:
|
||||
batch_size, sequence_length, _ = inputs_embeds.shape
|
||||
device = inputs_embeds.device
|
||||
else:
|
||||
batch_size, sequence_length = input_ids.shape
|
||||
device = input_ids.device
|
||||
|
||||
attention_mask = self.model._prepare_4d_causal_attention_mask_with_cache_position(
|
||||
attention_mask,
|
||||
sequence_length=sequence_length,
|
||||
target_length=past_key_values.get_max_cache_shape(),
|
||||
dtype=self.lm_head.weight.dtype,
|
||||
device=device,
|
||||
cache_position=cache_position,
|
||||
batch_size=batch_size,
|
||||
config=self.config,
|
||||
past_key_values=past_key_values,
|
||||
)
|
||||
|
||||
# Assemble all model inputs for generation
|
||||
model_inputs.update(
|
||||
{
|
||||
@@ -1749,6 +1834,7 @@ class WallXPolicy(PreTrainedPolicy):
|
||||
pretrained_name_or_path=config.pretrained_name_or_path,
|
||||
action_tokenizer_path=config.action_tokenizer_path,
|
||||
attn_implementation=config.attn_implementation,
|
||||
vision_attn_implementation=config.vision_attn_implementation,
|
||||
)
|
||||
self.model.to(config.device)
|
||||
self.model.to_bfloat16_for_selected_params()
|
||||
@@ -1768,6 +1854,8 @@ class WallXPolicy(PreTrainedPolicy):
|
||||
def preprocess_inputs(
|
||||
self,
|
||||
batch: dict[str, Any],
|
||||
*,
|
||||
compute_position_ids: bool = False,
|
||||
) -> BatchFeature:
|
||||
"""
|
||||
Convert a batch of LeRobot dataset items to Wall-X model input format.
|
||||
@@ -1789,50 +1877,21 @@ class WallXPolicy(PreTrainedPolicy):
|
||||
# Get batch size from state tensor
|
||||
batch_size = batch[OBS_STATE].shape[0]
|
||||
|
||||
# ==================== PROCESS ALL SAMPLES ====================
|
||||
all_image_inputs = []
|
||||
all_texts = []
|
||||
|
||||
# Find image keys in batch
|
||||
img_keys = [key for key in self.config.image_features if key in batch]
|
||||
if not img_keys:
|
||||
raise ValueError("Wall-X requires at least one image feature in each batch")
|
||||
|
||||
# Resize one camera batch at a time on the tensors' current device. Reassembling
|
||||
# sample-major keeps image_grid_thw aligned with each sample's image placeholders.
|
||||
all_image_inputs, dimensions_by_key = _prepare_wall_x_image_inputs(batch, img_keys)
|
||||
all_texts = []
|
||||
|
||||
# Preserve the existing grounding behavior for multi-camera inputs: the old camera
|
||||
# loop left these values set to the final configured camera's dimensions.
|
||||
orig_height, orig_width, resized_height, resized_width = dimensions_by_key[img_keys[-1]]
|
||||
|
||||
for i in range(batch_size):
|
||||
# Vision preprocessing per sample
|
||||
processed_frames = []
|
||||
orig_height, orig_width = None, None
|
||||
resized_height, resized_width = None, None
|
||||
|
||||
for key in img_keys:
|
||||
current_obs = batch[key][i].clone() # (C, H, W)
|
||||
if current_obs.dim() == 3:
|
||||
current_obs = current_obs.permute(1, 2, 0) # (H, W, C)
|
||||
|
||||
img_pil = Image.fromarray((current_obs * 255).to(torch.uint8).cpu().numpy())
|
||||
orig_width, orig_height = img_pil.size
|
||||
|
||||
target_size = RESOLUTION
|
||||
if target_size != -1:
|
||||
if orig_width > orig_height:
|
||||
new_width = target_size
|
||||
new_height = int(target_size * orig_height / orig_width)
|
||||
else:
|
||||
new_height = target_size
|
||||
new_width = int(target_size * orig_width / orig_height)
|
||||
img_pil = img_pil.resize((new_width, new_height))
|
||||
|
||||
current_width, current_height = img_pil.size
|
||||
resized_height, resized_width = smart_resize(
|
||||
current_height,
|
||||
current_width,
|
||||
factor=IMAGE_FACTOR,
|
||||
min_pixels=MIN_PIXELS,
|
||||
max_pixels=MAX_PIXELS,
|
||||
)
|
||||
resized_img = img_pil.resize((resized_width, resized_height))
|
||||
processed_frames.append(resized_img)
|
||||
|
||||
all_image_inputs.append(processed_frames)
|
||||
|
||||
# Text preprocessing
|
||||
task_text = batch["task"][i] if isinstance(batch["task"], list) else batch["task"]
|
||||
instruction_info = {"instruction": task_text}
|
||||
@@ -1859,8 +1918,8 @@ class WallXPolicy(PreTrainedPolicy):
|
||||
agent_pos_mask = (~torch.isnan(agent_pos)).float()
|
||||
agent_pos = agent_pos.nan_to_num(nan=0.0)
|
||||
|
||||
if agent_pos.shape[-1] != 20:
|
||||
pad_size = 20 - agent_pos.shape[-1]
|
||||
if agent_pos.shape[-1] < self.config.max_state_dim:
|
||||
pad_size = self.config.max_state_dim - agent_pos.shape[-1]
|
||||
agent_pos = torch.cat(
|
||||
[
|
||||
agent_pos,
|
||||
@@ -1880,6 +1939,10 @@ class WallXPolicy(PreTrainedPolicy):
|
||||
],
|
||||
dim=-1,
|
||||
)
|
||||
elif agent_pos.shape[-1] > self.config.max_state_dim:
|
||||
raise ValueError(
|
||||
f"State dimension {agent_pos.shape[-1]} exceeds max_state_dim {self.config.max_state_dim}"
|
||||
)
|
||||
|
||||
# ==================== PROCESS ACTIONS ====================
|
||||
action = batch.get(ACTION) # (batch_size, chunk_size, action_dim)
|
||||
@@ -1889,8 +1952,8 @@ class WallXPolicy(PreTrainedPolicy):
|
||||
dof_mask = (~torch.isnan(action)).float()
|
||||
action = action.nan_to_num(nan=0.0)
|
||||
|
||||
if action.shape[-1] != 20:
|
||||
pad_size = 20 - action.shape[-1]
|
||||
if action.shape[-1] < self.config.max_action_dim:
|
||||
pad_size = self.config.max_action_dim - action.shape[-1]
|
||||
action = torch.cat(
|
||||
[action, torch.zeros(action.shape[0], action.shape[1], pad_size, device=action.device)],
|
||||
dim=-1,
|
||||
@@ -1902,6 +1965,10 @@ class WallXPolicy(PreTrainedPolicy):
|
||||
],
|
||||
dim=-1,
|
||||
)
|
||||
elif action.shape[-1] > self.config.max_action_dim:
|
||||
raise ValueError(
|
||||
f"Action dimension {action.shape[-1]} exceeds max_action_dim {self.config.max_action_dim}"
|
||||
)
|
||||
else:
|
||||
action_dim = self.config.output_features[ACTION].shape[0]
|
||||
dof_mask = torch.cat(
|
||||
@@ -1910,7 +1977,10 @@ class WallXPolicy(PreTrainedPolicy):
|
||||
batch_size, self.config.chunk_size, action_dim, device=batch[OBS_STATE].device
|
||||
),
|
||||
torch.zeros(
|
||||
batch_size, self.config.chunk_size, 20 - action_dim, device=batch[OBS_STATE].device
|
||||
batch_size,
|
||||
self.config.chunk_size,
|
||||
self.config.max_action_dim - action_dim,
|
||||
device=batch[OBS_STATE].device,
|
||||
),
|
||||
],
|
||||
dim=-1,
|
||||
@@ -1930,12 +2000,26 @@ class WallXPolicy(PreTrainedPolicy):
|
||||
text=all_texts,
|
||||
images=all_image_inputs,
|
||||
videos=None,
|
||||
device=batch[OBS_STATE].device,
|
||||
padding=True,
|
||||
truncation=True,
|
||||
return_tensors="pt",
|
||||
max_length=TOKENIZER_MAX_LENGTH,
|
||||
)
|
||||
|
||||
if compute_position_ids:
|
||||
# Qwen's RoPE indexing uses Python list/scalar conversions. Run it while the
|
||||
# tokenizer and grid metadata are still on CPU, then move the compact result.
|
||||
position_ids, rope_deltas = self.model.get_rope_index(
|
||||
inputs.input_ids,
|
||||
inputs.get("image_grid_thw"),
|
||||
inputs.get("video_grid_thw"),
|
||||
inputs.get("second_per_grid_ts"),
|
||||
inputs.attention_mask,
|
||||
)
|
||||
inputs["position_ids"] = position_ids
|
||||
inputs["rope_deltas"] = rope_deltas
|
||||
|
||||
# ==================== ADDITIONAL INPUTS ====================
|
||||
action_token_id = self.model.processor.tokenizer.convert_tokens_to_ids("<|action|>")
|
||||
moe_token_types = inputs.input_ids == action_token_id
|
||||
@@ -1952,7 +2036,7 @@ class WallXPolicy(PreTrainedPolicy):
|
||||
)
|
||||
|
||||
# Move all tensors to the correct device
|
||||
device = self.config.device
|
||||
device = batch[OBS_STATE].device
|
||||
for key, value in inputs.items():
|
||||
if isinstance(value, torch.Tensor):
|
||||
inputs[key] = value.to(device)
|
||||
@@ -1972,9 +2056,7 @@ class WallXPolicy(PreTrainedPolicy):
|
||||
Returns:
|
||||
tuple: (loss, loss_dict)
|
||||
"""
|
||||
batch = self.preprocess_inputs(
|
||||
batch,
|
||||
)
|
||||
batch = self.preprocess_inputs(batch, compute_position_ids=True)
|
||||
|
||||
# Call the underlying model's forward with mode="train"
|
||||
outputs = self.model(**batch, mode="train")
|
||||
@@ -1982,19 +2064,19 @@ class WallXPolicy(PreTrainedPolicy):
|
||||
# Extract losses from output
|
||||
loss = outputs.loss
|
||||
loss_dict = {
|
||||
"loss": loss.item() if loss is not None else 0.0,
|
||||
"loss": loss.detach() if loss is not None else 0.0,
|
||||
}
|
||||
|
||||
if outputs.flow_loss is not None:
|
||||
loss_dict["flow_loss"] = outputs.flow_loss.item()
|
||||
loss_dict["flow_loss"] = outputs.flow_loss.detach()
|
||||
if outputs.cross_entropy_loss is not None:
|
||||
loss_dict["cross_entropy_loss"] = outputs.cross_entropy_loss.item()
|
||||
loss_dict["cross_entropy_loss"] = outputs.cross_entropy_loss.detach()
|
||||
|
||||
# Add channel losses if available
|
||||
if outputs.channel_loss_dict is not None:
|
||||
for key, value in outputs.channel_loss_dict.items():
|
||||
if isinstance(value, torch.Tensor):
|
||||
loss_dict[f"channel_{key}"] = value.item()
|
||||
loss_dict[f"channel_{key}"] = value.detach()
|
||||
|
||||
return loss, loss_dict
|
||||
|
||||
|
||||
@@ -20,19 +20,13 @@ import torch
|
||||
|
||||
from lerobot.configs import PipelineFeatureType, PolicyFeature
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
ComplementaryDataProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
ProcessorStepRegistry,
|
||||
RenameObservationsProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
make_default_policy_processor_steps,
|
||||
make_policy_processor_pipelines,
|
||||
)
|
||||
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
|
||||
|
||||
from .configuration_wall_x import WallXConfig
|
||||
|
||||
@@ -65,37 +59,22 @@ def make_wall_x_pre_post_processors(
|
||||
A tuple containing the configured pre-processor and post-processor pipelines
|
||||
"""
|
||||
|
||||
steps = make_default_policy_processor_steps(config, dataset_stats)
|
||||
|
||||
input_steps = [
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
steps.rename_observations,
|
||||
steps.add_batch_dim,
|
||||
WallXTaskProcessor(), # Process task description
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
steps.normalize,
|
||||
steps.to_device,
|
||||
]
|
||||
|
||||
output_steps = [
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
steps.unnormalize,
|
||||
steps.to_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,
|
||||
),
|
||||
)
|
||||
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
|
||||
|
||||
|
||||
@ProcessorStepRegistry.register(name="wall_x_task_processor")
|
||||
|
||||
@@ -0,0 +1,45 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2025 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 .configuration_qwen2_5_vl import (
|
||||
Qwen2_5_VLConfig,
|
||||
Qwen2_5_VLTextConfig,
|
||||
Qwen2_5_VLVisionConfig,
|
||||
)
|
||||
from .qwen2_5_vl_moe import (
|
||||
BlockSparseMLP,
|
||||
Qwen2_5_VLACausalLMOutputWithPast,
|
||||
Qwen2_5_VLDecoderLayer_with_MoE,
|
||||
Qwen2_5_VLMoEModel,
|
||||
SparseMoeBlock,
|
||||
)
|
||||
from .vision_attention import (
|
||||
WallXVisionAttention,
|
||||
configure_wall_x_vision_attention,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"BlockSparseMLP",
|
||||
"Qwen2_5_VLACausalLMOutputWithPast",
|
||||
"Qwen2_5_VLConfig",
|
||||
"Qwen2_5_VLDecoderLayer_with_MoE",
|
||||
"Qwen2_5_VLMoEModel",
|
||||
"Qwen2_5_VLTextConfig",
|
||||
"Qwen2_5_VLVisionConfig",
|
||||
"SparseMoeBlock",
|
||||
"WallXVisionAttention",
|
||||
"configure_wall_x_vision_attention",
|
||||
]
|
||||
@@ -1,250 +1,114 @@
|
||||
from transformers.configuration_utils import PretrainedConfig
|
||||
from transformers.modeling_rope_utils import rope_config_validation
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2025 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.
|
||||
|
||||
"""Wall-X configuration extensions for the native Transformers Qwen2.5-VL config."""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from huggingface_hub.dataclasses import strict
|
||||
|
||||
from lerobot.utils.import_utils import _transformers_available
|
||||
|
||||
if TYPE_CHECKING or _transformers_available:
|
||||
from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import (
|
||||
Qwen2_5_VLConfig as TransformersQwen2_5_VLConfig,
|
||||
Qwen2_5_VLTextConfig as TransformersQwen2_5_VLTextConfig,
|
||||
Qwen2_5_VLVisionConfig,
|
||||
)
|
||||
else:
|
||||
|
||||
@dataclass
|
||||
class _TransformersConfigFallback:
|
||||
"""Import-safe stand-in used only when Transformers is unavailable."""
|
||||
|
||||
TransformersQwen2_5_VLConfig = _TransformersConfigFallback
|
||||
TransformersQwen2_5_VLTextConfig = _TransformersConfigFallback
|
||||
Qwen2_5_VLVisionConfig = None
|
||||
|
||||
# Wall-X checkpoints pre0.6.0 use the legacy, flat Qwen2.5-VL config layout. The native
|
||||
# ``Qwen2_5_VLConfig`` accepts that layout and moves text-model fields into its
|
||||
# ``text_config`` sub-config, so only the Wall-X-specific MoE fields need to be
|
||||
# declared here.
|
||||
_LEGACY_TEXT_ATTRIBUTES = {
|
||||
"attention_dropout",
|
||||
"attention_moe",
|
||||
"dim_inputs",
|
||||
"dof_config",
|
||||
"experts",
|
||||
"hidden_act",
|
||||
"hidden_size",
|
||||
"initializer_range",
|
||||
"intermediate_size",
|
||||
"layer_types",
|
||||
"max_position_embeddings",
|
||||
"max_window_layers",
|
||||
"mlp_moe",
|
||||
"noise_scheduler",
|
||||
"num_attention_heads",
|
||||
"num_experts",
|
||||
"num_hidden_layers",
|
||||
"num_key_value_heads",
|
||||
"pad_token_id",
|
||||
"rms_norm_eps",
|
||||
"sliding_window",
|
||||
"use_cache",
|
||||
"use_sliding_window",
|
||||
"vocab_size",
|
||||
}
|
||||
|
||||
|
||||
class Qwen2_5_VLVisionConfig(PretrainedConfig):
|
||||
model_type = "qwen2_5_vl"
|
||||
base_config_key = "vision_config"
|
||||
@strict
|
||||
class Qwen2_5_VLTextConfig(TransformersQwen2_5_VLTextConfig): # noqa: N801
|
||||
"""Native Qwen2.5-VL text config plus Wall-X's hard-routed MoE settings."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
depth=32,
|
||||
hidden_size=3584,
|
||||
hidden_act="silu",
|
||||
intermediate_size=3420,
|
||||
num_heads=16,
|
||||
in_channels=3,
|
||||
patch_size=14,
|
||||
spatial_merge_size=2,
|
||||
temporal_patch_size=2,
|
||||
tokens_per_second=4,
|
||||
window_size=112,
|
||||
out_hidden_size=3584,
|
||||
fullatt_block_indexes=[7, 15, 23, 31],
|
||||
initializer_range=0.02,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(**kwargs)
|
||||
num_experts: int = 4
|
||||
experts: list[dict] | None = None
|
||||
dof_config: dict | None = None
|
||||
noise_scheduler: dict | None = None
|
||||
dim_inputs: tuple[int, ...] | list[int] = (1536, 1536)
|
||||
attention_moe: bool = False
|
||||
mlp_moe: bool = False
|
||||
|
||||
self.depth = depth
|
||||
self.hidden_size = hidden_size
|
||||
self.hidden_act = hidden_act
|
||||
self.intermediate_size = intermediate_size
|
||||
self.num_heads = num_heads
|
||||
self.in_channels = in_channels
|
||||
self.patch_size = patch_size
|
||||
self.spatial_merge_size = spatial_merge_size
|
||||
self.temporal_patch_size = temporal_patch_size
|
||||
self.tokens_per_second = tokens_per_second
|
||||
self.window_size = window_size
|
||||
self.fullatt_block_indexes = fullatt_block_indexes
|
||||
self.out_hidden_size = out_hidden_size
|
||||
self.initializer_range = initializer_range
|
||||
def __post_init__(self, **kwargs):
|
||||
self.dim_inputs = tuple(self.dim_inputs)
|
||||
super().__post_init__(**kwargs)
|
||||
|
||||
|
||||
class Qwen2_5_VLConfig(PretrainedConfig):
|
||||
r"""
|
||||
This is the configuration class to store the configuration of a [`Qwen2_5_VLModel`]. It is used to instantiate a
|
||||
Qwen2-VL model according to the specified arguments, defining the model architecture. Instantiating a configuration
|
||||
with the defaults will yield a similar configuration to that of
|
||||
Qwen2-VL-7B-Instruct [Qwen/Qwen2-VL-7B-Instruct](https://huggingface.co/Qwen/Qwen2-VL-7B-Instruct).
|
||||
@strict
|
||||
class Qwen2_5_VLConfig(TransformersQwen2_5_VLConfig): # noqa: N801
|
||||
"""Native composite Qwen2.5-VL config with a Wall-X text sub-config.
|
||||
|
||||
Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
|
||||
documentation from [`PretrainedConfig`] for more information.
|
||||
The native composite loader supports both current nested configs and the
|
||||
flat layout used by existing ``wall-oss-flow`` checkpoints.
|
||||
"""
|
||||
|
||||
|
||||
Args:
|
||||
vocab_size (`int`, *optional*, defaults to 152064):
|
||||
Vocabulary size of the Qwen2_5_VL model. Defines the number of different tokens that can be represented by the
|
||||
`inputs_ids` passed when calling [`Qwen2_5_VLModel`]
|
||||
hidden_size (`int`, *optional*, defaults to 8192):
|
||||
Dimension of the hidden representations.
|
||||
intermediate_size (`int`, *optional*, defaults to 29568):
|
||||
Dimension of the MLP representations.
|
||||
num_hidden_layers (`int`, *optional*, defaults to 80):
|
||||
Number of hidden layers in the Transformer encoder.
|
||||
num_attention_heads (`int`, *optional*, defaults to 64):
|
||||
Number of attention heads for each attention layer in the Transformer encoder.
|
||||
num_key_value_heads (`int`, *optional*, defaults to 8):
|
||||
This is the number of key_value heads that should be used to implement Grouped Query Attention. If
|
||||
`num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
|
||||
`num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
|
||||
converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
|
||||
by meanpooling all the original heads within that group. For more details checkout [this
|
||||
paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to `32`.
|
||||
hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
|
||||
The non-linear activation function (function or string) in the decoder.
|
||||
max_position_embeddings (`int`, *optional*, defaults to 32768):
|
||||
The maximum sequence length that this model might ever be used with.
|
||||
initializer_range (`float`, *optional*, defaults to 0.02):
|
||||
The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
|
||||
rms_norm_eps (`float`, *optional*, defaults to 1e-05):
|
||||
The epsilon used by the rms normalization layers.
|
||||
use_cache (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not the model should return the last key/values attentions (not used by all models). Only
|
||||
relevant if `config.is_decoder=True`.
|
||||
tie_word_embeddings (`bool`, *optional*, defaults to `False`):
|
||||
Whether the model's input and output word embeddings should be tied.
|
||||
rope_theta (`float`, *optional*, defaults to 1000000.0):
|
||||
The base period of the RoPE embeddings.
|
||||
use_sliding_window (`bool`, *optional*, defaults to `False`):
|
||||
Whether to use sliding window attention.
|
||||
sliding_window (`int`, *optional*, defaults to 4096):
|
||||
Sliding window attention (SWA) window size. If not specified, will default to `4096`.
|
||||
max_window_layers (`int`, *optional*, defaults to 80):
|
||||
The number of layers that use SWA (Sliding Window Attention). The bottom layers use SWA while the top use full attention.
|
||||
attention_dropout (`float`, *optional*, defaults to 0.0):
|
||||
The dropout ratio for the attention probabilities.
|
||||
vision_config (`Dict`, *optional*):
|
||||
The config for the visual encoder initialization.
|
||||
rope_scaling (`Dict`, *optional*):
|
||||
Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type
|
||||
and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value
|
||||
accordingly.
|
||||
Expected contents:
|
||||
`rope_type` (`str`):
|
||||
The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',
|
||||
'llama3'], with 'default' being the original RoPE implementation.
|
||||
`factor` (`float`, *optional*):
|
||||
Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In
|
||||
most scaling types, a `factor` of x will enable the model to handle sequences of length x *
|
||||
original maximum pre-trained length.
|
||||
`original_max_position_embeddings` (`int`, *optional*):
|
||||
Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during
|
||||
pretraining.
|
||||
`attention_factor` (`float`, *optional*):
|
||||
Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention
|
||||
computation. If unspecified, it defaults to value recommended by the implementation, using the
|
||||
`factor` field to infer the suggested value.
|
||||
`beta_fast` (`float`, *optional*):
|
||||
Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear
|
||||
ramp function. If unspecified, it defaults to 32.
|
||||
`beta_slow` (`float`, *optional*):
|
||||
Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear
|
||||
ramp function. If unspecified, it defaults to 1.
|
||||
`short_factor` (`List[float]`, *optional*):
|
||||
Only used with 'longrope'. The scaling factor to be applied to short contexts (<
|
||||
`original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
|
||||
size divided by the number of attention heads divided by 2
|
||||
`long_factor` (`List[float]`, *optional*):
|
||||
Only used with 'longrope'. The scaling factor to be applied to long contexts (<
|
||||
`original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
|
||||
size divided by the number of attention heads divided by 2
|
||||
`low_freq_factor` (`float`, *optional*):
|
||||
Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE
|
||||
`high_freq_factor` (`float`, *optional*):
|
||||
Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE
|
||||
|
||||
```python
|
||||
>>> from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2_5_VLConfig
|
||||
|
||||
>>> # Initializing a Qwen2_5_VL style configuration
|
||||
>>> configuration = Qwen2_5_VLConfig()
|
||||
|
||||
>>> # Initializing a model from the Qwen2-VL-7B style configuration
|
||||
>>> model = Qwen2_5_VLForConditionalGeneration(configuration)
|
||||
|
||||
>>> # Accessing the model configuration
|
||||
>>> configuration = model.config
|
||||
```"""
|
||||
|
||||
model_type = "qwen2_5_vl"
|
||||
sub_configs = {"vision_config": Qwen2_5_VLVisionConfig}
|
||||
keys_to_ignore_at_inference = ["past_key_values"]
|
||||
# Default tensor parallel plan for base model `Qwen2_5_VL`
|
||||
base_model_tp_plan = {
|
||||
"layers.*.self_attn.q_proj": "colwise",
|
||||
"layers.*.self_attn.k_proj": "colwise",
|
||||
"layers.*.self_attn.v_proj": "colwise",
|
||||
"layers.*.self_attn.o_proj": "rowwise",
|
||||
"layers.*.mlp.gate_proj": "colwise",
|
||||
"layers.*.mlp.up_proj": "colwise",
|
||||
"layers.*.mlp.down_proj": "rowwise",
|
||||
}
|
||||
base_model_pp_plan = {
|
||||
"embed_tokens": (["input_ids"], ["inputs_embeds"]),
|
||||
"layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
|
||||
"norm": (["hidden_states"], ["hidden_states"]),
|
||||
sub_configs = {
|
||||
"vision_config": Qwen2_5_VLVisionConfig,
|
||||
"text_config": Qwen2_5_VLTextConfig,
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vocab_size=152064,
|
||||
hidden_size=8192,
|
||||
intermediate_size=29568,
|
||||
num_hidden_layers=80,
|
||||
num_attention_heads=64,
|
||||
num_key_value_heads=8,
|
||||
hidden_act="silu",
|
||||
max_position_embeddings=32768,
|
||||
initializer_range=0.02,
|
||||
rms_norm_eps=1e-05,
|
||||
use_cache=True,
|
||||
tie_word_embeddings=False,
|
||||
rope_theta=1000000.0,
|
||||
use_sliding_window=False,
|
||||
sliding_window=4096,
|
||||
max_window_layers=80,
|
||||
attention_dropout=0.0,
|
||||
vision_config=None,
|
||||
rope_scaling=None,
|
||||
num_experts=4,
|
||||
experts=None,
|
||||
dof_config=None,
|
||||
noise_scheduler=None,
|
||||
dim_inputs=(1536, 1536),
|
||||
attention_moe=False,
|
||||
mlp_moe=False,
|
||||
**kwargs,
|
||||
):
|
||||
if isinstance(vision_config, dict):
|
||||
self.vision_config = self.sub_configs["vision_config"](**vision_config)
|
||||
elif vision_config is None:
|
||||
self.vision_config = self.sub_configs["vision_config"]()
|
||||
def __getattr__(self, name):
|
||||
"""Keep legacy direct access to fields now owned by ``text_config``.
|
||||
|
||||
self.vocab_size = vocab_size
|
||||
self.max_position_embeddings = max_position_embeddings
|
||||
self.hidden_size = hidden_size
|
||||
self.intermediate_size = intermediate_size
|
||||
self.num_hidden_layers = num_hidden_layers
|
||||
self.num_attention_heads = num_attention_heads
|
||||
self.use_sliding_window = use_sliding_window
|
||||
self.sliding_window = sliding_window
|
||||
self.max_window_layers = max_window_layers
|
||||
self.layer_types = ["dense"] * num_hidden_layers
|
||||
|
||||
# for backward compatibility
|
||||
if num_key_value_heads is None:
|
||||
num_key_value_heads = num_attention_heads
|
||||
|
||||
self.num_key_value_heads = num_key_value_heads
|
||||
self.hidden_act = hidden_act
|
||||
self.initializer_range = initializer_range
|
||||
self.rms_norm_eps = rms_norm_eps
|
||||
self.use_cache = use_cache
|
||||
self.rope_theta = rope_theta
|
||||
self.attention_dropout = attention_dropout
|
||||
self.rope_scaling = rope_scaling
|
||||
|
||||
self.num_experts = num_experts
|
||||
self.experts = experts
|
||||
self.dof_config = dof_config
|
||||
self.noise_scheduler = noise_scheduler
|
||||
self.dim_inputs = tuple(dim_inputs)
|
||||
self.attention_moe = attention_moe
|
||||
self.mlp_moe = mlp_moe
|
||||
|
||||
if self.rope_scaling is not None and "type" in self.rope_scaling:
|
||||
if self.rope_scaling["type"] == "mrope":
|
||||
self.rope_scaling["type"] = "default"
|
||||
self.rope_scaling["rope_type"] = self.rope_scaling["type"]
|
||||
rope_config_validation(self, ignore_keys={"mrope_section"})
|
||||
|
||||
super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
|
||||
|
||||
@property
|
||||
def text_config(self):
|
||||
return self
|
||||
|
||||
|
||||
__all__ = ["Qwen2_5_VLConfig"]
|
||||
Wall-X historically used a flat config and accesses fields such as
|
||||
``hidden_size`` and ``num_experts`` directly. Forwarding unknown
|
||||
attributes preserves that API without duplicating the native config.
|
||||
"""
|
||||
text_config = self.__dict__.get("text_config")
|
||||
if name in _LEGACY_TEXT_ATTRIBUTES and text_config is not None and hasattr(text_config, name):
|
||||
return getattr(text_config, name)
|
||||
raise AttributeError(f"{type(self).__name__!s} has no attribute {name!r}")
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,208 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2025 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.
|
||||
|
||||
"""Wall-X vision attention backends.
|
||||
|
||||
Qwen2.5-VL's native non-Flash vision path splits a packed image sequence into
|
||||
Python-level chunks before calling attention. Wall-X batches many camera frames,
|
||||
so that path launches thousands of tiny attention operations per training step.
|
||||
This module keeps the native SDPA path as a portable fallback and adds a packed
|
||||
``torch.nn.attention.varlen`` path that consumes Qwen's existing ``cu_seqlens``
|
||||
metadata directly.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import inspect
|
||||
import logging
|
||||
from functools import lru_cache
|
||||
from typing import TYPE_CHECKING, Literal
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from lerobot.utils.import_utils import _transformers_available
|
||||
|
||||
if TYPE_CHECKING or _transformers_available:
|
||||
from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
|
||||
Qwen2_5_VLVisionAttention,
|
||||
apply_rotary_pos_emb_vision,
|
||||
)
|
||||
else:
|
||||
Qwen2_5_VLVisionAttention = nn.Module
|
||||
apply_rotary_pos_emb_vision = None
|
||||
|
||||
try:
|
||||
from torch.nn.attention.varlen import varlen_attn as _varlen_attn
|
||||
except ImportError: # torch<2.10
|
||||
_varlen_attn = None
|
||||
|
||||
_VARLEN_USES_WINDOW_SIZE = (
|
||||
_varlen_attn is not None and "window_size" in inspect.signature(_varlen_attn).parameters
|
||||
)
|
||||
|
||||
|
||||
VisionAttentionBackend = Literal["auto", "sdpa", "varlen"]
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@lru_cache
|
||||
def _log_resolved_backend(requested: str, resolved: str) -> None:
|
||||
logger.info("Wall-X vision attention backend: %s (requested: %s)", resolved, requested)
|
||||
|
||||
|
||||
def _varlen_unavailable_reason(
|
||||
hidden_states: torch.Tensor,
|
||||
position_embeddings: tuple[torch.Tensor, torch.Tensor] | None,
|
||||
) -> str | None:
|
||||
if _varlen_attn is None:
|
||||
return "torch.nn.attention.varlen is unavailable (PyTorch 2.10 or newer is required)"
|
||||
if position_embeddings is None:
|
||||
return "precomputed vision position embeddings were not provided"
|
||||
if hidden_states.device.type != "cuda" or torch.version.cuda is None:
|
||||
return "packed varlen attention requires an NVIDIA CUDA device"
|
||||
if hidden_states.dtype not in {torch.float16, torch.bfloat16}:
|
||||
return f"packed varlen attention requires float16 or bfloat16 inputs, got {hidden_states.dtype}"
|
||||
major, _minor = torch.cuda.get_device_capability(hidden_states.device)
|
||||
if major < 8:
|
||||
return "packed varlen attention requires an NVIDIA Ampere GPU or newer"
|
||||
return None
|
||||
|
||||
|
||||
def _supports_varlen_attention(
|
||||
hidden_states: torch.Tensor,
|
||||
position_embeddings: tuple[torch.Tensor, torch.Tensor] | None,
|
||||
) -> bool:
|
||||
return _varlen_unavailable_reason(hidden_states, position_embeddings) is None
|
||||
|
||||
|
||||
class WallXVisionAttention(Qwen2_5_VLVisionAttention):
|
||||
"""Qwen2.5-VL vision attention with packed varlen and native SDPA fallback."""
|
||||
|
||||
def __init__(self, config, backend: VisionAttentionBackend):
|
||||
super().__init__(config)
|
||||
self.wallx_backend = backend
|
||||
self._resolved_backend_key = None
|
||||
self._resolved_backend = None
|
||||
|
||||
def _resolve_backend(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
position_embeddings: tuple[torch.Tensor, torch.Tensor] | None,
|
||||
) -> str:
|
||||
key = (
|
||||
hidden_states.device.type,
|
||||
hidden_states.device.index,
|
||||
hidden_states.dtype,
|
||||
position_embeddings is not None,
|
||||
)
|
||||
if self._resolved_backend_key == key:
|
||||
return self._resolved_backend
|
||||
|
||||
use_varlen = self.wallx_backend != "sdpa" and _supports_varlen_attention(
|
||||
hidden_states, position_embeddings
|
||||
)
|
||||
if self.wallx_backend == "varlen" and not use_varlen:
|
||||
reason = _varlen_unavailable_reason(hidden_states, position_embeddings)
|
||||
raise RuntimeError(f"Wall-X vision_attn_implementation='varlen' cannot be used: {reason}")
|
||||
|
||||
resolved_backend = "varlen" if use_varlen else "sdpa"
|
||||
self._resolved_backend_key = key
|
||||
self._resolved_backend = resolved_backend
|
||||
_log_resolved_backend(self.wallx_backend, resolved_backend)
|
||||
return resolved_backend
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
cu_seqlens: torch.Tensor,
|
||||
rotary_pos_emb: torch.Tensor | None = None,
|
||||
position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
|
||||
**kwargs,
|
||||
) -> torch.Tensor:
|
||||
del rotary_pos_emb
|
||||
|
||||
if self._resolve_backend(hidden_states, position_embeddings) == "sdpa":
|
||||
return super().forward(
|
||||
hidden_states=hidden_states,
|
||||
cu_seqlens=cu_seqlens,
|
||||
position_embeddings=position_embeddings,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
seq_length = hidden_states.shape[0]
|
||||
query_states, key_states, value_states = (
|
||||
self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0)
|
||||
)
|
||||
|
||||
cos, sin = position_embeddings
|
||||
query_states, key_states = apply_rotary_pos_emb_vision(
|
||||
query_states,
|
||||
key_states,
|
||||
cos,
|
||||
sin,
|
||||
)
|
||||
|
||||
if cu_seqlens.dtype != torch.int32:
|
||||
cu_seqlens = cu_seqlens.to(dtype=torch.int32)
|
||||
max_seqlen = int((cu_seqlens[1:] - cu_seqlens[:-1]).max().item())
|
||||
varlen_kwargs = {"scale": self.scaling}
|
||||
if _VARLEN_USES_WINDOW_SIZE:
|
||||
varlen_kwargs["window_size"] = (-1, -1)
|
||||
else: # Stable PyTorch 2.10 API; pre-release variants used window_size.
|
||||
varlen_kwargs["is_causal"] = False
|
||||
attn_output = _varlen_attn(
|
||||
query_states,
|
||||
key_states,
|
||||
value_states,
|
||||
cu_seqlens,
|
||||
cu_seqlens,
|
||||
max_seqlen,
|
||||
max_seqlen,
|
||||
**varlen_kwargs,
|
||||
)
|
||||
attn_output = attn_output.reshape(seq_length, -1).contiguous()
|
||||
return self.proj(attn_output)
|
||||
|
||||
|
||||
def configure_wall_x_vision_attention(
|
||||
vision_model: nn.Module,
|
||||
backend: VisionAttentionBackend,
|
||||
) -> None:
|
||||
"""Install Wall-X's scoped packed attention without changing checkpoint keys."""
|
||||
if backend == "sdpa":
|
||||
_log_resolved_backend(backend, "sdpa")
|
||||
return
|
||||
if backend == "varlen" and _varlen_attn is None:
|
||||
raise RuntimeError(
|
||||
"Wall-X vision_attn_implementation='varlen' requires torch.nn.attention.varlen "
|
||||
"from PyTorch 2.10 or newer"
|
||||
)
|
||||
if backend == "auto" and _varlen_attn is None:
|
||||
_log_resolved_backend(backend, "sdpa")
|
||||
return
|
||||
|
||||
for block in vision_model.blocks:
|
||||
previous_attention = block.attn
|
||||
replacement = WallXVisionAttention(previous_attention.config, backend=backend)
|
||||
replacement.to(
|
||||
device=previous_attention.qkv.weight.device,
|
||||
dtype=previous_attention.qkv.weight.dtype,
|
||||
)
|
||||
replacement.load_state_dict(previous_attention.state_dict(), strict=True)
|
||||
replacement.train(previous_attention.training)
|
||||
block.attn = replacement
|
||||
@@ -116,6 +116,7 @@ def preprocesser_call(
|
||||
images: list | Any | None = None,
|
||||
text: str | list[str] | None = None,
|
||||
videos: list | Any | None = None,
|
||||
device: torch.device | str | None = None,
|
||||
padding: bool | str = False,
|
||||
truncation: bool | None = None,
|
||||
max_length: int | None = None,
|
||||
@@ -134,6 +135,7 @@ def preprocesser_call(
|
||||
images: Input images (PIL, numpy arrays, or torch tensors)
|
||||
text: Text or list of texts to tokenize
|
||||
videos: Input videos (numpy arrays or torch tensors)
|
||||
device: Device on which image/video preprocessing should run
|
||||
padding: Whether to pad sequences to same length
|
||||
truncation: Whether to truncate sequences longer than max_length
|
||||
max_length: Maximum length for truncation/padding
|
||||
@@ -151,7 +153,11 @@ def preprocesser_call(
|
||||
"""
|
||||
# Process image inputs
|
||||
if images is not None and len(images) > 0:
|
||||
image_inputs = processor.image_processor(images=images, return_tensors=return_tensors)
|
||||
image_inputs = processor.image_processor(
|
||||
images=images,
|
||||
return_tensors=return_tensors,
|
||||
device=device,
|
||||
)
|
||||
image_grid_thw = image_inputs["image_grid_thw"]
|
||||
else:
|
||||
image_inputs = {}
|
||||
@@ -159,7 +165,11 @@ def preprocesser_call(
|
||||
|
||||
# Process video inputs
|
||||
if videos is not None:
|
||||
videos_inputs = processor.image_processor(videos=videos, return_tensors=return_tensors)
|
||||
videos_inputs = processor.image_processor(
|
||||
videos=videos,
|
||||
return_tensors=return_tensors,
|
||||
device=device,
|
||||
)
|
||||
video_grid_thw = videos_inputs["video_grid_thw"]
|
||||
else:
|
||||
videos_inputs = {}
|
||||
@@ -413,10 +423,7 @@ def get_task_instruction(
|
||||
}
|
||||
)
|
||||
|
||||
if priority_order is not None:
|
||||
priority_order = OrderedDict(priority_order)
|
||||
else:
|
||||
priority_order = default_priority_order
|
||||
priority_order = OrderedDict(priority_order) if priority_order is not None else default_priority_order
|
||||
|
||||
got_instruction = False
|
||||
task_instruction = ""
|
||||
@@ -424,9 +431,8 @@ def get_task_instruction(
|
||||
# Sample instruction components based on priority probabilities
|
||||
for key, prob in priority_order.items():
|
||||
if key in frame_instruction_info and frame_instruction_info[key] != "":
|
||||
if got_instruction:
|
||||
if random.random() >= prob:
|
||||
continue
|
||||
if got_instruction and random.random() >= prob:
|
||||
continue
|
||||
|
||||
task_instruction += f"\n{frame_instruction_info[key]}"
|
||||
got_instruction = True
|
||||
@@ -538,10 +544,7 @@ def img_key_mapping(img_keys: list[str]) -> list[str]:
|
||||
if key in CAMERA_NAME_MAPPING:
|
||||
key = CAMERA_NAME_MAPPING[key]
|
||||
else:
|
||||
if "view" in key:
|
||||
key = key.replace("_", " ")
|
||||
else:
|
||||
key = key + " view"
|
||||
key = key.replace("_", " ") if "view" in key else key + " view"
|
||||
processed_img_keys.append(key)
|
||||
return processed_img_keys
|
||||
|
||||
|
||||
@@ -22,19 +22,14 @@ import torch
|
||||
|
||||
from lerobot.configs import PipelineFeatureType, PolicyFeature
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
ObservationProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
ProcessorStep,
|
||||
ProcessorStepRegistry,
|
||||
RenameObservationsProcessorStep,
|
||||
TokenizerProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
make_default_policy_processor_steps,
|
||||
make_policy_processor_pipelines,
|
||||
)
|
||||
from lerobot.types import EnvTransition, TransitionKey
|
||||
from lerobot.utils.constants import (
|
||||
@@ -42,8 +37,6 @@ from lerobot.utils.constants import (
|
||||
OBS_IMAGES,
|
||||
OBS_PREFIX,
|
||||
OBS_STATE,
|
||||
POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
)
|
||||
|
||||
from .configuration_xvla import XVLAConfig
|
||||
@@ -61,10 +54,11 @@ def make_xvla_pre_post_processors(
|
||||
Build the LeRobot processor pipelines for XVLA.
|
||||
"""
|
||||
|
||||
features = {**config.input_features, **config.output_features}
|
||||
steps = make_default_policy_processor_steps(config, dataset_stats)
|
||||
|
||||
input_steps = [
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
steps.rename_observations,
|
||||
steps.add_batch_dim,
|
||||
TokenizerProcessorStep(
|
||||
tokenizer_name=config.tokenizer_name,
|
||||
max_length=config.tokenizer_max_length,
|
||||
@@ -74,32 +68,15 @@ def make_xvla_pre_post_processors(
|
||||
XVLAImageToFloatProcessorStep(),
|
||||
XVLAImageNetNormalizeProcessorStep(),
|
||||
XVLAAddDomainIdProcessorStep(),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
NormalizerProcessorStep(
|
||||
features=features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
steps.to_device,
|
||||
steps.normalize,
|
||||
]
|
||||
output_steps = [
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features,
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
steps.unnormalize,
|
||||
steps.to_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,
|
||||
),
|
||||
)
|
||||
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
|
||||
|
||||
|
||||
# Custom XVLA processor steps
|
||||
|
||||
@@ -42,10 +42,14 @@ from .delta_action_processor import MapDeltaActionToRobotActionStep, MapTensorTo
|
||||
from .device_processor import DeviceProcessorStep
|
||||
from .env_processor import IsaaclabArenaProcessorStep, LiberoProcessorStep
|
||||
from .factory import (
|
||||
DefaultPolicyProcessorSteps,
|
||||
make_default_policy_processor_steps,
|
||||
make_default_pre_post_processors,
|
||||
make_default_processors,
|
||||
make_default_robot_action_processor,
|
||||
make_default_robot_observation_processor,
|
||||
make_default_teleop_action_processor,
|
||||
make_policy_processor_pipelines,
|
||||
)
|
||||
from .gym_action_processor import (
|
||||
Numpy2TorchActionProcessorStep,
|
||||
@@ -129,10 +133,14 @@ __all__ = [
|
||||
"ImageCropResizeProcessorStep",
|
||||
"InfoProcessorStep",
|
||||
"InterventionActionProcessorStep",
|
||||
"DefaultPolicyProcessorSteps",
|
||||
"make_default_policy_processor_steps",
|
||||
"make_default_pre_post_processors",
|
||||
"make_default_processors",
|
||||
"make_default_teleop_action_processor",
|
||||
"make_default_robot_action_processor",
|
||||
"make_default_robot_observation_processor",
|
||||
"make_policy_processor_pipelines",
|
||||
"AbsoluteActionsProcessorStep",
|
||||
"RelativeActionsProcessorStep",
|
||||
"MapDeltaActionToRobotActionStep",
|
||||
|
||||
@@ -175,6 +175,9 @@ class AddBatchDimensionComplementaryDataStep(ComplementaryDataProcessorStep):
|
||||
if isinstance(task_index_value, Tensor) and task_index_value.dim() == 0:
|
||||
complementary_data["task_index"] = task_index_value.unsqueeze(0)
|
||||
|
||||
complementary_data.pop("language_persistent", None)
|
||||
complementary_data.pop("language_events", None)
|
||||
|
||||
if "messages" in complementary_data:
|
||||
messages = complementary_data["messages"]
|
||||
if isinstance(messages, list) and (not messages or isinstance(messages[0], dict)):
|
||||
|
||||
@@ -14,15 +14,33 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from lerobot.types import RobotAction, RobotObservation
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from lerobot.configs.policies import PreTrainedConfig
|
||||
from lerobot.types import PolicyAction, RobotAction, RobotObservation
|
||||
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
|
||||
|
||||
from .batch_processor import AddBatchDimensionProcessorStep
|
||||
from .converters import (
|
||||
observation_to_transition,
|
||||
policy_action_to_transition,
|
||||
robot_action_observation_to_transition,
|
||||
transition_to_observation,
|
||||
transition_to_policy_action,
|
||||
transition_to_robot_action,
|
||||
)
|
||||
from .pipeline import IdentityProcessorStep, RobotProcessorPipeline
|
||||
from .device_processor import DeviceProcessorStep
|
||||
from .normalize_processor import NormalizerProcessorStep, UnnormalizerProcessorStep
|
||||
from .pipeline import (
|
||||
IdentityProcessorStep,
|
||||
PolicyProcessorPipeline,
|
||||
ProcessorStep,
|
||||
RobotProcessorPipeline,
|
||||
)
|
||||
from .rename_processor import RenameObservationsProcessorStep
|
||||
|
||||
|
||||
def make_default_teleop_action_processor() -> RobotProcessorPipeline[
|
||||
@@ -61,3 +79,97 @@ def make_default_processors():
|
||||
robot_action_processor = make_default_robot_action_processor()
|
||||
robot_observation_processor = make_default_robot_observation_processor()
|
||||
return (teleop_action_processor, robot_action_processor, robot_observation_processor)
|
||||
|
||||
|
||||
@dataclass
|
||||
class DefaultPolicyProcessorSteps:
|
||||
"""The canonical processor steps shared by most policies' pre/post pipelines.
|
||||
|
||||
Policies compose these in their own order (step ORDER is a Hub-serialized contract
|
||||
and intentionally stays explicit per policy) and interleave their custom steps.
|
||||
"""
|
||||
|
||||
rename_observations: RenameObservationsProcessorStep
|
||||
add_batch_dim: AddBatchDimensionProcessorStep
|
||||
to_device: DeviceProcessorStep
|
||||
normalize: NormalizerProcessorStep
|
||||
unnormalize: UnnormalizerProcessorStep
|
||||
to_cpu: DeviceProcessorStep
|
||||
|
||||
|
||||
def make_default_policy_processor_steps(
|
||||
config: PreTrainedConfig,
|
||||
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
|
||||
*,
|
||||
normalizer_device: torch.device | str | None = None,
|
||||
) -> DefaultPolicyProcessorSteps:
|
||||
"""Construct the canonical policy processor steps from a policy config.
|
||||
|
||||
Args:
|
||||
config: A `PreTrainedConfig` providing `device`, `input_features`,
|
||||
`output_features` and `normalization_mapping`.
|
||||
dataset_stats: Dataset statistics used for (un)normalization.
|
||||
normalizer_device: Device passed to `NormalizerProcessorStep` (some policies pin
|
||||
their normalization stats to the policy device; most leave it unset).
|
||||
"""
|
||||
return DefaultPolicyProcessorSteps(
|
||||
rename_observations=RenameObservationsProcessorStep(rename_map={}),
|
||||
add_batch_dim=AddBatchDimensionProcessorStep(),
|
||||
to_device=DeviceProcessorStep(device=config.device),
|
||||
normalize=NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
device=normalizer_device,
|
||||
),
|
||||
unnormalize=UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
to_cpu=DeviceProcessorStep(device="cpu"),
|
||||
)
|
||||
|
||||
|
||||
def make_policy_processor_pipelines(
|
||||
input_steps: list[ProcessorStep],
|
||||
output_steps: list[ProcessorStep],
|
||||
) -> tuple[
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction],
|
||||
]:
|
||||
"""Wrap pre/post step lists into the canonical policy pipeline pair.
|
||||
|
||||
Uses the standard pipeline names (which determine the serialized JSON filenames on
|
||||
the Hub) and the standard policy-action converters on the postprocessor.
|
||||
"""
|
||||
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 make_default_pre_post_processors(
|
||||
config: PreTrainedConfig,
|
||||
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
|
||||
*,
|
||||
normalizer_device: torch.device | str | None = None,
|
||||
) -> tuple[
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction],
|
||||
]:
|
||||
"""The pure-scaffold policy pipeline pair: Rename -> Batch -> Device -> Normalize,
|
||||
and Unnormalize -> Device(cpu). Policies with custom steps or a different step order
|
||||
compose `make_default_policy_processor_steps` themselves instead.
|
||||
"""
|
||||
s = make_default_policy_processor_steps(config, dataset_stats, normalizer_device=normalizer_device)
|
||||
return make_policy_processor_pipelines(
|
||||
input_steps=[s.rename_observations, s.add_batch_dim, s.to_device, s.normalize],
|
||||
output_steps=[s.unnormalize, s.to_cpu],
|
||||
)
|
||||
|
||||
@@ -41,7 +41,7 @@ from pathlib import Path
|
||||
from typing import Any, TypedDict, TypeVar, cast
|
||||
|
||||
import torch
|
||||
from huggingface_hub import hf_hub_download, snapshot_download
|
||||
from huggingface_hub import hf_hub_download
|
||||
from safetensors.torch import load_file, save_file
|
||||
|
||||
from lerobot.configs import PipelineFeatureType, PolicyFeature
|
||||
@@ -205,10 +205,6 @@ class ProcessorStep(ABC):
|
||||
"""
|
||||
return None
|
||||
|
||||
def save_artifacts(self, save_directory: Path) -> dict[str, str]:
|
||||
"""Save non-tensor assets and map constructor arguments to relative paths."""
|
||||
return {}
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Resets the internal state of the processor step, if any."""
|
||||
return None
|
||||
@@ -553,22 +549,6 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
|
||||
pipeline_config = self.get_config()
|
||||
pipeline_state_dict = self.state_dict()
|
||||
|
||||
for processor_step, step_entry in zip(self.steps, pipeline_config["steps"], strict=True):
|
||||
artifacts = processor_step.save_artifacts(save_directory)
|
||||
if artifacts:
|
||||
for config_key, relative_path in artifacts.items():
|
||||
artifact_path = Path(relative_path)
|
||||
if artifact_path.is_absolute() or ".." in artifact_path.parts:
|
||||
raise ValueError(
|
||||
f"Processor artifact path must be relative to the checkpoint: {relative_path!r}"
|
||||
)
|
||||
if not (save_directory / artifact_path).exists():
|
||||
raise FileNotFoundError(
|
||||
f"Processor step did not save declared artifact '{relative_path}'"
|
||||
)
|
||||
step_entry["config"][config_key] = artifact_path.as_posix()
|
||||
step_entry["artifacts"] = artifacts
|
||||
|
||||
for state_key, step_state_dict in pipeline_state_dict.items():
|
||||
state_filename = f"{state_key}.safetensors"
|
||||
save_file(step_state_dict, save_directory / state_filename)
|
||||
@@ -751,12 +731,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
|
||||
|
||||
# 3. Build steps with overrides
|
||||
steps, validated_overrides = cls._build_steps_with_overrides(
|
||||
loaded_config,
|
||||
overrides or {},
|
||||
model_id,
|
||||
base_path,
|
||||
config_filename,
|
||||
hub_download_kwargs,
|
||||
loaded_config, overrides or {}, model_id, base_path, hub_download_kwargs
|
||||
)
|
||||
|
||||
# 4. Validate that all overrides were used
|
||||
@@ -945,7 +920,6 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
|
||||
overrides: dict[str, Any],
|
||||
model_id: str,
|
||||
base_path: Path | None,
|
||||
config_filename: str,
|
||||
hub_download_kwargs: dict[str, Any],
|
||||
) -> tuple[list[ProcessorStep], set[str]]:
|
||||
"""Build all processor steps with overrides and state loading.
|
||||
@@ -998,67 +972,13 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
|
||||
ImportError: If a step class cannot be imported or found in registry
|
||||
ValueError: If a step cannot be instantiated with its configuration
|
||||
"""
|
||||
loaded_config = deepcopy(loaded_config)
|
||||
cls._resolve_artifact_paths(
|
||||
loaded_config,
|
||||
model_id,
|
||||
base_path,
|
||||
config_filename,
|
||||
hub_download_kwargs,
|
||||
)
|
||||
steps, remaining_override_keys = cls._build_steps_from_config(loaded_config, overrides)
|
||||
|
||||
for step_instance, step_entry in zip(steps, loaded_config["steps"], strict=True):
|
||||
cls._load_step_state(
|
||||
step_instance,
|
||||
step_entry,
|
||||
model_id,
|
||||
base_path,
|
||||
config_filename,
|
||||
hub_download_kwargs,
|
||||
)
|
||||
cls._load_step_state(step_instance, step_entry, model_id, base_path, hub_download_kwargs)
|
||||
|
||||
return steps, remaining_override_keys
|
||||
|
||||
@classmethod
|
||||
def _resolve_artifact_paths(
|
||||
cls,
|
||||
loaded_config: dict[str, Any],
|
||||
model_id: str,
|
||||
base_path: Path | None,
|
||||
config_filename: str,
|
||||
hub_download_kwargs: dict[str, Any],
|
||||
) -> None:
|
||||
"""Resolve declared relative processor artifacts before step construction."""
|
||||
is_local = Path(model_id).is_dir() or Path(model_id).is_file()
|
||||
|
||||
for step_entry in loaded_config["steps"]:
|
||||
artifacts = step_entry.get("artifacts", {})
|
||||
for config_key, relative_path in artifacts.items():
|
||||
artifact_path = Path(relative_path)
|
||||
if artifact_path.is_absolute() or ".." in artifact_path.parts:
|
||||
raise ValueError(
|
||||
f"Processor artifact path must be relative to the checkpoint: {relative_path!r}"
|
||||
)
|
||||
|
||||
resolved_path = base_path / artifact_path if base_path is not None else artifact_path
|
||||
if not resolved_path.exists() and not is_local:
|
||||
repository_path = Path(config_filename).parent / artifact_path
|
||||
snapshot_download(
|
||||
repo_id=model_id,
|
||||
repo_type="model",
|
||||
allow_patterns=f"{repository_path.as_posix()}/**",
|
||||
**hub_download_kwargs,
|
||||
)
|
||||
|
||||
if not resolved_path.exists():
|
||||
step_name = step_entry.get("registry_name", step_entry.get("class", "unknown"))
|
||||
raise FileNotFoundError(
|
||||
f"Missing processor artifact '{relative_path}' for step '{step_name}' "
|
||||
f"next to '{config_filename}'. Checkpoint artifacts are incomplete."
|
||||
)
|
||||
step_entry["config"][config_key] = str(resolved_path)
|
||||
|
||||
@classmethod
|
||||
def _build_steps_from_config(
|
||||
cls,
|
||||
@@ -1218,7 +1138,6 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
|
||||
step_entry: dict[str, Any],
|
||||
model_id: str,
|
||||
base_path: Path | None,
|
||||
config_filename: str,
|
||||
hub_download_kwargs: dict[str, Any],
|
||||
) -> None:
|
||||
"""Load state dictionary for a processor step if available.
|
||||
@@ -1276,7 +1195,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
|
||||
# Download from Hub
|
||||
state_path = hf_hub_download(
|
||||
repo_id=model_id,
|
||||
filename=(Path(config_filename).parent / state_filename).as_posix(),
|
||||
filename=state_filename,
|
||||
repo_type="model",
|
||||
**hub_download_kwargs,
|
||||
)
|
||||
|
||||
@@ -16,7 +16,7 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import asdict, dataclass
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from lerobot.configs import PipelineFeatureType, PolicyFeature
|
||||
@@ -32,18 +32,17 @@ from .pipeline import ProcessorStep, ProcessorStepRegistry
|
||||
@dataclass
|
||||
@ProcessorStepRegistry.register(name="render_messages_processor")
|
||||
class RenderMessagesStep(ProcessorStep):
|
||||
"""Render language columns into recipe-defined messages and supervision metadata."""
|
||||
"""Processor step that turns raw language columns into rendered chat messages.
|
||||
|
||||
Reads ``language_persistent`` and ``language_events`` from the transition's
|
||||
complementary data, renders them through ``recipe`` at the sample timestamp,
|
||||
and replaces the raw columns with the resulting ``messages`` /
|
||||
``message_streams`` / ``target_message_indices`` keys.
|
||||
"""
|
||||
|
||||
recipe: TrainingRecipe
|
||||
dataset_ctx: Any | None = None
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if isinstance(self.recipe, dict):
|
||||
self.recipe = TrainingRecipe.from_dict(self.recipe)
|
||||
|
||||
def get_config(self) -> dict[str, Any]:
|
||||
return {"recipe": asdict(self.recipe)}
|
||||
|
||||
def __call__(self, transition: EnvTransition) -> EnvTransition | None:
|
||||
"""Render messages for a single transition; return ``None`` to drop it."""
|
||||
complementary_data = transition.get(TransitionKey.COMPLEMENTARY_DATA) or {}
|
||||
@@ -51,17 +50,7 @@ class RenderMessagesStep(ProcessorStep):
|
||||
events = complementary_data.get(LANGUAGE_EVENTS) or []
|
||||
|
||||
if not persistent and not events:
|
||||
rendered = _fallback_low_level_render(complementary_data.get("task"))
|
||||
if rendered is None:
|
||||
return transition
|
||||
new_transition = transition.copy()
|
||||
new_complementary_data = dict(new_transition.get(TransitionKey.COMPLEMENTARY_DATA) or {})
|
||||
new_complementary_data.update(rendered)
|
||||
new_transition[TransitionKey.COMPLEMENTARY_DATA] = new_complementary_data
|
||||
return new_transition
|
||||
|
||||
if _is_batched_language(persistent) or _is_batched_language(events):
|
||||
return self._call_batch(transition, complementary_data, persistent, events)
|
||||
return transition
|
||||
|
||||
timestamp = complementary_data.get("timestamp")
|
||||
if timestamp is None:
|
||||
@@ -78,147 +67,18 @@ class RenderMessagesStep(ProcessorStep):
|
||||
dataset_ctx=self.dataset_ctx,
|
||||
)
|
||||
if rendered is None:
|
||||
rendered = _fallback_low_level_render(complementary_data.get("task"))
|
||||
if rendered is None:
|
||||
return None
|
||||
return None
|
||||
|
||||
new_transition = transition.copy()
|
||||
new_complementary_data = dict(new_transition.get(TransitionKey.COMPLEMENTARY_DATA) or {})
|
||||
new_complementary_data = dict(complementary_data)
|
||||
new_complementary_data.pop(LANGUAGE_PERSISTENT, None)
|
||||
new_complementary_data.pop(LANGUAGE_EVENTS, None)
|
||||
new_complementary_data.update(rendered)
|
||||
new_transition[TransitionKey.COMPLEMENTARY_DATA] = new_complementary_data
|
||||
return new_transition
|
||||
|
||||
def _call_batch(
|
||||
self,
|
||||
transition: EnvTransition,
|
||||
complementary_data: dict[str, Any],
|
||||
persistent_batch: list,
|
||||
events_batch: list,
|
||||
) -> EnvTransition | None:
|
||||
timestamp = complementary_data.get("timestamp")
|
||||
if timestamp is None:
|
||||
raise KeyError("RenderMessagesStep requires sample timestamp in complementary data.")
|
||||
|
||||
batch_size = max(len(persistent_batch), len(events_batch))
|
||||
messages: list[list[dict[str, Any]]] = []
|
||||
message_streams: list[list[str | None]] = []
|
||||
target_message_indices: list[list[int]] = []
|
||||
keep_indices: list[int] = []
|
||||
|
||||
for i in range(batch_size):
|
||||
rendered = render_sample(
|
||||
recipe=self.recipe,
|
||||
persistent=persistent_batch[i] if i < len(persistent_batch) else [],
|
||||
events=events_batch[i] if i < len(events_batch) else [],
|
||||
t=_batch_value(timestamp, i),
|
||||
sample_idx=int(_batch_value(complementary_data.get("index", 0), i)),
|
||||
task=_batch_value(complementary_data.get("task"), i),
|
||||
dataset_ctx=self.dataset_ctx,
|
||||
)
|
||||
if rendered is None:
|
||||
rendered = _fallback_low_level_render(_batch_value(complementary_data.get("task"), i))
|
||||
if rendered is None:
|
||||
continue
|
||||
keep_indices.append(i)
|
||||
messages.append(rendered["messages"])
|
||||
message_streams.append(rendered["message_streams"])
|
||||
target_message_indices.append(rendered["target_message_indices"])
|
||||
|
||||
if not messages:
|
||||
return None
|
||||
|
||||
new_transition = (
|
||||
_select_batch_indices(transition, keep_indices)
|
||||
if len(keep_indices) != batch_size
|
||||
else transition.copy()
|
||||
)
|
||||
new_complementary_data = dict(new_transition.get(TransitionKey.COMPLEMENTARY_DATA) or {})
|
||||
new_complementary_data.pop(LANGUAGE_PERSISTENT, None)
|
||||
new_complementary_data.pop(LANGUAGE_EVENTS, None)
|
||||
new_complementary_data["messages"] = messages
|
||||
new_complementary_data["message_streams"] = message_streams
|
||||
new_complementary_data["target_message_indices"] = target_message_indices
|
||||
new_transition[TransitionKey.COMPLEMENTARY_DATA] = new_complementary_data
|
||||
return new_transition
|
||||
|
||||
def transform_features(
|
||||
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
||||
"""Pass features through unchanged; rendering only touches complementary data."""
|
||||
return features
|
||||
|
||||
|
||||
def _scalar(value: Any) -> float | int:
|
||||
"""Unwrap a tensor/array/single-element list into a Python scalar."""
|
||||
if hasattr(value, "item"):
|
||||
return value.item()
|
||||
if isinstance(value, list):
|
||||
if len(value) != 1:
|
||||
raise ValueError(f"Expected a scalar, got list of length {len(value)}: {value!r}")
|
||||
return _scalar(value[0])
|
||||
return value
|
||||
|
||||
|
||||
def _is_batched_language(value: Any) -> bool:
|
||||
return isinstance(value, list) and bool(value) and isinstance(value[0], list)
|
||||
|
||||
|
||||
def _batch_value(value: Any, index: int) -> Any:
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, list):
|
||||
return value[index]
|
||||
if hasattr(value, "ndim") and value.ndim > 0:
|
||||
return _scalar(value[index])
|
||||
return _scalar(value)
|
||||
|
||||
|
||||
def _select_batch_indices(transition: EnvTransition, indices: list[int]) -> EnvTransition:
|
||||
selected = transition.copy()
|
||||
for key in (TransitionKey.OBSERVATION, TransitionKey.COMPLEMENTARY_DATA):
|
||||
data = selected.get(key)
|
||||
if isinstance(data, dict):
|
||||
selected[key] = {k: _select_value(v, indices) for k, v in data.items()}
|
||||
action = selected.get(TransitionKey.ACTION)
|
||||
if action is not None:
|
||||
selected[TransitionKey.ACTION] = _select_value(action, indices)
|
||||
return selected
|
||||
|
||||
|
||||
def _select_value(value: Any, indices: list[int]) -> Any:
|
||||
if isinstance(value, list) and len(value) >= len(indices):
|
||||
return [value[i] for i in indices]
|
||||
if hasattr(value, "index_select") and hasattr(value, "new_tensor") and getattr(value, "ndim", 0) > 0:
|
||||
return value.index_select(0, value.new_tensor(indices).long())
|
||||
return value
|
||||
|
||||
|
||||
def _fallback_low_level_render(task: Any) -> dict[str, Any] | None:
|
||||
"""Keep action-only samples trainable when no recipe branch matches."""
|
||||
if hasattr(task, "item"):
|
||||
task = task.item()
|
||||
if isinstance(task, list):
|
||||
messages = []
|
||||
message_streams = []
|
||||
target_message_indices = []
|
||||
for t in task:
|
||||
rendered = _fallback_low_level_render(t)
|
||||
if rendered is None:
|
||||
return None
|
||||
messages.append(rendered["messages"])
|
||||
message_streams.append(rendered["message_streams"])
|
||||
target_message_indices.append(rendered["target_message_indices"])
|
||||
return {
|
||||
"messages": messages,
|
||||
"message_streams": message_streams,
|
||||
"target_message_indices": target_message_indices,
|
||||
}
|
||||
if not isinstance(task, str) or not task:
|
||||
return None
|
||||
return {
|
||||
"messages": [{"role": "user", "content": task}],
|
||||
"message_streams": ["low_level"],
|
||||
"target_message_indices": [],
|
||||
}
|
||||
|
||||
@@ -25,7 +25,6 @@ from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import torch
|
||||
@@ -33,7 +32,6 @@ import torch
|
||||
from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
|
||||
from lerobot.types import EnvTransition, RobotObservation, TransitionKey
|
||||
from lerobot.utils.constants import (
|
||||
ACTION_CODE_TOKEN_MASK,
|
||||
ACTION_TOKEN_MASK,
|
||||
ACTION_TOKENS,
|
||||
OBS_LANGUAGE_ATTENTION_MASK,
|
||||
@@ -351,7 +349,6 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
|
||||
max_action_tokens: int = 256
|
||||
fast_skip_tokens: int = 128
|
||||
paligemma_tokenizer_name: str = "google/paligemma-3b-pt-224"
|
||||
allow_truncation: bool = True
|
||||
# Internal tokenizer instance (not part of the config)
|
||||
action_tokenizer: Any = field(default=None, init=False, repr=False)
|
||||
_paligemma_tokenizer: Any = field(default=None, init=False, repr=False)
|
||||
@@ -415,15 +412,14 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
|
||||
# During inference, no action is available, skip tokenization
|
||||
return new_transition
|
||||
|
||||
# Tokenize and get masks for the full formatted sequence and the discrete action codes.
|
||||
tokens, mask, code_mask = self._tokenize_action(action)
|
||||
# Tokenize and get both tokens and mask
|
||||
tokens, mask = self._tokenize_action(action)
|
||||
|
||||
# Store mask in complementary data
|
||||
complementary_data = new_transition.get(TransitionKey.COMPLEMENTARY_DATA, {})
|
||||
if complementary_data is None:
|
||||
complementary_data = {}
|
||||
complementary_data[ACTION_TOKEN_MASK] = mask
|
||||
complementary_data[ACTION_CODE_TOKEN_MASK] = code_mask
|
||||
complementary_data[ACTION_TOKENS] = tokens
|
||||
new_transition[TransitionKey.COMPLEMENTARY_DATA] = complementary_data
|
||||
return new_transition
|
||||
@@ -434,7 +430,7 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
|
||||
"""
|
||||
return self._paligemma_tokenizer.vocab_size - 1 - self.fast_skip_tokens - tokens
|
||||
|
||||
def _tokenize_action(self, action: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
def _tokenize_action(self, action: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Tokenizes the action tensor and creates a mask.
|
||||
|
||||
@@ -463,7 +459,6 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
|
||||
# The fast tokenizer expects action data and returns token IDs
|
||||
tokens_list = []
|
||||
masks_list = []
|
||||
code_masks_list = []
|
||||
|
||||
for i in range(batch_size):
|
||||
# Tokenize single action (move to CPU first as tokenizer uses scipy which requires numpy)
|
||||
@@ -481,82 +476,65 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
|
||||
if tokens.dim() > 1:
|
||||
tokens = tokens.flatten()
|
||||
|
||||
action_code_tokens = self._act_tokens_to_paligemma_tokens(tokens)
|
||||
bos_id = self._paligemma_tokenizer.bos_token_id
|
||||
prompt_tokens = torch.tensor(
|
||||
self._paligemma_tokenizer.encode("Action: ", add_special_tokens=False),
|
||||
device=action.device,
|
||||
)
|
||||
end_tokens = torch.tensor(self._paligemma_tokenizer.encode("|"), device=action.device)
|
||||
|
||||
code_start = 1 + len(prompt_tokens)
|
||||
code_end = code_start + len(action_code_tokens)
|
||||
# add bos
|
||||
tokens = torch.cat(
|
||||
[
|
||||
torch.tensor([bos_id], device=action.device),
|
||||
prompt_tokens,
|
||||
action_code_tokens,
|
||||
end_tokens,
|
||||
torch.tensor(
|
||||
self._paligemma_tokenizer.encode("Action: ", add_special_tokens=False),
|
||||
device=action.device,
|
||||
),
|
||||
self._act_tokens_to_paligemma_tokens(tokens),
|
||||
torch.tensor(self._paligemma_tokenizer.encode("|"), device=action.device),
|
||||
]
|
||||
)
|
||||
code_mask = torch.zeros(len(tokens), dtype=torch.bool, device=action.device)
|
||||
code_mask[code_start:code_end] = True
|
||||
|
||||
# Truncate or pad to max_action_tokens
|
||||
if len(tokens) > self.max_action_tokens:
|
||||
if not self.allow_truncation:
|
||||
raise ValueError(
|
||||
f"FAST action sequence has {len(tokens)} tokens, exceeding "
|
||||
f"max_action_tokens={self.max_action_tokens}."
|
||||
)
|
||||
logging.warning(
|
||||
f"Token length ({len(tokens)}) exceeds max length ({self.max_action_tokens}), truncating. "
|
||||
"Consider increasing the `max_action_tokens` in your model config if this happens frequently."
|
||||
)
|
||||
tokens = tokens[: self.max_action_tokens]
|
||||
code_mask = code_mask[: self.max_action_tokens]
|
||||
mask = torch.ones(self.max_action_tokens, dtype=torch.bool, device=action.device)
|
||||
else:
|
||||
pad_len = self.max_action_tokens - len(tokens)
|
||||
mask = torch.cat(
|
||||
[
|
||||
torch.ones(len(tokens), dtype=torch.bool, device=action.device),
|
||||
torch.zeros(pad_len, dtype=torch.bool, device=action.device),
|
||||
torch.zeros(
|
||||
self.max_action_tokens - len(tokens), dtype=torch.bool, device=action.device
|
||||
),
|
||||
]
|
||||
)
|
||||
code_mask = torch.nn.functional.pad(code_mask, (0, pad_len), value=False)
|
||||
# Pad tokens with zeros
|
||||
tokens = torch.nn.functional.pad(tokens, (0, pad_len), value=0)
|
||||
tokens = torch.nn.functional.pad(tokens, (0, self.max_action_tokens - len(tokens)), value=0)
|
||||
|
||||
tokens_list.append(tokens)
|
||||
masks_list.append(mask)
|
||||
code_masks_list.append(code_mask)
|
||||
|
||||
# Stack into batched tensors
|
||||
tokens_batch = torch.stack(tokens_list, dim=0) # (B, max_action_tokens)
|
||||
masks_batch = torch.stack(masks_list, dim=0) # (B, max_action_tokens)
|
||||
code_masks_batch = torch.stack(code_masks_list, dim=0) # (B, max_action_tokens)
|
||||
|
||||
# Remove batch dimension if input was single sample
|
||||
if single_sample:
|
||||
tokens_batch = tokens_batch.squeeze(0)
|
||||
masks_batch = masks_batch.squeeze(0)
|
||||
code_masks_batch = code_masks_batch.squeeze(0)
|
||||
|
||||
# Move to the same device as the input
|
||||
if device is not None:
|
||||
tokens_batch = tokens_batch.to(device)
|
||||
masks_batch = masks_batch.to(device)
|
||||
code_masks_batch = code_masks_batch.to(device)
|
||||
|
||||
return tokens_batch, masks_batch, code_masks_batch
|
||||
return tokens_batch, masks_batch
|
||||
|
||||
def action(self, action: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
This method is not used since we override __call__.
|
||||
Required by ActionProcessorStep ABC.
|
||||
"""
|
||||
tokens, _, _ = self._tokenize_action(action)
|
||||
tokens, _ = self._tokenize_action(action)
|
||||
return tokens
|
||||
|
||||
def get_config(self) -> dict[str, Any]:
|
||||
@@ -572,9 +550,6 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
|
||||
config = {
|
||||
"trust_remote_code": self.trust_remote_code,
|
||||
"max_action_tokens": self.max_action_tokens,
|
||||
"fast_skip_tokens": self.fast_skip_tokens,
|
||||
"paligemma_tokenizer_name": self.paligemma_tokenizer_name,
|
||||
"allow_truncation": self.allow_truncation,
|
||||
}
|
||||
|
||||
# Only save tokenizer_name if it was used to create the tokenizer
|
||||
@@ -583,14 +558,6 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
|
||||
|
||||
return config
|
||||
|
||||
def save_artifacts(self, save_directory: Path) -> dict[str, str]:
|
||||
artifact_path = Path("action_tokenizer")
|
||||
save_pretrained = getattr(self.action_tokenizer, "save_pretrained", None)
|
||||
if save_pretrained is None:
|
||||
raise TypeError("Action tokenizer must implement save_pretrained() to save a portable pipeline.")
|
||||
save_pretrained(save_directory / artifact_path)
|
||||
return {"action_tokenizer_name": artifact_path.as_posix()}
|
||||
|
||||
def transform_features(
|
||||
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
||||
|
||||
@@ -21,8 +21,6 @@ from pathlib import Path
|
||||
from tempfile import TemporaryDirectory
|
||||
from typing import TYPE_CHECKING, Any, TypeVar
|
||||
|
||||
import packaging
|
||||
import safetensors
|
||||
from huggingface_hub import HfApi, ModelCard, ModelCardData, hf_hub_download
|
||||
from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
|
||||
from huggingface_hub.errors import HfHubHTTPError
|
||||
@@ -30,6 +28,7 @@ from safetensors.torch import load_model as load_model_as_safetensor, save_model
|
||||
from torch import Tensor, nn
|
||||
|
||||
from lerobot.configs.rewards import RewardModelConfig
|
||||
from lerobot.utils.device_utils import resolve_safetensors_device
|
||||
from lerobot.utils.hub import HubMixin
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -129,29 +128,13 @@ class PreTrainedRewardModel(nn.Module, HubMixin, abc.ABC):
|
||||
|
||||
@classmethod
|
||||
def _load_as_safetensor(cls, model: T, model_file: str, map_location: str, strict: bool) -> T:
|
||||
# Create base kwargs
|
||||
kwargs = {"strict": strict}
|
||||
|
||||
# Add device parameter for newer versions that support it
|
||||
if packaging.version.parse(safetensors.__version__) >= packaging.version.parse("0.4.3"):
|
||||
kwargs["device"] = map_location
|
||||
|
||||
# Load the model with appropriate kwargs
|
||||
missing_keys, unexpected_keys = load_model_as_safetensor(model, model_file, **kwargs)
|
||||
missing_keys, unexpected_keys = load_model_as_safetensor(
|
||||
model, model_file, strict=strict, device=resolve_safetensors_device(map_location)
|
||||
)
|
||||
if missing_keys:
|
||||
logging.warning(f"Missing key(s) when loading model: {missing_keys}")
|
||||
if unexpected_keys:
|
||||
logging.warning(f"Unexpected key(s) when loading model: {unexpected_keys}")
|
||||
|
||||
# For older versions, manually move to device if needed
|
||||
if "device" not in kwargs and map_location != "cpu":
|
||||
logging.warning(
|
||||
"Loading model weights on other devices than 'cpu' is not supported natively in your version of safetensors."
|
||||
" This means that the model is loaded on 'cpu' first and then copied to the device."
|
||||
" This leads to a slower loading time."
|
||||
" Please update safetensors to version 0.4.3 or above for improved performance."
|
||||
)
|
||||
model.to(map_location)
|
||||
return model
|
||||
|
||||
def get_optim_params(self):
|
||||
|
||||
@@ -21,8 +21,6 @@ from lerobot.utils.import_utils import make_device_from_device_class
|
||||
from .config import RobotConfig
|
||||
from .robot import Robot
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def make_robot_from_config(config: RobotConfig) -> Robot:
|
||||
# TODO(Steven): Consider just using the make_device_from_device_class for all types
|
||||
@@ -120,7 +118,7 @@ def ensure_safe_goal_position(
|
||||
}
|
||||
|
||||
if warnings_dict:
|
||||
logger.warning(
|
||||
logging.warning(
|
||||
"Relative goal position magnitude had to be clamped to be safe.\n"
|
||||
f"{pformat(warnings_dict, indent=4)}"
|
||||
)
|
||||
|
||||
@@ -1,38 +0,0 @@
|
||||
# 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.
|
||||
|
||||
"""Policy-agnostic runtime for language-conditioned policies.
|
||||
|
||||
Adapters registered in :mod:`lerobot.runtime.registry` are served by ``lerobot-rollout --language``.
|
||||
"""
|
||||
|
||||
from .adapter import BaseLanguageAdapter, GenerationConfig, LanguageDiagnostics
|
||||
from .language_runtime import (
|
||||
LanguageConditionedPolicyAdapter,
|
||||
LanguageConditionedRuntime,
|
||||
RuntimeState,
|
||||
Tick,
|
||||
TickClock,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"BaseLanguageAdapter",
|
||||
"GenerationConfig",
|
||||
"LanguageConditionedPolicyAdapter",
|
||||
"LanguageConditionedRuntime",
|
||||
"LanguageDiagnostics",
|
||||
"RuntimeState",
|
||||
"Tick",
|
||||
"TickClock",
|
||||
]
|
||||
@@ -1,165 +0,0 @@
|
||||
# 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.
|
||||
|
||||
"""Policy adapters for the language runtime.
|
||||
|
||||
The base adapter owns generation control and diagnostics while subclasses provide policy-specific actions and text.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
from .language_runtime import RuntimeState
|
||||
|
||||
_SAY_RE = re.compile(r"<\s*say\s*>(.*?)<\s*/\s*say\s*>", re.IGNORECASE | re.DOTALL)
|
||||
|
||||
|
||||
@dataclass
|
||||
class GenerationConfig:
|
||||
"""Text-generation settings fixed for the adapter's lifetime."""
|
||||
|
||||
min_new_tokens: int = 0
|
||||
temperature: float = 0.0
|
||||
top_p: float = 1.0
|
||||
chunks_per_regen: int = 1 # regenerate the language context every N action chunks
|
||||
enable_memory: bool = True # generate a running memory note on subtask change
|
||||
enable_subtask: bool = True # generate the low-level subtask (off => use the given text directly)
|
||||
|
||||
|
||||
@dataclass
|
||||
class LanguageDiagnostics:
|
||||
"""Runtime-panel generation counters keyed by text kind."""
|
||||
|
||||
last_raw: dict[str, str] = field(default_factory=dict)
|
||||
empty: dict[str, int] = field(default_factory=dict)
|
||||
repeat: int = 0
|
||||
|
||||
def _bump(self, table: dict[str, int], kind: str) -> int:
|
||||
table[kind] = table.get(kind, 0) + 1
|
||||
return table[kind]
|
||||
|
||||
|
||||
class BaseLanguageAdapter(ABC):
|
||||
"""Batteries-included adapter: generic high-level control, policy primitives abstract."""
|
||||
|
||||
def __init__(self, policy: Any, gen: GenerationConfig | None = None) -> None:
|
||||
self.policy = policy
|
||||
self.gen = gen or GenerationConfig()
|
||||
self.diag = LanguageDiagnostics()
|
||||
self._chunks_until_regen = 0
|
||||
|
||||
@abstractmethod
|
||||
def select_action(self, observation: dict[str, Any], state: RuntimeState) -> Any:
|
||||
"""Produce an action chunk from the observation + current language context."""
|
||||
|
||||
@abstractmethod
|
||||
def generate_text(
|
||||
self,
|
||||
kind: str,
|
||||
observation: dict[str, Any] | None,
|
||||
state: RuntimeState,
|
||||
user_text: str | None = None,
|
||||
) -> str:
|
||||
"""Generate one text stream (``kind``) and return the decoded string."""
|
||||
|
||||
def update_language_state(self, observation: dict[str, Any] | None, state: RuntimeState) -> None:
|
||||
"""Throttled regeneration of the language context (subtask / memory / ...)."""
|
||||
if self._chunks_until_regen > 0:
|
||||
self._chunks_until_regen -= 1
|
||||
return
|
||||
self._chunks_until_regen = max(1, self.gen.chunks_per_regen) - 1
|
||||
self._regenerate_context(observation, state)
|
||||
|
||||
def handle_interjection(
|
||||
self, user_text: str, observation: dict[str, Any] | None, state: RuntimeState
|
||||
) -> None:
|
||||
"""React to a mid-run user message by regenerating the plan."""
|
||||
out = self.generate_text("interjection", observation, state, user_text=user_text)
|
||||
plan = self.plan_from_text(out)
|
||||
if plan:
|
||||
state.set_context("plan", plan, label="plan")
|
||||
|
||||
def plan_from_text(self, text: str) -> str:
|
||||
"""Strip ``<say>`` speech markers from a generated plan."""
|
||||
plan, _speech = split_plan_and_say(text)
|
||||
return plan
|
||||
|
||||
def _regenerate_context(self, observation: dict[str, Any] | None, state: RuntimeState) -> None:
|
||||
"""Default hierarchy: regenerate the subtask, then memory when it changes.
|
||||
|
||||
Override for a policy with a different language hierarchy.
|
||||
"""
|
||||
if not self.gen.enable_subtask:
|
||||
# Preserve operator-provided subtasks in direct mode.
|
||||
return
|
||||
subtask = self._generate_filtered("subtask", observation, state)
|
||||
if subtask is None:
|
||||
return
|
||||
previous = state.language_context.get("subtask")
|
||||
if not state.set_context("subtask", subtask, label="subtask"):
|
||||
self.diag.repeat += 1
|
||||
return
|
||||
self.diag.repeat = 0
|
||||
if previous:
|
||||
state.extra["prior_subtask"] = previous
|
||||
if not self.gen.enable_memory:
|
||||
return
|
||||
memory = self._generate_filtered("memory", observation, state)
|
||||
if memory is not None:
|
||||
state.set_context("memory", memory, label="memory")
|
||||
|
||||
def _generate_filtered(
|
||||
self, kind: str, observation: dict[str, Any] | None, state: RuntimeState
|
||||
) -> str | None:
|
||||
"""Generate one ``kind``, record diagnostics, and drop empty output."""
|
||||
text = self.generate_text(kind, observation, state)
|
||||
self.diag.last_raw[kind] = text or ""
|
||||
if not text:
|
||||
count = self.diag._bump(self.diag.empty, kind)
|
||||
if count == 1 or count % 5 == 0:
|
||||
state.log(f" [info] {kind} gen returned empty (x{count})")
|
||||
return None
|
||||
return text
|
||||
|
||||
|
||||
class DirectTaskPolicyAdapter(BaseLanguageAdapter):
|
||||
"""Adapter for flat policies whose preprocessors condition actions on the operator's task."""
|
||||
|
||||
def select_action(self, observation: dict[str, Any], state: RuntimeState) -> Any:
|
||||
return self.policy.predict_action_chunk(observation)
|
||||
|
||||
def generate_text(
|
||||
self,
|
||||
kind: str,
|
||||
observation: dict[str, Any] | None,
|
||||
state: RuntimeState,
|
||||
user_text: str | None = None,
|
||||
) -> str:
|
||||
return ""
|
||||
|
||||
|
||||
def split_plan_and_say(text: str) -> tuple[str, str]:
|
||||
"""Split ``plan <say>speech</say>`` into ``(plan, speech)``."""
|
||||
if not text:
|
||||
return "", ""
|
||||
match = _SAY_RE.search(text)
|
||||
if not match:
|
||||
return text.strip(), ""
|
||||
speech = match.group(1).strip().strip('"').strip("'")
|
||||
plan = (text[: match.start()] + text[match.end() :]).strip()
|
||||
return plan, speech
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,349 +0,0 @@
|
||||
# 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.
|
||||
|
||||
"""Small reusable runtime for language-conditioned robot policies."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
from collections import deque
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Protocol
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RuntimeState:
|
||||
"""Explicit state shared by the runtime and policy adapter."""
|
||||
|
||||
task: str = ""
|
||||
language_context: dict[str, str] = field(default_factory=dict)
|
||||
action_queue: deque[Any] = field(default_factory=deque)
|
||||
events: set[str] = field(default_factory=set)
|
||||
log_lines: list[str] = field(default_factory=list)
|
||||
mode: str = "action"
|
||||
stop: bool = False
|
||||
tick: Tick | None = None
|
||||
actions_dispatched: int = 0
|
||||
action_deadline: float | None = None
|
||||
extra: dict[str, Any] = field(default_factory=dict)
|
||||
revision: int = 0
|
||||
lock: Any = field(default_factory=threading.RLock, repr=False)
|
||||
|
||||
def emit(self, event_name: str) -> None:
|
||||
self.events.add(event_name)
|
||||
|
||||
def take_event(self, event_name: str) -> bool:
|
||||
if event_name not in self.events:
|
||||
return False
|
||||
self.events.remove(event_name)
|
||||
return True
|
||||
|
||||
def log(self, line: str) -> None:
|
||||
self.log_lines.append(line)
|
||||
|
||||
def set_context(self, key: str, value: str | None, *, label: str | None = None) -> bool:
|
||||
with self.lock:
|
||||
previous = self.language_context.get(key)
|
||||
if previous == value:
|
||||
return False
|
||||
if value is None:
|
||||
self.language_context.pop(key, None)
|
||||
else:
|
||||
self.language_context[key] = value
|
||||
self.revision += 1
|
||||
if label is not None and value:
|
||||
self.log(f" {label}: {value}")
|
||||
return True
|
||||
|
||||
def get(self, key: str, default: Any = None) -> Any:
|
||||
try:
|
||||
return self[key]
|
||||
except KeyError:
|
||||
return default
|
||||
|
||||
def setdefault(self, key: str, default: Any = None) -> Any:
|
||||
current = self.get(key, None)
|
||||
if current is not None:
|
||||
return current
|
||||
self[key] = default
|
||||
return default
|
||||
|
||||
def __getitem__(self, key: str) -> Any:
|
||||
if hasattr(self, key):
|
||||
return getattr(self, key)
|
||||
if key in self.extra:
|
||||
return self.extra[key]
|
||||
raise KeyError(key)
|
||||
|
||||
def __setitem__(self, key: str, value: Any) -> None:
|
||||
with self.lock:
|
||||
if hasattr(self, key):
|
||||
if key == "mode" and self.mode != value:
|
||||
self.revision += 1
|
||||
setattr(self, key, value)
|
||||
else:
|
||||
self.extra[key] = value
|
||||
|
||||
|
||||
class LanguageConditionedPolicyAdapter(Protocol):
|
||||
"""Runtime policy contract, implemented directly or through ``BaseLanguageAdapter``."""
|
||||
|
||||
def select_action(self, observation: dict[str, Any], state: RuntimeState) -> Any: ...
|
||||
|
||||
def update_language_state(self, observation: dict[str, Any] | None, state: RuntimeState) -> None: ...
|
||||
|
||||
def handle_interjection(
|
||||
self, user_text: str, observation: dict[str, Any] | None, state: RuntimeState
|
||||
) -> None: ...
|
||||
|
||||
|
||||
@dataclass
|
||||
class Tick:
|
||||
index: int
|
||||
monotonic_seconds: float
|
||||
|
||||
|
||||
@dataclass
|
||||
class TickClock:
|
||||
max_rate_hz: float = 50.0
|
||||
_index: int = field(default=0, init=False)
|
||||
_last_seconds: float | None = field(default=None, init=False)
|
||||
|
||||
def advance(self) -> Tick:
|
||||
period = 1.0 / max(self.max_rate_hz, 0.1)
|
||||
now = time.monotonic()
|
||||
if self._last_seconds is not None:
|
||||
sleep_for = (self._last_seconds + period) - now
|
||||
if sleep_for > 0:
|
||||
time.sleep(sleep_for)
|
||||
now = time.monotonic()
|
||||
self._last_seconds = now
|
||||
self._index += 1
|
||||
return Tick(index=self._index, monotonic_seconds=now)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _RateGate:
|
||||
hz: float
|
||||
_last_seconds: float | None = None
|
||||
|
||||
def due(self, tick: Tick, *, force: bool = False) -> bool:
|
||||
if force:
|
||||
self._last_seconds = tick.monotonic_seconds
|
||||
return True
|
||||
period = 1.0 / max(self.hz, 1e-6)
|
||||
if self._last_seconds is None or tick.monotonic_seconds - self._last_seconds >= period:
|
||||
self._last_seconds = tick.monotonic_seconds
|
||||
return True
|
||||
return False
|
||||
|
||||
def rearm(self) -> None:
|
||||
self._last_seconds = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class LanguageConditionedRuntime:
|
||||
"""Generic tick loop for language-conditioned robot policies."""
|
||||
|
||||
policy_adapter: LanguageConditionedPolicyAdapter
|
||||
observation_provider: Callable[[], dict[str, Any] | None] | None = None
|
||||
action_executor: Callable[[Any], None] | None = None
|
||||
event_collector: Callable[[RuntimeState], None] | None = None
|
||||
chunk_hz: float = 4.0
|
||||
ctrl_hz: float = 50.0
|
||||
high_level_hz: float = 1.0
|
||||
max_rate_hz: float = 50.0
|
||||
|
||||
state: RuntimeState = field(default_factory=RuntimeState)
|
||||
_chunk_gate: _RateGate = field(init=False)
|
||||
_ctrl_gate: _RateGate = field(init=False)
|
||||
_language_gate: _RateGate = field(init=False)
|
||||
_stop: bool = field(default=False, init=False)
|
||||
_last_dispatch_seconds: float | None = field(default=None, init=False)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self._chunk_gate = _RateGate(self.chunk_hz)
|
||||
self._ctrl_gate = _RateGate(self.ctrl_hz)
|
||||
self._language_gate = _RateGate(self.high_level_hz)
|
||||
|
||||
@property
|
||||
def policy(self) -> Any:
|
||||
return getattr(self.policy_adapter, "policy", self.policy_adapter)
|
||||
|
||||
def set_task(self, task: str) -> None:
|
||||
with self.state.lock:
|
||||
if self.state.task != task:
|
||||
self.state.revision += 1
|
||||
self.state.task = task
|
||||
self.state.log(f"Task: {task}")
|
||||
|
||||
def stop(self) -> None:
|
||||
self._stop = True
|
||||
self.state.stop = True
|
||||
|
||||
def run(self, *, max_ticks: int | None = None) -> None:
|
||||
clock = TickClock(max_rate_hz=self.max_rate_hz)
|
||||
while not self._stop:
|
||||
tick = clock.advance()
|
||||
self._run_tick(tick)
|
||||
self._flush_logs()
|
||||
if self.state.stop:
|
||||
self._stop = True
|
||||
if max_ticks is not None and tick.index >= max_ticks:
|
||||
break
|
||||
self._on_shutdown()
|
||||
|
||||
def step_once(self) -> list[str]:
|
||||
previous = self.state.tick.index if self.state.tick is not None else 0
|
||||
tick = Tick(index=previous + 1, monotonic_seconds=time.monotonic())
|
||||
self._run_tick(tick, force_rates=True)
|
||||
return list(self.state.log_lines)
|
||||
|
||||
def _run_tick(self, tick: Tick, *, force_rates: bool = False) -> None:
|
||||
self.state.tick = tick
|
||||
self.state.log_lines = []
|
||||
if self.event_collector is not None:
|
||||
self.event_collector(self.state)
|
||||
self._handle_action_deadline()
|
||||
if self.state.stop:
|
||||
return
|
||||
self.maybe_update_language_state(force=force_rates)
|
||||
self.maybe_handle_user_events()
|
||||
self.maybe_enqueue_action_chunk(force=force_rates)
|
||||
self.dispatch_action(force=force_rates)
|
||||
self.state.events.clear()
|
||||
|
||||
def _current_observation(self) -> dict[str, Any] | None:
|
||||
if self.observation_provider is None:
|
||||
return None
|
||||
try:
|
||||
return self.observation_provider()
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.debug("observation_provider failed: %s", exc)
|
||||
return None
|
||||
|
||||
def maybe_update_language_state(self, *, force: bool = False) -> None:
|
||||
if self.state.mode != "action" or not self.state.task:
|
||||
return
|
||||
if self.state.action_queue:
|
||||
self._language_gate.rearm()
|
||||
return
|
||||
if self.state.tick is None or not self._language_gate.due(self.state.tick, force=force):
|
||||
return
|
||||
observation = self._current_observation()
|
||||
try:
|
||||
self.policy_adapter.update_language_state(observation, self.state)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning("language update failed: %s", exc, exc_info=logger.isEnabledFor(logging.DEBUG))
|
||||
self.state.log(f" [warn] language update failed: {type(exc).__name__}: {exc}")
|
||||
|
||||
def maybe_handle_user_events(self) -> None:
|
||||
if self.state.take_event("user_interjection"):
|
||||
self._handle_user_interjection()
|
||||
|
||||
def _handle_user_interjection(self) -> None:
|
||||
text = str(self.state.extra.get("recent_interjection") or "")
|
||||
if not text:
|
||||
return
|
||||
observation = self._current_observation()
|
||||
self.policy_adapter.handle_interjection(text, observation, self.state)
|
||||
self.state.extra["recent_interjection"] = None
|
||||
|
||||
def maybe_enqueue_action_chunk(self, *, force: bool = False) -> None:
|
||||
with self.state.lock:
|
||||
if self.state.mode != "action" or not self.state.task:
|
||||
return
|
||||
if self.state.action_queue:
|
||||
return
|
||||
if self.state.tick is None or not self._chunk_gate.due(self.state.tick, force=force):
|
||||
return
|
||||
revision = self.state.revision
|
||||
observation = self._current_observation()
|
||||
if observation is None:
|
||||
return
|
||||
try:
|
||||
chunk = self.policy_adapter.select_action(observation, self.state)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning("select_action failed: %s", exc, exc_info=logger.isEnabledFor(logging.DEBUG))
|
||||
self.state.log(f" [warn] select_action failed: {type(exc).__name__}: {exc}")
|
||||
return
|
||||
with self.state.lock:
|
||||
if (
|
||||
self.state.revision != revision
|
||||
or self.state.mode != "action"
|
||||
or self.state.stop
|
||||
or self._stop
|
||||
):
|
||||
logger.info("Discarded an action chunk invalidated during inference.")
|
||||
return
|
||||
self._enqueue_chunk(chunk)
|
||||
|
||||
def _enqueue_chunk(self, chunk: Any) -> None:
|
||||
if chunk is None:
|
||||
return
|
||||
chunk_iter = chunk[0] if getattr(chunk, "ndim", None) == 3 else chunk
|
||||
if getattr(chunk_iter, "ndim", None) == 1:
|
||||
chunk_iter = chunk_iter.unsqueeze(0)
|
||||
for step in chunk_iter:
|
||||
self.state.action_queue.append(step.unsqueeze(0) if hasattr(step, "unsqueeze") else step)
|
||||
try:
|
||||
self.state.extra["last_chunk_size"] = int(chunk_iter.shape[0])
|
||||
except Exception: # noqa: BLE001
|
||||
self.state.extra["last_chunk_size"] = len(self.state.action_queue)
|
||||
|
||||
def dispatch_action(self, *, force: bool = False) -> None:
|
||||
if self.state.mode != "action":
|
||||
self._last_dispatch_seconds = None
|
||||
return
|
||||
if self.state.tick is None or not self._ctrl_gate.due(self.state.tick, force=force):
|
||||
return
|
||||
queue = self.state.action_queue
|
||||
if not queue:
|
||||
self._last_dispatch_seconds = None
|
||||
return
|
||||
now = time.monotonic()
|
||||
if self._last_dispatch_seconds is None or self.ctrl_hz <= 0:
|
||||
n_to_pop = 1
|
||||
else:
|
||||
n_to_pop = max(1, min(len(queue), int(round((now - self._last_dispatch_seconds) * self.ctrl_hz))))
|
||||
self._last_dispatch_seconds = now
|
||||
latest = None
|
||||
for _ in range(n_to_pop):
|
||||
if not queue:
|
||||
break
|
||||
latest = queue.popleft()
|
||||
self.state.actions_dispatched += 1
|
||||
if latest is not None and self.action_executor is not None:
|
||||
self.action_executor(latest)
|
||||
|
||||
def _handle_action_deadline(self) -> None:
|
||||
deadline = self.state.action_deadline
|
||||
if self.state.mode == "action" and deadline is not None and time.monotonic() >= deadline:
|
||||
self.state.mode = "paused"
|
||||
self.state.action_deadline = None
|
||||
self.state.action_queue.clear()
|
||||
self.state.log("timed action elapsed — paused")
|
||||
|
||||
def _flush_logs(self) -> None:
|
||||
for line in self.state.log_lines:
|
||||
print(f"[runtime] {line}", flush=True)
|
||||
|
||||
def _on_shutdown(self) -> None:
|
||||
self.state.action_queue.clear()
|
||||
print("[runtime] stopped", flush=True)
|
||||
@@ -1,39 +0,0 @@
|
||||
# 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.
|
||||
|
||||
"""Lazy mapping from policy types to language-runtime adapters."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
_ADAPTERS: dict[str, str] = {
|
||||
"pi052": "lerobot.policies.pi052.inference.pi052_adapter:PI052PolicyAdapter",
|
||||
"pi05": "lerobot.runtime.adapter:DirectTaskPolicyAdapter",
|
||||
"molmoact2": "lerobot.runtime.adapter:DirectTaskPolicyAdapter",
|
||||
}
|
||||
|
||||
|
||||
def get_language_adapter_factory(policy_type: str) -> Callable[..., Any]:
|
||||
"""Return the adapter class registered for ``policy_type``."""
|
||||
spec = _ADAPTERS.get(policy_type)
|
||||
if spec is None:
|
||||
raise ValueError(
|
||||
f"No language-runtime adapter registered for policy type {policy_type!r}. "
|
||||
f"Registered: {sorted(_ADAPTERS)}. Add an entry to lerobot.runtime.registry."
|
||||
)
|
||||
module_path, class_name = spec.split(":")
|
||||
return getattr(importlib.import_module(module_path), class_name)
|
||||
@@ -1,406 +0,0 @@
|
||||
# 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.
|
||||
|
||||
"""RoboCasa backend for interactive language-conditioned rollouts.
|
||||
|
||||
It reuses the eval observation/action pipeline while prompts control a persistent selected scene.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from collections.abc import Callable
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from lerobot.utils.io_utils import StreamingVideoWriter
|
||||
from lerobot.utils.video_annotation import annotate_frame
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _short_cam_name(cam: str) -> str:
|
||||
"""Human-friendly view label for a RoboCasa camera name."""
|
||||
c = cam.replace("robot0_", "")
|
||||
return {
|
||||
"agentview_left": "left",
|
||||
"agentview_right": "right",
|
||||
"eye_in_hand": "wrist",
|
||||
}.get(c, c)
|
||||
|
||||
|
||||
def _label_panel(img: np.ndarray, label: str) -> np.ndarray:
|
||||
"""Draw a small camera-view label in the bottom-left corner of a panel."""
|
||||
try:
|
||||
import cv2 # noqa: PLC0415
|
||||
except ImportError:
|
||||
return img
|
||||
y = img.shape[0] - 6
|
||||
cv2.putText(img, label, (5, y), cv2.FONT_HERSHEY_SIMPLEX, 0.45, (0, 0, 0), 3, cv2.LINE_AA)
|
||||
cv2.putText(img, label, (5, y), cv2.FONT_HERSHEY_SIMPLEX, 0.45, (255, 255, 0), 1, cv2.LINE_AA)
|
||||
return img
|
||||
|
||||
|
||||
# Two workers avoid broken single-worker EGL rendering; only env 0 is displayed.
|
||||
_SIM_N_ENVS = 2
|
||||
|
||||
|
||||
def create_sim_env(
|
||||
*,
|
||||
task: str,
|
||||
split: str | None,
|
||||
obj_registries: list[str],
|
||||
seed: int | None,
|
||||
render_size: int = 384,
|
||||
) -> tuple[Any, dict]:
|
||||
"""Create and reset the vectorized RoboCasa environment before CUDA initializes.
|
||||
|
||||
Two workers keep EGL stable, while only env 0 is driven and displayed.
|
||||
"""
|
||||
from lerobot.envs.configs import RoboCasaEnv as RoboCasaEnvConfig # noqa: PLC0415
|
||||
|
||||
# The policy resizes inputs, so render_size only affects display quality and cost.
|
||||
env_cfg = RoboCasaEnvConfig(
|
||||
task=task,
|
||||
split=split,
|
||||
obj_registries=list(obj_registries),
|
||||
observation_height=render_size,
|
||||
observation_width=render_size,
|
||||
)
|
||||
# Keep one kitchen alive across sequential prompts.
|
||||
envs = env_cfg.create_envs(
|
||||
n_envs=_SIM_N_ENVS,
|
||||
use_async_envs=True,
|
||||
terminate_on_success=False,
|
||||
horizon=100_000,
|
||||
)
|
||||
env = envs[next(iter(envs))][0]
|
||||
logger.info("[sim] resetting RoboCasa scene task=%r split=%r (n_envs=%d)", task, split, _SIM_N_ENVS)
|
||||
seeds = None if seed is None else [seed + i for i in range(_SIM_N_ENVS)]
|
||||
obs, _ = env.reset(seed=seeds)
|
||||
return env, obs
|
||||
|
||||
|
||||
def start_mjpeg_server(port: int, get_frame: Callable[[], np.ndarray | None]) -> Any:
|
||||
"""Start an MJPEG server that shows a placeholder until ``get_frame`` returns frames."""
|
||||
import io # noqa: PLC0415
|
||||
import threading # noqa: PLC0415
|
||||
import time # noqa: PLC0415
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer # noqa: PLC0415
|
||||
|
||||
from PIL import Image # noqa: PLC0415
|
||||
|
||||
_placeholder = Image.new("RGB", (256, 256), (17, 17, 17))
|
||||
|
||||
class _Handler(BaseHTTPRequestHandler):
|
||||
def log_message(self, *args): # silence per-request logging
|
||||
pass
|
||||
|
||||
def do_GET(self): # noqa: N802
|
||||
if self.path in ("/", "/index.html"):
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "text/html")
|
||||
self.end_headers()
|
||||
self.wfile.write(
|
||||
b"<html><body style='margin:0;background:#111;text-align:center'>"
|
||||
b"<img src='/stream' style='max-width:100vw;max-height:100vh;"
|
||||
b"image-rendering:pixelated'></body></html>"
|
||||
)
|
||||
return
|
||||
if self.path != "/stream":
|
||||
self.send_response(404)
|
||||
self.end_headers()
|
||||
return
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "multipart/x-mixed-replace; boundary=frame")
|
||||
self.end_headers()
|
||||
try:
|
||||
while True:
|
||||
frame = get_frame()
|
||||
buf = io.BytesIO()
|
||||
img = Image.fromarray(frame) if frame is not None else _placeholder
|
||||
img.save(buf, format="JPEG", quality=80)
|
||||
data = buf.getvalue()
|
||||
self.wfile.write(
|
||||
b"--frame\r\nContent-Type: image/jpeg\r\nContent-Length: "
|
||||
+ str(len(data)).encode()
|
||||
+ b"\r\n\r\n"
|
||||
+ data
|
||||
+ b"\r\n"
|
||||
)
|
||||
time.sleep(0.05)
|
||||
except (BrokenPipeError, ConnectionResetError):
|
||||
pass
|
||||
|
||||
try:
|
||||
# Bind all interfaces intentionally so the viewer remains reachable
|
||||
# through the documented SSH port-forwarding workflow.
|
||||
server = ThreadingHTTPServer(("0.0.0.0", port), _Handler) # nosec B104
|
||||
except OSError as exc:
|
||||
logger.warning("[sim] could not start live stream on port %d: %s", port, exc)
|
||||
print(f"[runtime] WARNING: live stream port {port} unavailable ({exc})", flush=True)
|
||||
return None
|
||||
threading.Thread(target=server.serve_forever, daemon=True, name="sim-mjpeg").start()
|
||||
print(
|
||||
f"[runtime] live view: http://localhost:{port} "
|
||||
f"(over SSH: ssh -L {port}:localhost:{port} <host>) — loading until scene is ready",
|
||||
flush=True,
|
||||
)
|
||||
return server
|
||||
|
||||
|
||||
class RoboCasaSimBackend:
|
||||
"""Expose a RoboCasa environment through the runtime observation/action contract.
|
||||
|
||||
The environment must be created before the policy initializes CUDA.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
env: Any,
|
||||
last_obs: dict,
|
||||
task: str,
|
||||
seed: int | None,
|
||||
device: str,
|
||||
preprocessor: Any,
|
||||
postprocessor: Any,
|
||||
record: bool = True,
|
||||
output_dir: str | None = None,
|
||||
view_cams: list[str] | None = None,
|
||||
) -> None:
|
||||
self.env = env
|
||||
self._last_obs = last_obs
|
||||
self._scene_task = task
|
||||
self._view_cams = view_cams or [
|
||||
"robot0_agentview_left",
|
||||
"robot0_eye_in_hand",
|
||||
"robot0_agentview_right",
|
||||
]
|
||||
self.device = torch.device(device) if isinstance(device, str) else device
|
||||
self.preprocessor = preprocessor
|
||||
self.postprocessor = postprocessor
|
||||
self.seed = seed
|
||||
self.record = record
|
||||
self.output_dir = Path(output_dir) if output_dir else Path("outputs/runtime_sim")
|
||||
|
||||
self._video_writer: StreamingVideoWriter | None = None
|
||||
self._video_path: Path | None = None
|
||||
self._live_counter = 0
|
||||
self._latest_frame: np.ndarray | None = None
|
||||
self._stream_server: Any = None
|
||||
self._reset_count = 0
|
||||
# Bind these after runtime construction for live annotations.
|
||||
self._task_getter: Callable[[], str | None] | None = None
|
||||
self._subtask_getter: Callable[[], str | None] | None = None
|
||||
self._memory_getter: Callable[[], str | None] | None = None
|
||||
logger.info("[sim] scene ready — task_description=%r", self._scene_description())
|
||||
|
||||
def bind_runtime(self, runtime: Any) -> None:
|
||||
"""Wire live task/subtask/memory getters from the runtime state."""
|
||||
self._task_getter = lambda: runtime.state.get("task")
|
||||
self._subtask_getter = lambda: runtime.state.language_context.get("subtask")
|
||||
self._memory_getter = lambda: (runtime.state.get("language_context") or {}).get("memory")
|
||||
|
||||
def _scene_description(self) -> str:
|
||||
try:
|
||||
return str(self.env.get_attr("task_description")[0]) or self._scene_task
|
||||
except Exception: # noqa: BLE001
|
||||
return self._scene_task
|
||||
|
||||
def _current_task(self) -> str:
|
||||
task = self._task_getter() if self._task_getter else None
|
||||
return task or self._scene_description() or self._scene_task
|
||||
|
||||
def reset_scene(self) -> None:
|
||||
"""Re-roll the kitchen: reset the env to a fresh scene (new layout/style).
|
||||
|
||||
Uses a new seed each call so ``/reset`` explores different kitchens.
|
||||
"""
|
||||
self._reset_count += 1
|
||||
n = self.env.num_envs
|
||||
if self.seed is None:
|
||||
seeds = None
|
||||
else:
|
||||
base = self.seed + self._reset_count * 1000
|
||||
seeds = [base + i for i in range(n)]
|
||||
obs, _ = self.env.reset(seed=seeds)
|
||||
self._last_obs = obs
|
||||
logger.info("[sim] scene reset (#%d)", self._reset_count)
|
||||
|
||||
def _env0_obs(self) -> dict:
|
||||
"""Slice env 0 out of the batched vec-env observation (batch of 1)."""
|
||||
raw = self._last_obs or {}
|
||||
pixels = raw.get("pixels")
|
||||
out: dict[str, Any] = {}
|
||||
if isinstance(pixels, dict):
|
||||
out["pixels"] = {k: np.asarray(v)[0:1] for k, v in pixels.items()}
|
||||
agent_pos = raw.get("agent_pos")
|
||||
if agent_pos is not None:
|
||||
out["agent_pos"] = np.asarray(agent_pos)[0:1]
|
||||
return out
|
||||
|
||||
def observation_provider(self) -> dict | None:
|
||||
from lerobot.envs.utils import preprocess_observation # noqa: PLC0415
|
||||
|
||||
try:
|
||||
obs = preprocess_observation(self._env0_obs())
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning("[sim] preprocess_observation failed: %s", exc)
|
||||
return None
|
||||
# The adapter later replaces this recipe input with its generated subtask.
|
||||
obs["task"] = [self._current_task()]
|
||||
if self.preprocessor is not None:
|
||||
try:
|
||||
obs = self.preprocessor(obs)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning("[sim] preprocessor failed: %s", exc)
|
||||
return None
|
||||
return {
|
||||
k: (v.to(self.device) if isinstance(v, torch.Tensor) else v)
|
||||
for k, v in obs.items()
|
||||
if isinstance(k, str) and k.startswith("observation.")
|
||||
}
|
||||
|
||||
def action_executor(self, action: Any) -> None:
|
||||
try:
|
||||
if self.postprocessor is not None:
|
||||
action = self.postprocessor(action)
|
||||
if isinstance(action, torch.Tensor):
|
||||
if action.ndim > 1 and action.shape[0] == 1:
|
||||
action = action.squeeze(0)
|
||||
action = action.detach().to("cpu").numpy()
|
||||
# Tile env 0's action because the extra workers exist only for EGL stability.
|
||||
action_row = np.asarray(action, dtype=np.float32).reshape(-1)
|
||||
action_np = np.tile(action_row, (self.env.num_envs, 1))
|
||||
obs, _reward, terminated, truncated, _info = self.env.step(action_np)
|
||||
self._last_obs = obs
|
||||
self._capture_frame()
|
||||
# AsyncVectorEnv resets terminated sub-environments automatically.
|
||||
if bool(np.any(terminated)) or bool(np.any(truncated)):
|
||||
logger.info("[sim] episode ended — scene auto-reset")
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.error("[sim] env.step failed: %s", exc, exc_info=True)
|
||||
|
||||
def _multiview_frame(self) -> np.ndarray | None:
|
||||
"""Label and compose env 0's existing observation views without extra rendering."""
|
||||
pixels = (self._last_obs or {}).get("pixels")
|
||||
if not isinstance(pixels, dict) or not pixels:
|
||||
return None
|
||||
panels: list[np.ndarray] = []
|
||||
for cam in self._view_cams:
|
||||
v = pixels.get(cam)
|
||||
if v is None:
|
||||
continue
|
||||
img = np.asarray(v)
|
||||
if img.ndim == 4: # (n_envs, H, W, C) -> env 0
|
||||
img = img[0]
|
||||
if img.ndim != 3 or img.shape[-1] != 3:
|
||||
continue
|
||||
panels.append(_label_panel(np.ascontiguousarray(img.astype(np.uint8)), _short_cam_name(cam)))
|
||||
if not panels:
|
||||
return None
|
||||
h = min(p.shape[0] for p in panels)
|
||||
panels = [p[:h] for p in panels]
|
||||
return np.concatenate(panels, axis=1)
|
||||
|
||||
def _capture_frame(self) -> None:
|
||||
frame = self._multiview_frame()
|
||||
if frame is None: # fallback to single env.render()
|
||||
try:
|
||||
rendered = self.env.call("render")[0]
|
||||
if isinstance(rendered, np.ndarray) and rendered.ndim == 3:
|
||||
frame = rendered
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.debug("[sim] render failed: %s", exc)
|
||||
if frame is None:
|
||||
return
|
||||
subtask = self._subtask_getter() if self._subtask_getter else None
|
||||
memory = self._memory_getter() if self._memory_getter else None
|
||||
annotated = annotate_frame(
|
||||
frame,
|
||||
(("Task", self._current_task()), ("Subtask", subtask), ("Memory", memory)),
|
||||
)
|
||||
self._latest_frame = annotated # served by the live MJPEG stream
|
||||
self._write_live_frame(annotated)
|
||||
if self.record:
|
||||
self._write_recording_frame(annotated)
|
||||
|
||||
def _write_recording_frame(self, frame: np.ndarray) -> None:
|
||||
try:
|
||||
if self._video_writer is None:
|
||||
self.output_dir.mkdir(parents=True, exist_ok=True)
|
||||
stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
self._video_path = self.output_dir / f"sim_{stamp}.mp4"
|
||||
fps = int((getattr(self.env, "metadata", None) or {}).get("render_fps", 20))
|
||||
self._video_writer = StreamingVideoWriter(self._video_path, fps)
|
||||
self._video_writer.add_frame(frame)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning("[sim] video encoding failed: %s", exc)
|
||||
self.record = False
|
||||
|
||||
def _write_live_frame(self, frame: np.ndarray) -> None:
|
||||
"""Write a rolling latest.png every few frames for live viewing over SSH.
|
||||
|
||||
Open ``{output_dir}/latest.png`` in an editor/viewer and refresh to watch
|
||||
the rollout in near-real-time without a GUI window. Written atomically
|
||||
(temp + replace) so a reader never sees a half-written file.
|
||||
"""
|
||||
self._live_counter += 1
|
||||
if self._live_counter % 3 != 0:
|
||||
return
|
||||
try:
|
||||
import os # noqa: PLC0415
|
||||
|
||||
from PIL import Image # noqa: PLC0415
|
||||
|
||||
self.output_dir.mkdir(parents=True, exist_ok=True)
|
||||
tmp = self.output_dir / ".latest.tmp.png"
|
||||
Image.fromarray(frame).save(tmp)
|
||||
os.replace(tmp, self.output_dir / "latest.png")
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.debug("[sim] live frame write failed: %s", exc)
|
||||
|
||||
def _flush_video(self) -> None:
|
||||
if self._video_writer is None:
|
||||
return
|
||||
writer = self._video_writer
|
||||
self._video_writer = None
|
||||
try:
|
||||
writer.close()
|
||||
logger.info("[sim] wrote video (%d frames) to %s", writer.frames_written, self._video_path)
|
||||
print(f"[runtime] sim video saved to {self._video_path}", flush=True)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning("[sim] video close failed: %s", exc)
|
||||
|
||||
def attach_stream_server(self, server: Any) -> None:
|
||||
"""Attach an already-running MJPEG server so disconnect() can stop it."""
|
||||
self._stream_server = server
|
||||
|
||||
def disconnect(self) -> None:
|
||||
"""Match the robot backend's cleanup contract."""
|
||||
if self._stream_server is not None:
|
||||
try:
|
||||
self._stream_server.shutdown()
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.debug("[sim] stream server shutdown raised %s", exc)
|
||||
self._flush_video()
|
||||
try:
|
||||
self.env.close()
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.debug("[sim] env.close raised %s", exc)
|
||||
@@ -28,7 +28,12 @@ For distributed runs, see ``examples/annotations/run_hf_job.py``.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from contextlib import suppress
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from huggingface_hub import HfApi, snapshot_download
|
||||
from huggingface_hub.errors import RevisionNotFoundError
|
||||
|
||||
from lerobot.annotations.steerable_pipeline.config import AnnotationPipelineConfig
|
||||
from lerobot.annotations.steerable_pipeline.executor import Executor
|
||||
@@ -42,6 +47,12 @@ from lerobot.annotations.steerable_pipeline.validator import StagingValidator
|
||||
from lerobot.annotations.steerable_pipeline.vlm_client import make_vlm_client
|
||||
from lerobot.annotations.steerable_pipeline.writer import LanguageColumnsWriter
|
||||
from lerobot.configs import parser
|
||||
from lerobot.utils.import_utils import _datasets_available, require_package
|
||||
|
||||
if TYPE_CHECKING or _datasets_available:
|
||||
from lerobot.datasets.dataset_metadata import CODEBASE_VERSION
|
||||
from lerobot.datasets.io_utils import load_info
|
||||
from lerobot.datasets.utils import create_lerobot_dataset_card
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -50,8 +61,6 @@ def _resolve_root(cfg: AnnotationPipelineConfig) -> Path:
|
||||
if cfg.root is not None:
|
||||
return Path(cfg.root)
|
||||
if cfg.repo_id is not None:
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
return Path(snapshot_download(repo_id=cfg.repo_id, repo_type="dataset"))
|
||||
raise ValueError("Either --root or --repo_id must be provided.")
|
||||
|
||||
@@ -125,10 +134,7 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
|
||||
|
||||
Pushes to ``cfg.new_repo_id`` when set, otherwise back to ``cfg.repo_id``.
|
||||
"""
|
||||
from huggingface_hub import HfApi # noqa: PLC0415
|
||||
|
||||
from lerobot.datasets.io_utils import load_info # noqa: PLC0415
|
||||
from lerobot.datasets.utils import create_lerobot_dataset_card # noqa: PLC0415
|
||||
require_package("datasets", "dataset")
|
||||
|
||||
repo_id = cfg.new_repo_id or cfg.repo_id
|
||||
commit_message = cfg.push_commit_message or "Add steerable annotations (lerobot-annotate)"
|
||||
@@ -163,8 +169,6 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
|
||||
# ``RevisionNotFoundError``. Read the version straight from the
|
||||
# dataset's own ``meta/info.json`` so we tag whatever the writer
|
||||
# actually wrote (no accidental drift if the codebase floor moves).
|
||||
from lerobot.datasets.dataset_metadata import CODEBASE_VERSION # noqa: PLC0415
|
||||
|
||||
version_tag = (
|
||||
dataset_info.codebase_version if dataset_info.codebase_version.startswith("v") else CODEBASE_VERSION
|
||||
)
|
||||
@@ -178,10 +182,6 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
|
||||
tag_kwargs["revision"] = revision
|
||||
|
||||
try:
|
||||
from contextlib import suppress # noqa: PLC0415
|
||||
|
||||
from huggingface_hub.errors import RevisionNotFoundError # noqa: PLC0415
|
||||
|
||||
with suppress(RevisionNotFoundError):
|
||||
api.delete_tag(repo_id, tag=version_tag, repo_type="dataset")
|
||||
api.create_tag(**tag_kwargs)
|
||||
|
||||
@@ -94,19 +94,6 @@ from lerobot.utils.utils import (
|
||||
init_logging,
|
||||
inside_slurm,
|
||||
)
|
||||
from lerobot.utils.video_annotation import annotate_frame
|
||||
|
||||
|
||||
def _annotate_eval_frames(frames: np.ndarray, task: str | None, subtask: str | None) -> np.ndarray:
|
||||
"""Overlay the high-level task and predicted subtask onto rendered frames.
|
||||
|
||||
``frames`` is ``(n_envs, H, W, C)`` uint8. Best-effort: if OpenCV isn't
|
||||
available the frames are returned unchanged so eval never fails over a
|
||||
visualization concern.
|
||||
"""
|
||||
if frames.ndim != 4 or frames.shape[-1] != 3:
|
||||
return frames
|
||||
return np.stack([annotate_frame(frame, (("Task", task), ("Subtask", subtask))) for frame in frames])
|
||||
|
||||
|
||||
def _env_features_to_dataset_features(env_features: dict) -> dict:
|
||||
@@ -487,36 +474,11 @@ def eval_policy(
|
||||
return
|
||||
n_to_render_now = min(max_episodes_rendered - n_episodes_rendered, env.num_envs)
|
||||
if isinstance(env, gym.vector.SyncVectorEnv):
|
||||
frames = np.stack([env.envs[i].render() for i in range(n_to_render_now)]) # noqa: B023
|
||||
ep_frames.append(np.stack([env.envs[i].render() for i in range(n_to_render_now)])) # noqa: B023
|
||||
elif hasattr(env, "call"):
|
||||
# Here we must render all frames and discard any we don't need.
|
||||
# Covers AsyncVectorEnv and _LazyAsyncVectorEnv (which wraps one).
|
||||
frames = np.stack(env.call("render")[:n_to_render_now])
|
||||
else:
|
||||
return
|
||||
|
||||
# Overlay the high-level task and (for hierarchical policies like
|
||||
# pi052) the predicted low-level subtask onto each frame. Both are
|
||||
# best-effort: missing values just skip that line.
|
||||
try:
|
||||
tasks = list(env.call("task_description"))
|
||||
except (AttributeError, NotImplementedError):
|
||||
try:
|
||||
tasks = list(env.call("task"))
|
||||
except (AttributeError, NotImplementedError):
|
||||
tasks = None
|
||||
subtasks = getattr(policy, "last_subtasks", None)
|
||||
annotated = []
|
||||
for i in range(frames.shape[0]):
|
||||
subtask_i = subtasks[i] if subtasks is not None and i < len(subtasks) else None
|
||||
annotated.append(
|
||||
_annotate_eval_frames(
|
||||
frames[i : i + 1],
|
||||
tasks[i] if tasks is not None and i < len(tasks) else None,
|
||||
subtask_i,
|
||||
)[0]
|
||||
)
|
||||
ep_frames.append(np.stack(annotated))
|
||||
ep_frames.append(np.stack(env.call("render")[:n_to_render_now]))
|
||||
|
||||
if max_episodes_rendered > 0:
|
||||
video_paths: list[str] = []
|
||||
|
||||
@@ -151,7 +151,6 @@ Usage examples
|
||||
"""
|
||||
|
||||
import logging
|
||||
import sys
|
||||
|
||||
from lerobot.cameras.opencv import OpenCVCameraConfig # noqa: F401
|
||||
from lerobot.cameras.realsense import RealSenseCameraConfig # noqa: F401
|
||||
@@ -242,69 +241,10 @@ def rollout(cfg: RolloutConfig):
|
||||
logger.info("Rollout finished")
|
||||
|
||||
|
||||
_LANGUAGE_RUNTIME_FLAGS = {
|
||||
"--language",
|
||||
"--no_robot",
|
||||
"--sim",
|
||||
"--direct_subtask",
|
||||
"--sim.direct_subtask",
|
||||
"--disable_memory",
|
||||
"--fp8",
|
||||
}
|
||||
_LANGUAGE_RUNTIME_PREFIXES = (
|
||||
"--sim.",
|
||||
"--chunk_hz",
|
||||
"--ctrl_hz",
|
||||
"--high_level_hz",
|
||||
"--subtask_chunks_per_gen",
|
||||
"--text_min_new_tokens",
|
||||
"--text_temperature",
|
||||
"--text_top_p",
|
||||
)
|
||||
|
||||
|
||||
def _uses_language_runtime(argv: list[str]) -> bool:
|
||||
"""Return whether *argv* selects the interactive language runtime.
|
||||
|
||||
``--language`` is the explicit selector for real-robot runs whose other
|
||||
options overlap with the standard rollout CLI. Language-only options also
|
||||
select it automatically, which keeps the former language-runtime examples
|
||||
working after replacing their command name with ``lerobot-rollout``.
|
||||
"""
|
||||
return any(
|
||||
arg.split("=", 1)[0] in _LANGUAGE_RUNTIME_FLAGS or arg.startswith(_LANGUAGE_RUNTIME_PREFIXES)
|
||||
for arg in argv
|
||||
)
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None):
|
||||
"""CLI entry point for ``lerobot-rollout``.
|
||||
|
||||
Standard policy deployment continues through :class:`RolloutConfig`.
|
||||
Interactive language-conditioned and RoboCasa runs share this entry point
|
||||
and are selected with ``--language`` or any language-runtime-only option.
|
||||
"""
|
||||
def main():
|
||||
"""CLI entry point for ``lerobot-rollout``."""
|
||||
register_third_party_plugins()
|
||||
cli_args = list(sys.argv[1:] if argv is None else argv)
|
||||
if _uses_language_runtime(cli_args):
|
||||
from lerobot.runtime.cli import run as run_language_runtime
|
||||
|
||||
# ``--language`` is a dispatcher flag, not part of the runtime's own
|
||||
# argparse surface. All other arguments pass through unchanged.
|
||||
runtime_args = [arg for arg in cli_args if arg != "--language"]
|
||||
return run_language_runtime(runtime_args, prog="lerobot-rollout")
|
||||
|
||||
if argv is None:
|
||||
return rollout()
|
||||
|
||||
# draccus reads sys.argv. Supporting an explicit argv keeps this entry
|
||||
# point easy to smoke-test and mirrors the language-runtime branch above.
|
||||
previous_argv = sys.argv
|
||||
try:
|
||||
sys.argv = [previous_argv[0], *cli_args]
|
||||
return rollout()
|
||||
finally:
|
||||
sys.argv = previous_argv
|
||||
rollout()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -20,11 +20,9 @@ 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
|
||||
|
||||
@@ -83,7 +81,6 @@ 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.
|
||||
@@ -101,7 +98,6 @@ 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:
|
||||
@@ -169,20 +165,15 @@ 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)
|
||||
return train_metrics, output_dict
|
||||
|
||||
|
||||
@@ -210,7 +201,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
|
||||
|
||||
require_package("accelerate", extra="training")
|
||||
from accelerate import Accelerator
|
||||
from accelerate.utils import DistributedDataParallelKwargs, DistributedType, InitProcessGroupKwargs
|
||||
from accelerate.utils import DistributedDataParallelKwargs, DistributedType
|
||||
|
||||
cfg.validate()
|
||||
|
||||
@@ -219,16 +210,7 @@ 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:
|
||||
# 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))
|
||||
ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
|
||||
# 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"
|
||||
@@ -238,7 +220,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
|
||||
accelerator = Accelerator(
|
||||
step_scheduler_with_optimizer=False,
|
||||
mixed_precision=mixed_precision,
|
||||
kwargs_handlers=[ddp_kwargs, ipg_kwargs],
|
||||
kwargs_handlers=[ddp_kwargs],
|
||||
cpu=force_cpu,
|
||||
)
|
||||
|
||||
@@ -334,14 +316,6 @@ 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:
|
||||
@@ -350,13 +324,6 @@ 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},
|
||||
@@ -453,17 +420,13 @@ 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(
|
||||
from_indices,
|
||||
to_indices,
|
||||
dataset.meta.episodes["dataset_from_index"],
|
||||
dataset.meta.episodes["dataset_to_index"],
|
||||
episode_indices_to_use=dataset.episodes,
|
||||
drop_n_last_frames=getattr(active_cfg, "drop_n_last_frames", 0),
|
||||
shuffle=True,
|
||||
seed=seed,
|
||||
seed=cfg.seed if cfg.seed is not None else 0,
|
||||
absolute_to_relative_idx=dataset.absolute_to_relative_idx,
|
||||
)
|
||||
if cfg.resume and step > 0:
|
||||
@@ -501,8 +464,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
|
||||
# declares language columns; otherwise stay on PyTorch's default
|
||||
# collate so non-language training runs are unaffected.
|
||||
collate_fn = lerobot_collate_fn if dataset.meta.has_language_columns else None
|
||||
# Allow spawn/forkserver workers where forking large rank processes exhausts memory.
|
||||
mp_context = os.environ.get("LEROBOT_DATALOADER_MP_CONTEXT") or None
|
||||
dataloader = torch.utils.data.DataLoader(
|
||||
dataset,
|
||||
num_workers=cfg.num_workers,
|
||||
@@ -514,7 +475,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
|
||||
collate_fn=collate_fn,
|
||||
prefetch_factor=cfg.prefetch_factor if cfg.num_workers > 0 else None,
|
||||
persistent_workers=cfg.persistent_workers and cfg.num_workers > 0,
|
||||
multiprocessing_context=mp_context if cfg.num_workers > 0 else None,
|
||||
)
|
||||
|
||||
# Build eval dataloader if a held-out split exists
|
||||
@@ -542,7 +502,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
|
||||
collate_fn=eval_collate_fn,
|
||||
prefetch_factor=cfg.prefetch_factor if cfg.num_workers > 0 else None,
|
||||
persistent_workers=cfg.persistent_workers and cfg.num_workers > 0,
|
||||
multiprocessing_context=mp_context if cfg.num_workers > 0 else None,
|
||||
)
|
||||
|
||||
# Prepare everything with accelerator
|
||||
@@ -616,10 +575,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
|
||||
batch = preprocessor(batch)
|
||||
train_tracker.dataloading_s = time.perf_counter() - start_time
|
||||
|
||||
# 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, _ = update_policy(
|
||||
train_tracker,
|
||||
policy,
|
||||
batch,
|
||||
@@ -628,7 +584,6 @@ 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
|
||||
@@ -653,9 +608,10 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
|
||||
train_tracker.samples_per_s = effective_batch_size / step_time
|
||||
logging.info(train_tracker)
|
||||
if wandb_logger:
|
||||
# Policy sub-losses (latent_loss, action_loss, ...) are aggregated into the
|
||||
# tracker by update_policy, so to_dict() already carries their windowed,
|
||||
# rank-reduced averages — no per-step output_dict passthrough needed.
|
||||
wandb_log_dict = train_tracker.to_dict()
|
||||
if output_dict:
|
||||
wandb_log_dict.update(output_dict)
|
||||
# Log sample weighting statistics if enabled
|
||||
if sample_weighter is not None:
|
||||
weighter_stats = sample_weighter.get_stats()
|
||||
@@ -728,11 +684,10 @@ 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=eval_target_policy,
|
||||
policy=accelerator.unwrap_model(policy),
|
||||
env_preprocessor=env_preprocessor,
|
||||
env_postprocessor=env_postprocessor,
|
||||
preprocessor=preprocessor,
|
||||
|
||||
@@ -22,7 +22,7 @@ from torch.utils.data._utils.collate import default_collate
|
||||
|
||||
from lerobot.datasets.language import LANGUAGE_COLUMNS
|
||||
|
||||
_PYTHON_LIST_KEYS = {"messages", "message_streams", "target_message_indices", *LANGUAGE_COLUMNS}
|
||||
_PYTHON_LIST_KEYS = {"messages", "message_streams", "target_message_indices"}
|
||||
|
||||
|
||||
def lerobot_collate_fn(batch: list[dict[str, Any] | None]) -> dict[str, Any] | None:
|
||||
|
||||
@@ -26,7 +26,6 @@ OBS_IMAGES = OBS_IMAGE + "s"
|
||||
OBS_LANGUAGE = OBS_STR + ".language"
|
||||
OBS_LANGUAGE_TOKENS = OBS_LANGUAGE + ".tokens"
|
||||
OBS_LANGUAGE_ATTENTION_MASK = OBS_LANGUAGE + ".attention_mask"
|
||||
OBS_LANGUAGE_CAUSAL_MARKS = OBS_LANGUAGE + ".causal_marks"
|
||||
OBS_LANGUAGE_SUBTASK = OBS_STR + ".subtask"
|
||||
OBS_LANGUAGE_SUBTASK_TOKENS = OBS_LANGUAGE_SUBTASK + ".tokens"
|
||||
OBS_LANGUAGE_SUBTASK_ATTENTION_MASK = OBS_LANGUAGE_SUBTASK + ".attention_mask"
|
||||
@@ -35,7 +34,6 @@ ACTION = "action"
|
||||
ACTION_PREFIX = ACTION + "."
|
||||
ACTION_TOKENS = ACTION + ".tokens"
|
||||
ACTION_TOKEN_MASK = ACTION + ".token_mask"
|
||||
ACTION_CODE_TOKEN_MASK = ACTION + ".code_token_mask"
|
||||
REWARD = "next.reward"
|
||||
TRUNCATED = "next.truncated"
|
||||
DONE = "next.done"
|
||||
|
||||
@@ -59,6 +59,20 @@ def get_safe_torch_device(try_device: str, log: bool = False) -> torch.device:
|
||||
return device
|
||||
|
||||
|
||||
def resolve_safetensors_device(map_location: str | torch.device) -> str:
|
||||
"""Resolve a device string for a safetensors load, working around a device-mapping quirk.
|
||||
|
||||
safetensors' load maps the bare string "cuda" to cuda:0 regardless of the current device
|
||||
(unlike torch's .to("cuda"), which honors torch.cuda.current_device()). Under multi-GPU
|
||||
accelerate/FSDP every rank would then load its weights onto GPU 0, OOMing it before sharding.
|
||||
Resolve "cuda" to the concrete current-device index so each rank loads onto its own GPU.
|
||||
"""
|
||||
map_location = str(map_location)
|
||||
if map_location == "cuda" and torch.cuda.is_available():
|
||||
return f"cuda:{torch.cuda.current_device()}"
|
||||
return map_location
|
||||
|
||||
|
||||
def get_safe_dtype(dtype: torch.dtype, device: str | torch.device):
|
||||
"""
|
||||
mps is currently not compatible with float64
|
||||
|
||||
@@ -23,46 +23,6 @@ logger = logging.getLogger(__name__)
|
||||
JsonLike = str | int | float | bool | None | list["JsonLike"] | dict[str, "JsonLike"] | tuple["JsonLike", ...]
|
||||
|
||||
|
||||
class StreamingVideoWriter:
|
||||
"""Incrementally encode RGB frames to an MP4 without retaining them in memory."""
|
||||
|
||||
def __init__(self, video_path: str | Path, fps: int) -> None:
|
||||
from .import_utils import require_package
|
||||
|
||||
require_package("av", extra="av-dep")
|
||||
import av
|
||||
|
||||
self._av = av
|
||||
self._container = av.open(str(video_path), mode="w")
|
||||
self._stream = self._container.add_stream("libx264", rate=fps)
|
||||
self._shape: tuple[int, int] | None = None
|
||||
self.frames_written = 0
|
||||
|
||||
def add_frame(self, frame_array) -> None:
|
||||
orig_height, orig_width = frame_array.shape[:2]
|
||||
height = orig_height - orig_height % 2
|
||||
width = orig_width - orig_width % 2
|
||||
if self._shape is None:
|
||||
self._shape = (height, width)
|
||||
self._stream.width = width
|
||||
self._stream.height = height
|
||||
self._stream.pix_fmt = "yuv420p"
|
||||
elif self._shape != (height, width):
|
||||
raise ValueError(f"Video frame shape changed from {self._shape} to {(height, width)}")
|
||||
frame = self._av.VideoFrame.from_ndarray(frame_array[:height, :width], format="rgb24")
|
||||
for packet in self._stream.encode(frame):
|
||||
self._container.mux(packet)
|
||||
self.frames_written += 1
|
||||
|
||||
def close(self) -> None:
|
||||
if self._container is None:
|
||||
return
|
||||
for packet in self._stream.encode():
|
||||
self._container.mux(packet)
|
||||
self._container.close()
|
||||
self._container = None
|
||||
|
||||
|
||||
def load_json(fpath: Path) -> Any:
|
||||
"""Load data from a JSON file.
|
||||
|
||||
@@ -98,12 +58,36 @@ def write_video(video_path: str | Path, stacked_frames: list, fps: int) -> None:
|
||||
stacked_frames: List of HWC uint8 numpy arrays (RGB).
|
||||
fps: Frames per second for the output video.
|
||||
"""
|
||||
writer = StreamingVideoWriter(video_path, fps)
|
||||
try:
|
||||
from .import_utils import require_package
|
||||
|
||||
require_package("av", extra="av-dep")
|
||||
import av
|
||||
|
||||
with av.open(str(video_path), mode="w") as container:
|
||||
orig_height, orig_width = stacked_frames[0].shape[:2]
|
||||
# yuv420p requires even dimensions; crop by one pixel if needed
|
||||
height = orig_height if orig_height % 2 == 0 else orig_height - 1
|
||||
width = orig_width if orig_width % 2 == 0 else orig_width - 1
|
||||
if height != orig_height or width != orig_width:
|
||||
logger.warning(
|
||||
"Frame dimensions %dx%d are not even; cropping to %dx%d for yuv420p compatibility.",
|
||||
orig_width,
|
||||
orig_height,
|
||||
width,
|
||||
height,
|
||||
)
|
||||
stream = container.add_stream("libx264", rate=fps)
|
||||
stream.width = width
|
||||
stream.height = height
|
||||
stream.pix_fmt = "yuv420p"
|
||||
for frame_array in stacked_frames:
|
||||
writer.add_frame(frame_array)
|
||||
finally:
|
||||
writer.close()
|
||||
if height != orig_height or width != orig_width:
|
||||
frame_array = frame_array[:height, :width]
|
||||
frame = av.VideoFrame.from_ndarray(frame_array, format="rgb24")
|
||||
for packet in stream.encode(frame):
|
||||
container.mux(packet)
|
||||
for packet in stream.encode():
|
||||
container.mux(packet)
|
||||
|
||||
|
||||
def deserialize_json_into_object[T: JsonLike](fpath: Path, obj: T) -> T:
|
||||
|
||||
@@ -104,8 +104,7 @@ class MetricsTracker:
|
||||
"episodes",
|
||||
"epochs",
|
||||
"accelerator",
|
||||
"_tensor_sums",
|
||||
"_tensor_counts",
|
||||
"_caller_metrics",
|
||||
]
|
||||
|
||||
def __init__(
|
||||
@@ -131,8 +130,9 @@ class MetricsTracker:
|
||||
self.episodes = self.samples / self._avg_samples_per_ep
|
||||
self.epochs = self.samples / self._num_frames
|
||||
self.accelerator = accelerator
|
||||
self._tensor_sums: dict[str, torch.Tensor] = {}
|
||||
self._tensor_counts: dict[str, int] = {}
|
||||
# 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)
|
||||
|
||||
def __getattr__(self, name: str) -> int | dict[str, AverageMeter] | AverageMeter | Any:
|
||||
if name in self.__dict__:
|
||||
@@ -160,21 +160,20 @@ 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 update_metrics(self, values: dict[str, Any]) -> None:
|
||||
"""Accumulate a dict of scalar metrics, auto-registering a meter for each new key.
|
||||
|
||||
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()
|
||||
Non-numeric values and bools are ignored.
|
||||
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)):
|
||||
continue
|
||||
if name in self._caller_metrics:
|
||||
continue
|
||||
if name not in self.metrics:
|
||||
self.metrics[name] = AverageMeter(name, ":.3f", reduction="mean")
|
||||
self.metrics[name].update(float(value))
|
||||
|
||||
def reduce_across_ranks(self) -> None:
|
||||
"""
|
||||
@@ -196,21 +195,10 @@ class MetricsTracker:
|
||||
if not buckets:
|
||||
return
|
||||
|
||||
# NB: don't use ``accelerator.reduce(..., reduction="max")`` — accelerate only implements
|
||||
# "sum"/"mean" (it always all-reduces with SUM and divides for "mean"), so "max" silently
|
||||
# returns the SUM across ranks, inflating every "max" metric by ``num_processes`` (e.g. a
|
||||
# 3.5s step reported as 28s on 8 GPUs). Gather per-rank values and reduce them explicitly.
|
||||
device = self.accelerator.device
|
||||
num_processes = self.accelerator.num_processes
|
||||
for reduction, names in buckets.items():
|
||||
local = torch.tensor([self.metrics[n].avg for n in names], dtype=torch.float32, device=device)
|
||||
gathered = self.accelerator.gather(local).view(num_processes, len(names))
|
||||
if reduction == "max":
|
||||
reduced = gathered.amax(dim=0)
|
||||
elif reduction == "sum":
|
||||
reduced = gathered.sum(dim=0)
|
||||
else: # "mean"
|
||||
reduced = gathered.mean(dim=0)
|
||||
tensor = torch.tensor([self.metrics[n].avg for n in names], dtype=torch.float32, device=device)
|
||||
reduced = self.accelerator.reduce(tensor, reduction=reduction)
|
||||
for name, value in zip(names, reduced.tolist(), strict=True):
|
||||
meter = self.metrics[name]
|
||||
# Preserve avg == sum / count so a later .update() on this meter accumulates
|
||||
@@ -247,5 +235,3 @@ class MetricsTracker:
|
||||
"""Resets average meters."""
|
||||
for m in self.metrics.values():
|
||||
m.reset()
|
||||
self._tensor_sums.clear()
|
||||
self._tensor_counts.clear()
|
||||
|
||||
@@ -38,10 +38,7 @@ def _is_scalar(x):
|
||||
|
||||
|
||||
def init_rerun(
|
||||
session_name: str = "lerobot_control_loop",
|
||||
ip: str | None = None,
|
||||
port: int | None = None,
|
||||
web_port: int | None = None,
|
||||
session_name: str = "lerobot_control_loop", ip: str | None = None, port: int | None = None
|
||||
) -> None:
|
||||
"""
|
||||
Initializes the Rerun SDK for visualizing the control loop.
|
||||
@@ -50,7 +47,6 @@ def init_rerun(
|
||||
session_name: Name of the Rerun session.
|
||||
ip: Optional IP for connecting to a Rerun server.
|
||||
port: Optional port for connecting to a Rerun server.
|
||||
web_port: Serve a headless web viewer on this port, using ``port`` for gRPC.
|
||||
"""
|
||||
|
||||
require_package("rerun-sdk", extra="viz", import_name="rerun")
|
||||
@@ -64,10 +60,6 @@ def init_rerun(
|
||||
memory_limit = os.getenv("LEROBOT_RERUN_MEMORY_LIMIT", "10%")
|
||||
if ip and port:
|
||||
rr.connect_grpc(url=f"rerun+http://{ip}:{port}/proxy")
|
||||
elif web_port is not None:
|
||||
grpc_port = port or 9876
|
||||
url = rr.serve_grpc(grpc_port=grpc_port)
|
||||
rr.serve_web_viewer(web_port=web_port, open_browser=False, connect_to=url)
|
||||
else:
|
||||
rr.spawn(memory_limit=memory_limit)
|
||||
|
||||
|
||||
@@ -1,71 +0,0 @@
|
||||
# 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.
|
||||
|
||||
"""Best-effort text overlays shared by evaluation and interactive rollouts."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterable
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def annotate_frame(frame: np.ndarray, fields: Iterable[tuple[str, str | None]]) -> np.ndarray:
|
||||
"""Return an RGB frame annotated with the non-empty labeled ``fields``."""
|
||||
if frame.ndim != 3 or frame.shape[-1] != 3:
|
||||
return frame
|
||||
try:
|
||||
import cv2 # noqa: PLC0415
|
||||
except ImportError:
|
||||
return frame
|
||||
|
||||
text_rows = [f"{label}: {value}" for label, value in fields if value]
|
||||
if not text_rows:
|
||||
return frame
|
||||
|
||||
image = np.ascontiguousarray(frame).copy()
|
||||
font, scale, thickness, margin = cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1, 6
|
||||
max_width = image.shape[1] - 2 * margin
|
||||
lines: list[str] = []
|
||||
for text in text_rows:
|
||||
current = ""
|
||||
for word in text.split():
|
||||
candidate = f"{current} {word}".strip()
|
||||
width = cv2.getTextSize(candidate, font, scale, thickness)[0][0]
|
||||
if width > max_width and current:
|
||||
lines.append(current)
|
||||
current = word
|
||||
else:
|
||||
current = candidate
|
||||
if current:
|
||||
lines.append(current)
|
||||
|
||||
line_height = 20
|
||||
header_height = min(image.shape[0], len(lines) * line_height + 6)
|
||||
backdrop = image.copy()
|
||||
cv2.rectangle(backdrop, (0, 0), (image.shape[1], header_height), (0, 0, 0), -1)
|
||||
cv2.addWeighted(backdrop, 0.55, image, 0.45, 0, dst=image)
|
||||
|
||||
for index, line in enumerate(lines):
|
||||
cv2.putText(
|
||||
image,
|
||||
line,
|
||||
(margin, 18 + index * line_height),
|
||||
font,
|
||||
scale,
|
||||
(255, 255, 255),
|
||||
thickness,
|
||||
cv2.LINE_AA,
|
||||
)
|
||||
return image
|
||||
@@ -29,13 +29,6 @@ def test_message_recipe_validates_unknown_binding():
|
||||
)
|
||||
|
||||
|
||||
def test_canonical_recipe_loads():
|
||||
"""The canonical PI052 blend YAML loads + validates."""
|
||||
recipe = TrainingRecipe.from_yaml(Path("src/lerobot/configs/recipes/subtask_mem_vqa_speech.yaml"))
|
||||
assert recipe.blend is not None
|
||||
assert sum(c.weight for c in recipe.blend.values()) == pytest.approx(1.0)
|
||||
|
||||
|
||||
def test_message_turn_requires_a_stream():
|
||||
"""Every turn must declare a stream — None is rejected at construction.
|
||||
|
||||
|
||||
@@ -343,84 +343,6 @@ def test_resolve_task_explicit_override_beats_rephrasings():
|
||||
assert rendered["messages"][0]["content"] == "explicit override wins"
|
||||
|
||||
|
||||
def test_flow_only_low_level_recipe_renders_without_target():
|
||||
"""Regression: a flow-only ``low_level`` recipe has no ``target`` turn —
|
||||
its supervision is the action-expert flow loss, not text-CE. It must
|
||||
still render (not ``None``), otherwise every blend draw of it is dropped
|
||||
and the action expert never receives a flow loss."""
|
||||
recipe = TrainingRecipe(
|
||||
messages=[
|
||||
MessageTurn(
|
||||
role="user",
|
||||
content="${subtask}",
|
||||
stream="low_level",
|
||||
if_present="subtask",
|
||||
),
|
||||
],
|
||||
bindings={"subtask": "active_at(t, style=subtask)"},
|
||||
)
|
||||
|
||||
rendered = render_sample(
|
||||
recipe=recipe,
|
||||
persistent=PERSISTENT,
|
||||
events=[],
|
||||
t=0.5,
|
||||
sample_idx=0,
|
||||
task="clean kitchen",
|
||||
)
|
||||
|
||||
assert rendered is not None
|
||||
assert rendered["messages"] == [{"role": "user", "content": "subtask 0"}]
|
||||
assert rendered["message_streams"] == ["low_level"]
|
||||
assert rendered["target_message_indices"] == []
|
||||
|
||||
|
||||
def test_vqa_frame_is_consumed_over_the_weighted_blend():
|
||||
"""A frame carrying a VQA annotation renders the ``ask_vqa*`` sub-recipe
|
||||
even when its blend weight is tiny — VQA annotations are sparse and must
|
||||
never be wasted on a subtask/action draw."""
|
||||
recipe = TrainingRecipe(
|
||||
blend={
|
||||
"high_level_subtask": TrainingRecipe(
|
||||
weight=0.99,
|
||||
messages=[
|
||||
MessageTurn(role="user", content="${task}", stream="high_level"),
|
||||
MessageTurn(role="assistant", content="a subtask", stream="high_level", target=True),
|
||||
],
|
||||
),
|
||||
"ask_vqa_top": TrainingRecipe(
|
||||
weight=0.01,
|
||||
bindings={
|
||||
"vqa_query": "emitted_at(t, style=vqa, role=user, camera=observation.images.top)",
|
||||
"vqa": "emitted_at(t, style=vqa, role=assistant, camera=observation.images.top)",
|
||||
},
|
||||
messages=[
|
||||
MessageTurn(
|
||||
role="user", content="${vqa_query}", stream="high_level", if_present="vqa_query"
|
||||
),
|
||||
MessageTurn(
|
||||
role="assistant",
|
||||
content="${vqa}",
|
||||
stream="high_level",
|
||||
target=True,
|
||||
if_present="vqa",
|
||||
),
|
||||
],
|
||||
),
|
||||
}
|
||||
)
|
||||
# A frame WITH a vqa event renders VQA on every sample_idx, despite the
|
||||
# ask_vqa weight being only 0.01.
|
||||
for sample_idx in range(20):
|
||||
rendered = render_sample(
|
||||
recipe=recipe, persistent=PERSISTENT, events=EVENTS_AT_1, t=1.0, sample_idx=sample_idx, task="x"
|
||||
)
|
||||
assert rendered["messages"][-1]["content"] == '{"count": 2}', sample_idx
|
||||
# A frame WITHOUT a vqa event falls back to the normal weighted blend.
|
||||
rendered = render_sample(recipe=recipe, persistent=PERSISTENT, events=[], t=1.0, sample_idx=0, task="x")
|
||||
assert rendered["messages"][-1]["content"] == "a subtask"
|
||||
|
||||
|
||||
def test_emitted_at_persistent_tolerates_small_timestamp_drift():
|
||||
"""Persistent ``emitted_at`` should match within EMITTED_AT_TOLERANCE_S
|
||||
so callers that derive ``t`` arithmetically (``frame_idx / fps``) still
|
||||
|
||||
@@ -25,7 +25,7 @@ from datasets import Dataset # noqa: E402
|
||||
from lerobot.datasets.io_utils import (
|
||||
hf_transform_to_torch,
|
||||
)
|
||||
from lerobot.datasets.sampler import EpisodeAwareSampler, compute_sampler_state
|
||||
from lerobot.datasets.sampler import EpisodeAwareSampler
|
||||
|
||||
|
||||
def calculate_episode_data_index(hf_dataset: Dataset) -> dict[str, torch.Tensor]:
|
||||
@@ -154,6 +154,8 @@ def test_partial_episode_drop_warns(caplog):
|
||||
|
||||
# --- seeded (seed, epoch) shuffling, resume, and state ---
|
||||
|
||||
from lerobot.datasets.sampler import compute_sampler_state # noqa: E402
|
||||
|
||||
EPISODE_BOUNDS = ([0, 2, 3], [2, 3, 6]) # episodes of 2, 1 and 3 frames
|
||||
|
||||
|
||||
|
||||
@@ -1,48 +0,0 @@
|
||||
#!/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
|
||||
|
||||
from lerobot.envs.robocasa import RoboCasaEnv, convert_action
|
||||
|
||||
|
||||
def test_robocasa_action_uses_openpi_checkpoint_order():
|
||||
action = np.arange(12, dtype=np.float32)
|
||||
|
||||
converted = convert_action(action)
|
||||
|
||||
np.testing.assert_array_equal(converted["action.end_effector_position"], [0, 1, 2])
|
||||
np.testing.assert_array_equal(converted["action.end_effector_rotation"], [3, 4, 5])
|
||||
np.testing.assert_array_equal(converted["action.gripper_close"], [6])
|
||||
np.testing.assert_array_equal(converted["action.base_motion"], [7, 8, 9, 10])
|
||||
np.testing.assert_array_equal(converted["action.control_mode"], [11])
|
||||
|
||||
|
||||
def test_robocasa_state_uses_openpi_checkpoint_order():
|
||||
env = object.__new__(RoboCasaEnv)
|
||||
env.obs_type = "pixels_agent_pos"
|
||||
env.camera_name = []
|
||||
raw_observation = {
|
||||
"state.end_effector_position_relative": np.arange(0, 3),
|
||||
"state.end_effector_rotation_relative": np.arange(3, 7),
|
||||
"state.base_position": np.arange(7, 10),
|
||||
"state.base_rotation": np.arange(10, 14),
|
||||
"state.gripper_qpos": np.arange(14, 16),
|
||||
}
|
||||
|
||||
observation = env._format_raw_obs(raw_observation)
|
||||
|
||||
np.testing.assert_array_equal(observation["agent_pos"], np.arange(16, dtype=np.float32))
|
||||
@@ -0,0 +1,245 @@
|
||||
#!/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.
|
||||
|
||||
"""Behavior-pinning tests for the shared flow-matching sampling primitives.
|
||||
|
||||
``euler_integrate`` is compared against a verbatim copy of the historical pi0/pi05/
|
||||
smolvla sampling loop (including its RTC hook semantics): any divergence from that
|
||||
reference is a behavior change for released checkpoints.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
from lerobot.policies.common.flow_matching import (
|
||||
_beta_distribution,
|
||||
euler_integrate,
|
||||
sample_beta,
|
||||
sample_noise,
|
||||
sample_time_beta,
|
||||
)
|
||||
|
||||
|
||||
def test_sample_beta_range_dtype_and_reproducibility():
|
||||
torch.manual_seed(0)
|
||||
s1 = sample_beta(1.5, 1.0, 4096, "cpu")
|
||||
torch.manual_seed(0)
|
||||
s2 = sample_beta(1.5, 1.0, 4096, "cpu")
|
||||
assert torch.equal(s1, s2)
|
||||
assert s1.shape == (4096,) and s1.dtype == torch.float32
|
||||
assert s1.min() >= 0.0 and s1.max() <= 1.0
|
||||
# Beta(1.5, 1.0) mean is 1.5/2.5 = 0.6.
|
||||
assert abs(s1.mean().item() - 0.6) < 0.02
|
||||
|
||||
|
||||
def test_sample_time_beta_openpi_convention():
|
||||
torch.manual_seed(1)
|
||||
time = sample_time_beta(4096, "cpu", alpha=1.5, beta=1.0, scale=0.999, offset=0.001)
|
||||
assert time.dtype == torch.float32
|
||||
assert time.min() >= 0.001 and time.max() <= 1.0
|
||||
# Exact composition: Beta sample * scale + offset, same RNG stream.
|
||||
torch.manual_seed(1)
|
||||
expected = sample_beta(1.5, 1.0, 4096, "cpu") * 0.999 + 0.001
|
||||
torch.testing.assert_close(time, expected, rtol=0, atol=0)
|
||||
|
||||
|
||||
def test_sample_noise_seeded():
|
||||
torch.manual_seed(2)
|
||||
n1 = sample_noise((2, 8, 4), "cpu")
|
||||
torch.manual_seed(2)
|
||||
n2 = sample_noise((2, 8, 4), "cpu")
|
||||
assert torch.equal(n1, n2)
|
||||
assert n1.dtype == torch.float32 and n1.shape == (2, 8, 4)
|
||||
|
||||
|
||||
def test_sample_noise_normal_is_default():
|
||||
torch.manual_seed(2)
|
||||
default = sample_noise((2, 8, 4), "cpu")
|
||||
torch.manual_seed(2)
|
||||
normal = sample_noise((2, 8, 4), "cpu", distribution="normal")
|
||||
assert torch.equal(default, normal)
|
||||
|
||||
|
||||
def test_sample_noise_uniform_evo1():
|
||||
torch.manual_seed(2)
|
||||
n = sample_noise((4096,), "cpu", distribution="uniform")
|
||||
assert n.dtype == torch.float32
|
||||
assert n.min() >= -1.0 and n.max() < 1.0
|
||||
# evo1's rand_like * 2 - 1 has mean ~0 over [-1, 1).
|
||||
assert abs(n.mean().item()) < 0.05
|
||||
# Exact match to the historical evo1 expression on the same RNG stream.
|
||||
torch.manual_seed(2)
|
||||
expected = torch.rand((4096,), dtype=torch.float32) * 2 - 1
|
||||
torch.testing.assert_close(n, expected, rtol=0, atol=0)
|
||||
|
||||
|
||||
def test_sample_noise_invalid_distribution():
|
||||
import pytest
|
||||
|
||||
with pytest.raises(ValueError, match="Unknown noise distribution"):
|
||||
sample_noise((2, 2), "cpu", distribution="bogus")
|
||||
|
||||
|
||||
def test_sample_time_beta_forward_complement_convention():
|
||||
# groot/wall_x forward convention: t = (1 - beta) * 0.999.
|
||||
torch.manual_seed(9)
|
||||
time = sample_time_beta(4096, "cpu", alpha=1.5, beta=1.0, scale=0.999, complement=True)
|
||||
torch.manual_seed(9)
|
||||
expected = (1.0 - sample_beta(1.5, 1.0, 4096, "cpu")) * 0.999
|
||||
torch.testing.assert_close(time, expected, rtol=0, atol=0)
|
||||
|
||||
|
||||
def test_sample_time_beta_evo1_clamp():
|
||||
# evo1: Beta(2, 2) clamped to [0.02, 0.98].
|
||||
torch.manual_seed(10)
|
||||
time = sample_time_beta(4096, "cpu", alpha=2.0, beta=2.0, clamp_min=0.02, clamp_max=0.98)
|
||||
assert time.min() >= 0.02 and time.max() <= 0.98
|
||||
|
||||
|
||||
def test_sample_beta_distribution_is_cached():
|
||||
a = _beta_distribution(1.5, 1.0)
|
||||
b = _beta_distribution(1.5, 1.0)
|
||||
assert a is b
|
||||
assert _beta_distribution(2.0, 2.0) is not a
|
||||
|
||||
|
||||
def test_euler_integrate_constant_velocity_is_exact():
|
||||
# With v_t == c constant, x_0 = x_1 + sum(dt * c) = x_1 - c exactly (num_steps * dt = -1).
|
||||
noise = torch.randn(3, 5, 2)
|
||||
c = torch.randn(3, 5, 2)
|
||||
out = euler_integrate(lambda x_t, time: c, noise, num_steps=10)
|
||||
torch.testing.assert_close(out, noise - c, rtol=0, atol=1e-6)
|
||||
|
||||
|
||||
def _reference_pi0_loop(denoise_fn, noise, num_steps, rtc_enabled, rtc_processor, kw):
|
||||
"""Verbatim structure of the historical pi0/pi05/smolvla sample_actions loop."""
|
||||
bsize = noise.shape[0]
|
||||
device = noise.device
|
||||
dt = -1.0 / num_steps
|
||||
x_t = noise
|
||||
for step in range(num_steps):
|
||||
time = 1.0 + step * dt
|
||||
time_tensor = torch.tensor(time, dtype=torch.float32, device=device).expand(bsize)
|
||||
|
||||
def denoise_step_partial_call(input_x_t, current_timestep=time_tensor):
|
||||
return denoise_fn(input_x_t, current_timestep)
|
||||
|
||||
if rtc_enabled:
|
||||
v_t = rtc_processor.denoise_step(
|
||||
x_t=x_t,
|
||||
prev_chunk_left_over=kw.get("prev_chunk_left_over"),
|
||||
inference_delay=kw.get("inference_delay"),
|
||||
time=time,
|
||||
original_denoise_step_partial=denoise_step_partial_call,
|
||||
execution_horizon=kw.get("execution_horizon"),
|
||||
)
|
||||
else:
|
||||
v_t = denoise_step_partial_call(x_t)
|
||||
x_t = x_t + dt * v_t
|
||||
if rtc_processor is not None and rtc_processor.is_debug_enabled():
|
||||
rtc_processor.track(time=time, x_t=x_t, v_t=v_t)
|
||||
return x_t
|
||||
|
||||
|
||||
class _StubRTCProcessor:
|
||||
def __init__(self, debug_enabled: bool):
|
||||
self._debug = debug_enabled
|
||||
self.tracked = []
|
||||
self.guidance_calls = []
|
||||
|
||||
def is_debug_enabled(self):
|
||||
return self._debug
|
||||
|
||||
def denoise_step(
|
||||
self,
|
||||
x_t,
|
||||
prev_chunk_left_over,
|
||||
inference_delay,
|
||||
time,
|
||||
original_denoise_step_partial,
|
||||
execution_horizon,
|
||||
):
|
||||
self.guidance_calls.append(
|
||||
{
|
||||
"time": time,
|
||||
"inference_delay": inference_delay,
|
||||
"execution_horizon": execution_horizon,
|
||||
"x_t": x_t.clone(),
|
||||
}
|
||||
)
|
||||
return original_denoise_step_partial(x_t) * 0.5
|
||||
|
||||
def track(self, time, x_t, v_t):
|
||||
self.tracked.append({"time": time, "x_t": x_t.clone(), "v_t": v_t.clone()})
|
||||
|
||||
|
||||
def _make_denoise_fn():
|
||||
weight = torch.randn(4, 4) * 0.1
|
||||
|
||||
def denoise_fn(x_t, time_tensor):
|
||||
return x_t @ weight + time_tensor[:, None, None]
|
||||
|
||||
return denoise_fn
|
||||
|
||||
|
||||
def test_euler_integrate_matches_historical_loop():
|
||||
torch.manual_seed(3)
|
||||
denoise_fn = _make_denoise_fn()
|
||||
noise = torch.randn(2, 6, 4)
|
||||
ref = _reference_pi0_loop(denoise_fn, noise, 10, rtc_enabled=False, rtc_processor=None, kw={})
|
||||
out = euler_integrate(denoise_fn, noise, 10)
|
||||
assert torch.equal(out, ref)
|
||||
|
||||
|
||||
def test_euler_integrate_rtc_guidance_and_kwarg_forwarding():
|
||||
torch.manual_seed(4)
|
||||
denoise_fn = _make_denoise_fn()
|
||||
noise = torch.randn(2, 6, 4)
|
||||
leftover = torch.randn(2, 6, 4)
|
||||
kw = {"inference_delay": 3, "prev_chunk_left_over": leftover, "execution_horizon": 25}
|
||||
|
||||
ref_proc, new_proc = _StubRTCProcessor(False), _StubRTCProcessor(False)
|
||||
ref = _reference_pi0_loop(denoise_fn, noise, 6, rtc_enabled=True, rtc_processor=ref_proc, kw=kw)
|
||||
out = euler_integrate(
|
||||
denoise_fn,
|
||||
noise,
|
||||
6,
|
||||
rtc_processor=new_proc,
|
||||
rtc_enabled=True,
|
||||
inference_delay=3,
|
||||
prev_chunk_left_over=leftover,
|
||||
execution_horizon=25,
|
||||
)
|
||||
assert torch.equal(out, ref)
|
||||
assert len(new_proc.guidance_calls) == 6
|
||||
for ref_call, new_call in zip(ref_proc.guidance_calls, new_proc.guidance_calls, strict=True):
|
||||
assert ref_call["time"] == new_call["time"]
|
||||
assert new_call["inference_delay"] == 3 and new_call["execution_horizon"] == 25
|
||||
# Guidance sees the PRE-update x_t.
|
||||
assert torch.equal(ref_call["x_t"], new_call["x_t"])
|
||||
|
||||
|
||||
def test_euler_integrate_debug_tracking_fires_even_when_rtc_disabled():
|
||||
# Historical behavior: track() fires whenever the processor exists and has debugging
|
||||
# enabled, independent of whether RTC guidance is active.
|
||||
torch.manual_seed(5)
|
||||
denoise_fn = _make_denoise_fn()
|
||||
noise = torch.randn(2, 6, 4)
|
||||
proc = _StubRTCProcessor(True)
|
||||
out = euler_integrate(denoise_fn, noise, 4, rtc_processor=proc, rtc_enabled=False)
|
||||
assert len(proc.guidance_calls) == 0
|
||||
assert len(proc.tracked) == 4
|
||||
# track() receives the POST-update x_t; the last one is the returned sample.
|
||||
assert torch.equal(proc.tracked[-1]["x_t"], out)
|
||||
@@ -0,0 +1,195 @@
|
||||
#!/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.
|
||||
|
||||
"""Behavior-pinning tests for the shared VLA helpers.
|
||||
|
||||
These helpers are the canonical versions of functions that used to be copy-pasted across
|
||||
the openpi-derived policies (pi0, pi05, pi0_fast, smolvla, eo1, xvla). The expected
|
||||
values below encode the historical per-policy behavior exactly; a failure here means a
|
||||
behavior change that would silently affect released checkpoints.
|
||||
"""
|
||||
|
||||
import math
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from lerobot.policies.common.vla_utils import (
|
||||
create_sinusoidal_pos_embedding,
|
||||
make_att_2d_masks,
|
||||
pad_vector,
|
||||
prepare_attention_masks_4d,
|
||||
resize_with_pad,
|
||||
resize_with_pad_torch,
|
||||
)
|
||||
from lerobot.utils.constants import OPENPI_ATTENTION_MASK_VALUE
|
||||
|
||||
|
||||
def test_create_sinusoidal_pos_embedding_matches_openpi_formula():
|
||||
time = torch.tensor([0.0, 0.25, 1.0])
|
||||
dim, min_period, max_period = 8, 4e-3, 4.0
|
||||
emb = create_sinusoidal_pos_embedding(time, dim, min_period, max_period, device=torch.device("cpu"))
|
||||
|
||||
assert emb.shape == (3, dim)
|
||||
# Independent recomputation of the openpi formula in float64.
|
||||
fraction = torch.linspace(0.0, 1.0, dim // 2, dtype=torch.float64)
|
||||
period = min_period * (max_period / min_period) ** fraction
|
||||
scaling = 1.0 / period * 2 * math.pi
|
||||
sin_input = scaling[None, :] * time.to(torch.float64)[:, None]
|
||||
expected = torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1)
|
||||
torch.testing.assert_close(emb, expected, rtol=1e-9, atol=1e-9)
|
||||
|
||||
|
||||
def test_create_sinusoidal_pos_embedding_validation():
|
||||
with pytest.raises(ValueError, match="divisible by 2"):
|
||||
create_sinusoidal_pos_embedding(torch.zeros(2), 7, 4e-3, 4.0, device=torch.device("cpu"))
|
||||
with pytest.raises(ValueError, match="batch_size"):
|
||||
create_sinusoidal_pos_embedding(torch.zeros(2, 2), 8, 4e-3, 4.0, device=torch.device("cpu"))
|
||||
|
||||
|
||||
def test_make_att_2d_masks_docstring_cases():
|
||||
# Pure causal attention: [[1 1 1]]
|
||||
pad = torch.ones(1, 3, dtype=torch.bool)
|
||||
att = torch.tensor([[1, 1, 1]], dtype=torch.int32)
|
||||
expected = torch.tensor([[[1, 0, 0], [1, 1, 0], [1, 1, 1]]], dtype=torch.bool)
|
||||
assert torch.equal(make_att_2d_masks(pad, att), expected)
|
||||
|
||||
# Prefix-LM: [[0 0 1 1]] -> first two tokens attend bidirectionally, rest causal.
|
||||
att = torch.tensor([[0, 0, 1, 1]], dtype=torch.int32)
|
||||
pad = torch.ones(1, 4, dtype=torch.bool)
|
||||
expected = torch.tensor([[[1, 1, 0, 0], [1, 1, 0, 0], [1, 1, 1, 0], [1, 1, 1, 1]]], dtype=torch.bool)
|
||||
assert torch.equal(make_att_2d_masks(pad, att), expected)
|
||||
|
||||
# Padding removes rows and columns.
|
||||
pad = torch.tensor([[True, True, False]])
|
||||
att = torch.tensor([[0, 1, 1]], dtype=torch.int32)
|
||||
out = make_att_2d_masks(pad, att)
|
||||
assert not out[0, :, 2].any() and not out[0, 2, :].any()
|
||||
|
||||
|
||||
def test_make_att_2d_masks_validation():
|
||||
with pytest.raises(ValueError):
|
||||
make_att_2d_masks(torch.ones(3, dtype=torch.bool), torch.ones(1, 3, dtype=torch.int32))
|
||||
with pytest.raises(ValueError):
|
||||
make_att_2d_masks(torch.ones(1, 3, dtype=torch.bool), torch.ones(3, dtype=torch.int32))
|
||||
|
||||
|
||||
def test_prepare_attention_masks_4d():
|
||||
masks = torch.tensor([[[True, False], [False, True]]])
|
||||
out = prepare_attention_masks_4d(masks)
|
||||
assert out.shape == (1, 1, 2, 2)
|
||||
expected = torch.tensor([[[[0.0, OPENPI_ATTENTION_MASK_VALUE], [OPENPI_ATTENTION_MASK_VALUE, 0.0]]]])
|
||||
assert torch.equal(out, expected)
|
||||
|
||||
out_bf16 = prepare_attention_masks_4d(masks, dtype=torch.bfloat16)
|
||||
assert out_bf16.dtype == torch.bfloat16
|
||||
assert torch.equal(out_bf16, expected.to(torch.bfloat16))
|
||||
|
||||
|
||||
def test_pad_vector_openpi_semantics():
|
||||
v = torch.arange(6.0).reshape(2, 3)
|
||||
padded = pad_vector(v, 5)
|
||||
assert padded.shape == (2, 5)
|
||||
assert torch.equal(padded[:, :3], v) and not padded[:, 3:].any()
|
||||
# Already large enough (>=): returned unchanged, same object.
|
||||
assert pad_vector(v, 3) is v
|
||||
assert pad_vector(v, 2) is v
|
||||
# 3D input.
|
||||
v3 = torch.ones(2, 4, 3)
|
||||
assert pad_vector(v3, 7).shape == (2, 4, 7)
|
||||
|
||||
|
||||
def test_pad_vector_truncate_semantics():
|
||||
v = torch.arange(6.0).reshape(2, 3)
|
||||
out = pad_vector(v, 2, truncate=True)
|
||||
assert out.shape == (2, 2) and torch.equal(out, v[:, :2])
|
||||
out = pad_vector(v, 5, truncate=True)
|
||||
assert out.shape == (2, 5) and torch.equal(out[:, :3], v) and not out[:, 3:].any()
|
||||
assert pad_vector(v, 0, truncate=True).shape == (2, 0)
|
||||
assert pad_vector(v, 3, truncate=True) is v
|
||||
|
||||
|
||||
@pytest.mark.parametrize("channels_last", [True, False])
|
||||
def test_resize_with_pad_torch_centered(channels_last):
|
||||
img = torch.rand(2, 3, 30, 60) if not channels_last else torch.rand(2, 30, 60, 3)
|
||||
out = resize_with_pad_torch(img, 64, 64)
|
||||
if channels_last:
|
||||
assert out.shape == (2, 64, 64, 3)
|
||||
# Aspect ratio preserved: 30x60 -> 32x64, padded 16 top and 16 bottom (centered).
|
||||
assert not out[:, :16].any() and not out[:, -16:].any()
|
||||
assert out[:, 16:48].abs().sum() > 0
|
||||
else:
|
||||
assert out.shape == (2, 3, 64, 64)
|
||||
assert not out[:, :, :16].any() and not out[:, :, -16:].any()
|
||||
|
||||
|
||||
def test_resize_with_pad_torch_uint8_roundtrip():
|
||||
img = (torch.rand(1, 3, 20, 20) * 255).to(torch.uint8)
|
||||
out = resize_with_pad_torch(img, 40, 40)
|
||||
assert out.dtype == torch.uint8 and out.shape == (1, 3, 40, 40)
|
||||
with pytest.raises(ValueError, match="Unsupported image dtype"):
|
||||
resize_with_pad_torch(torch.rand(1, 3, 8, 8, dtype=torch.float64), 16, 16)
|
||||
|
||||
|
||||
def test_resize_with_pad_top_left():
|
||||
img = torch.rand(2, 3, 30, 60)
|
||||
out = resize_with_pad(img, 64, 64, pad_value=-1.0)
|
||||
assert out.shape == (2, 3, 64, 64)
|
||||
# 30x60 -> 32x64; this variant pads on the TOP only (32 rows of pad_value).
|
||||
assert torch.equal(out[:, :, :32], torch.full((2, 3, 32, 64), -1.0))
|
||||
assert out[:, :, 32:].min() >= 0
|
||||
# No-op fast path returns the same object.
|
||||
assert resize_with_pad(img, 30, 60, pad_value=0.0) is img
|
||||
with pytest.raises(ValueError, match="expected"):
|
||||
resize_with_pad(torch.rand(3, 8, 8), 16, 16, pad_value=0.0)
|
||||
|
||||
|
||||
def test_clone_past_key_values():
|
||||
pytest.importorskip("transformers")
|
||||
from transformers import DynamicCache
|
||||
|
||||
from lerobot.policies.common.vla_utils import clone_past_key_values
|
||||
|
||||
cache = DynamicCache()
|
||||
keys, values = torch.rand(1, 2, 4, 8), torch.rand(1, 2, 4, 8)
|
||||
cache.update(keys, values, 0)
|
||||
cloned = clone_past_key_values(cache)
|
||||
(ck, cv, _), (ok, ov, _) = next(iter(cloned)), next(iter(cache))
|
||||
assert torch.equal(ck, ok) and torch.equal(cv, ov)
|
||||
# Deep copy: mutating the clone must not touch the original.
|
||||
ck.zero_()
|
||||
assert not torch.equal(ck, ok)
|
||||
|
||||
|
||||
def test_clone_past_key_values_is_fullgraph_compilable():
|
||||
pytest.importorskip("transformers")
|
||||
from transformers import DynamicCache
|
||||
|
||||
from lerobot.policies.common.vla_utils import clone_past_key_values
|
||||
|
||||
cache = DynamicCache()
|
||||
keys, values = torch.rand(1, 2, 4, 8), torch.rand(1, 2, 4, 8)
|
||||
cache.update(keys, values, 0)
|
||||
|
||||
compiled_clone = torch.compile(clone_past_key_values, backend="eager", fullgraph=True)
|
||||
cloned = compiled_clone(cache)
|
||||
|
||||
(cloned_keys, cloned_values, _), (original_keys, original_values, _) = (
|
||||
next(iter(cloned)),
|
||||
next(iter(cache)),
|
||||
)
|
||||
assert torch.equal(cloned_keys, original_keys)
|
||||
assert torch.equal(cloned_values, original_values)
|
||||
@@ -1,151 +0,0 @@
|
||||
#!/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
|
||||
@@ -1,146 +0,0 @@
|
||||
#!/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
|
||||
@@ -1,152 +0,0 @@
|
||||
#!/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,
|
||||
)
|
||||
@@ -1,162 +0,0 @@
|
||||
#!/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)
|
||||
@@ -1,122 +0,0 @@
|
||||
#!/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)
|
||||
@@ -1,65 +0,0 @@
|
||||
#!/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
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user