feat(pi052): add language-supervised policy

tests/policies/pi052/test_pi052_attention_masking.py

@@ -0,0 +1,151 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Attention-masking tests for the PI052 (π0.5 v2) text head.
+
+Regression coverage for the text-CE collapse bug: PaliGemma's
+``embed_prefix`` flags every language token ``att=0``, which
+``make_att_2d_masks`` turns into one fully *bidirectional* block. Under
+that mask the text cross-entropy degenerates into a copy task — a
+supervised target token attends to the tokens it is trained to predict —
+and the LM head never learns causal generation, so ``select_message``
+collapses at inference.
+
+``_mark_target_span_causal`` sets ``att=1`` on the supervised target
+language positions so each target token attends causally among the
+targets while staying bidirectional to images + the user prompt. These
+tests pin that behaviour for the PaliGemma prefix layout.
+"""
+
+import pytest
+import torch
+
+# modeling_pi052 / modeling_pi05 import transformers transitively.
+pytest.importorskip("transformers")
+
+from lerobot.policies.pi05.modeling_pi05 import make_att_2d_masks  # noqa: E402
+from lerobot.policies.pi052.modeling_pi052 import (  # noqa: E402
+    _mark_target_span_causal,
+    _shifted_lin_ce,
+)
+
+
+def _shifted_ce(logits, labels):
+    """Adapter: ``_shifted_lin_ce`` is Liger-fused (hidden @ lm_head_weightᵀ).
+
+    An identity ``lm_head_weight`` makes the computed logits equal ``logits``.
+    Liger's Triton kernel is GPU-only, so inputs run on CUDA; the loss is
+    returned on CPU so grad still flows back to the CPU ``logits`` leaf.
+    """
+    if not torch.cuda.is_available():
+        pytest.skip("Liger fused CE requires CUDA")
+    vocab_size = logits.shape[-1]
+    eye = torch.eye(vocab_size, dtype=logits.dtype, device="cuda")
+    return _shifted_lin_ce(logits.cuda(), eye, labels.cuda()).cpu()
+
+
+# Synthetic prefix: two image tokens, three prompt tokens, and four supervised target tokens.
+# Text labels mask the prompt with -100 and cover the target through the prefix end.
+N_IMAGE = 2
+N_PROMPT = 3
+N_TARGET = 4
+LANG_START = N_IMAGE
+LANG_END = N_IMAGE + N_PROMPT + N_TARGET  # = prefix length
+PREFIX_LEN = LANG_END
+
+
+def _embed_prefix_att_masks() -> torch.Tensor:
+    """Mimic PaliGemma ``embed_prefix``: images + lang all att=0."""
+    return torch.zeros(1, PREFIX_LEN, dtype=torch.bool)
+
+
+def _text_labels() -> torch.Tensor:
+    """-100 over the prompt span, real ids over the target span."""
+    labels = torch.full((1, N_PROMPT + N_TARGET), -100, dtype=torch.long)
+    labels[0, N_PROMPT:] = torch.arange(10, 10 + N_TARGET)
+    return labels
+
+
+def _attends(prefix_att_masks: torch.Tensor) -> torch.Tensor:
+    """2D boolean attendance matrix; ``[i, j]`` True ⇒ i attends to j."""
+    pad = torch.ones(1, PREFIX_LEN, dtype=torch.bool)
+    return make_att_2d_masks(pad, prefix_att_masks)[0]
+
+
+def test_mark_sets_att_on_targets_only():
+    """Only the supervised target language positions flip to att=1."""
+    marked = _mark_target_span_causal(_embed_prefix_att_masks(), _text_labels(), LANG_START, LANG_END)
+    expected = [False] * PREFIX_LEN
+    for i in range(LANG_START + N_PROMPT, LANG_END):  # target span
+        expected[i] = True
+    assert marked[0].tolist() == expected
+
+
+def test_target_tokens_attend_causally_among_themselves():
+    """A target token must NOT attend to later targets, but must attend
+    to earlier ones — genuine causal next-token prediction."""
+    marked = _mark_target_span_causal(_embed_prefix_att_masks(), _text_labels(), LANG_START, LANG_END)
+    attends = _attends(marked)
+    tgt = range(LANG_START + N_PROMPT, LANG_END)
+    for i in tgt:
+        for j in tgt:
+            if j > i:
+                assert not attends[i, j], f"target {i} must not see future target {j}"
+            else:
+                assert attends[i, j], f"target {i} must see earlier/self target {j}"
+
+
+def test_target_tokens_attend_prompt_and_images_bidirectionally():
+    """Targets keep full visibility of images + the user prompt."""
+    marked = _mark_target_span_causal(_embed_prefix_att_masks(), _text_labels(), LANG_START, LANG_END)
+    attends = _attends(marked)
+    context = list(range(0, LANG_START + N_PROMPT))  # images + prompt
+    for i in range(LANG_START + N_PROMPT, LANG_END):
+        for j in context:
+            assert attends[i, j], f"target {i} must attend context {j}"
+
+
+def test_non_target_subtask_stays_bidirectional():
+    """A flow-only / non-target language span (all -100 labels) leaves the
+    mask untouched — the action expert reads it bidirectionally."""
+    all_ignored = torch.full((1, N_PROMPT + N_TARGET), -100, dtype=torch.long)
+    marked = _mark_target_span_causal(_embed_prefix_att_masks(), all_ignored, LANG_START, LANG_END)
+    assert torch.equal(marked, _embed_prefix_att_masks())
+
+
+def test_unmarked_mask_is_bidirectional_the_bug():
+    """Documents the bug the fix prevents: without ``_mark_target_span_causal``
+    a target token attends *bidirectionally* to later targets — the
+    text-CE can copy the answer it is trained to predict."""
+    attends = _attends(_embed_prefix_att_masks())
+    first_tgt = LANG_START + N_PROMPT
+    last_tgt = LANG_END - 1
+    assert attends[first_tgt, last_tgt], (
+        "raw embed_prefix mask is bidirectional over language — the first "
+        "target token can see the last, which is the collapse bug"
+    )
+
+
+def test_shifted_ce_returns_zero_when_no_text_positions_are_supervised():
+    pytest.importorskip("liger_kernel")
+    logits = torch.randn(2, 4, 8, requires_grad=True)
+    labels = torch.full((2, 4), -100, dtype=torch.long)
+
+    loss = _shifted_ce(logits, labels)
+
+    assert loss.item() == 0
+    loss.backward()
+    assert logits.grad is not None

tests/policies/pi052/test_pi052_bucketed_ce.py

@@ -0,0 +1,146 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+
+pytest.importorskip("transformers")
+
+from lerobot.policies.pi052.modeling_pi052 import _lin_ce_flat, _shifted_lin_ce
+
+
+def test_shifted_ce_none_retains_distinct_per_sample_losses():
+    hidden = torch.tensor(
+        [
+            [[8.0, 0.0], [0.0, 8.0], [0.0, 0.0]],
+            [[0.0, 8.0], [8.0, 0.0], [0.0, 0.0]],
+        ]
+    )
+    labels = torch.tensor([[0, 0, 1], [0, 0, 1]])
+    losses = _shifted_lin_ce(hidden, torch.eye(2), labels, reduction="none")
+
+    assert losses.shape == (2,)
+    assert losses[0] < losses[1]
+
+
+def test_checkpoint_resolution_forwards_explicit_hub_options(monkeypatch, tmp_path):
+    import lerobot.policies.pi05.modeling_pi05 as modeling_pi05
+
+    checkpoint = tmp_path / "model.safetensors"
+    checkpoint.touch()
+    calls = []
+
+    def fake_cached_file(model_id, filename, **kwargs):
+        calls.append((model_id, filename, kwargs))
+        return None if filename.endswith("index.json") else str(checkpoint)
+
+    monkeypatch.setattr(modeling_pi05, "cached_file", fake_cached_file)
+    files = modeling_pi05._resolve_weight_files(
+        "org/model",
+        force_download=True,
+        resume_download=True,
+        proxies={"https": "proxy"},
+        token="secret",
+        cache_dir=tmp_path / "cache",
+        local_files_only=True,
+        revision="commit",
+    )
+
+    assert files == [checkpoint]
+    for _model_id, _filename, kwargs in calls:
+        assert kwargs["revision"] == "commit"
+        assert kwargs["cache_dir"] == tmp_path / "cache"
+        assert kwargs["force_download"] is True
+        assert kwargs["resume_download"] is True
+        assert kwargs["proxies"] == {"https": "proxy"}
+        assert kwargs["token"] == "secret"
+        assert kwargs["local_files_only"] is True
+
+
+def test_checkpoint_resolution_rejects_local_directory_without_weights(tmp_path):
+    import lerobot.policies.pi05.modeling_pi05 as modeling_pi05
+
+    with pytest.raises(FileNotFoundError, match="model.safetensors"):
+        modeling_pi05._resolve_weight_files(
+            tmp_path,
+            force_download=False,
+            resume_download=None,
+            proxies=None,
+            token=None,
+            cache_dir=None,
+            local_files_only=False,
+            revision=None,
+        )
+
+
+@pytest.mark.parametrize("z_loss_weight", [0.0, 1e-4])
+@pytest.mark.parametrize("rows,valid_rows", [(24, 9), (48, 25)])
+def test_bucketed_ce_matches_dense_loss_and_gradients(z_loss_weight, rows, valid_rows):
+    generator = torch.Generator().manual_seed(23)
+    hidden_size, vocab_size = 7, 19
+    hidden_ref = torch.randn(rows, hidden_size, generator=generator, dtype=torch.float64, requires_grad=True)
+    weight_ref = torch.randn(
+        vocab_size, hidden_size, generator=generator, dtype=torch.float64, requires_grad=True
+    )
+    labels = torch.full((rows,), -100, dtype=torch.long)
+    valid_indices = torch.randperm(rows, generator=generator)[:valid_rows]
+    labels[valid_indices] = torch.randint(0, vocab_size, (valid_rows,), generator=generator)
+    hidden_bucketed = hidden_ref.detach().clone().requires_grad_(True)
+    weight_bucketed = weight_ref.detach().clone().requires_grad_(True)
+
+    import lerobot.policies.pi052.modeling_pi052 as modeling_pi052
+
+    loss_ref = _lin_ce_flat(hidden_ref, weight_ref, labels, z_loss_weight=z_loss_weight)
+    old_limit = modeling_pi052._LOGITS_CE_MAX_POSITIONS
+    modeling_pi052._LOGITS_CE_MAX_POSITIONS = 16
+    try:
+        loss_bucketed = _lin_ce_flat(
+            hidden_bucketed,
+            weight_bucketed,
+            labels,
+            z_loss_weight=z_loss_weight,
+        )
+    finally:
+        modeling_pi052._LOGITS_CE_MAX_POSITIONS = old_limit
+
+    loss_ref.backward()
+    loss_bucketed.backward()
+
+    torch.testing.assert_close(loss_bucketed, loss_ref, rtol=1e-6, atol=1e-6)
+    torch.testing.assert_close(hidden_bucketed.grad, hidden_ref.grad, rtol=1e-12, atol=1e-12)
+    torch.testing.assert_close(weight_bucketed.grad, weight_ref.grad, rtol=1e-12, atol=1e-12)
+
+
+def test_bucketed_ce_all_ignored_preserves_zero_gradients():
+    hidden = torch.randn(24, 7, dtype=torch.float64, requires_grad=True)
+    weight = torch.randn(19, 7, dtype=torch.float64, requires_grad=True)
+    labels = torch.full((24,), -100, dtype=torch.long)
+
+    import lerobot.policies.pi052.modeling_pi052 as modeling_pi052
+
+    old_limit = modeling_pi052._LOGITS_CE_MAX_POSITIONS
+    modeling_pi052._LOGITS_CE_MAX_POSITIONS = 16
+    try:
+        loss = _lin_ce_flat(hidden, weight, labels)
+    finally:
+        modeling_pi052._LOGITS_CE_MAX_POSITIONS = old_limit
+    loss.backward()
+
+    assert loss.item() == 0.0
+    assert hidden.grad is not None
+    assert weight.grad is not None
+    assert torch.count_nonzero(hidden.grad) == 0
+    assert torch.count_nonzero(weight.grad) == 0

tests/policies/pi052/test_pi052_checkpoint_portability.py

@@ -0,0 +1,152 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import shutil
+from dataclasses import asdict
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+
+from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature
+from lerobot.configs.recipe import MessageTurn, TrainingRecipe
+from lerobot.policies import make_pre_post_processors
+from lerobot.processor import ActionTokenizerProcessorStep, DataProcessorPipeline, NormalizerProcessorStep
+from lerobot.processor.converters import identity_transition
+from lerobot.processor.render_messages_processor import RenderMessagesStep
+from lerobot.utils.constants import ACTION
+
+
+class _ActionTokenizer:
+    def __call__(self, actions):
+        return np.asarray(actions).round().astype(np.int64)
+
+    def save_pretrained(self, path):
+        path.mkdir(parents=True)
+        (path / "processor_config.json").write_text('{"processor_class": "_ActionTokenizer"}\n')
+
+
+class _PaligemmaTokenizer:
+    vocab_size = 4096
+    bos_token_id = 2
+
+    def encode(self, text, **kwargs):
+        return [10, 11] if text == "Action: " else [12]
+
+
+def _make_pipeline(action_tokenizer_path):
+    recipe = TrainingRecipe(
+        messages=[
+            MessageTurn(role="user", content="${task}", stream="high_level"),
+            MessageTurn(role="assistant", content="${subtask}", stream="low_level", target=True),
+        ]
+    )
+    stats = {ACTION: {"min": torch.tensor([-1.0, -2.0]), "max": torch.tensor([1.0, 2.0])}}
+    normalizer = NormalizerProcessorStep(
+        features={ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(2,))},
+        norm_map={FeatureType.ACTION: NormalizationMode.MIN_MAX},
+        stats=stats,
+    )
+    action_tokenizer = ActionTokenizerProcessorStep(
+        action_tokenizer_name=str(action_tokenizer_path),
+        max_action_tokens=16,
+        fast_skip_tokens=128,
+    )
+    return DataProcessorPipeline(
+        [normalizer, RenderMessagesStep(recipe), action_tokenizer],
+        name="policy_preprocessor",
+        to_transition=identity_transition,
+        to_output=identity_transition,
+    )
+
+
+def test_pi052_pipeline_embeds_and_loads_fitted_action_tokenizer(tmp_path, monkeypatch):
+    original_cache = tmp_path / "original_fast_cache"
+    original_cache.mkdir()
+    tokenizer = _ActionTokenizer()
+    monkeypatch.setattr(
+        "lerobot.processor.tokenizer_processor.AutoProcessor.from_pretrained",
+        lambda path, **kwargs: tokenizer,
+    )
+    monkeypatch.setattr(
+        "lerobot.processor.tokenizer_processor.AutoTokenizer.from_pretrained",
+        lambda *args, **kwargs: _PaligemmaTokenizer(),
+    )
+    monkeypatch.setattr(
+        "lerobot.policies.pi052.fit_fast_tokenizer.fit_fast_tokenizer",
+        lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("FAST fitting must not run")),
+    )
+
+    pipeline = _make_pipeline(original_cache)
+    expected_tokens = pipeline.steps[-1]._tokenize_action(torch.tensor([[[0.2, 0.8]]]))[0]
+    expected_recipe = asdict(pipeline.steps[1].recipe)
+    expected_state = pipeline.steps[0].state_dict()
+    checkpoint = tmp_path / "checkpoint"
+    pipeline.save_pretrained(checkpoint)
+    DataProcessorPipeline(
+        [],
+        name="policy_postprocessor",
+        to_transition=identity_transition,
+        to_output=identity_transition,
+    ).save_pretrained(checkpoint)
+
+    saved_config = json.loads((checkpoint / "policy_preprocessor.json").read_text())
+    tokenizer_step = saved_config["steps"][2]
+    assert tokenizer_step["config"]["action_tokenizer_name"] == "action_tokenizer"
+    assert tokenizer_step["artifacts"] == {"action_tokenizer_name": "action_tokenizer"}
+    assert (checkpoint / "action_tokenizer" / "processor_config.json").is_file()
+
+    shutil.rmtree(original_cache)
+    loaded, _ = make_pre_post_processors(
+        SimpleNamespace(type="pi052", auto_fit_fast_tokenizer=True),
+        pretrained_path=str(checkpoint),
+        dataset_repo_id="org/dataset-that-must-not-be-read",
+    )
+
+    assert asdict(loaded.steps[1].recipe) == expected_recipe
+    for key, tensor in expected_state.items():
+        torch.testing.assert_close(loaded.steps[0].state_dict()[key], tensor)
+    torch.testing.assert_close(
+        loaded.steps[-1]._tokenize_action(torch.tensor([[[0.2, 0.8]]]))[0],
+        expected_tokens,
+    )
+
+
+def test_pi052_pipeline_rejects_missing_fitted_action_tokenizer(tmp_path, monkeypatch):
+    tokenizer = _ActionTokenizer()
+    monkeypatch.setattr(
+        "lerobot.processor.tokenizer_processor.AutoProcessor.from_pretrained",
+        lambda path, **kwargs: tokenizer,
+    )
+    monkeypatch.setattr(
+        "lerobot.processor.tokenizer_processor.AutoTokenizer.from_pretrained",
+        lambda *args, **kwargs: _PaligemmaTokenizer(),
+    )
+
+    pipeline = _make_pipeline(tmp_path / "original_fast_cache")
+    checkpoint = tmp_path / "checkpoint"
+    pipeline.save_pretrained(checkpoint)
+    shutil.rmtree(checkpoint / "action_tokenizer")
+
+    with pytest.raises(FileNotFoundError, match="Checkpoint artifacts are incomplete"):
+        DataProcessorPipeline.from_pretrained(
+            checkpoint,
+            config_filename="policy_preprocessor.json",
+            to_transition=identity_transition,
+            to_output=identity_transition,
+        )

tests/policies/pi052/test_pi052_fast_action_loss.py

@@ -0,0 +1,162 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Regression tests for PI052 FAST action-code supervision."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+from torch.nn import functional as F  # noqa: N812
+
+pytest.importorskip("transformers")
+pytest.importorskip("liger_kernel")
+
+from lerobot.policies.pi052.modeling_pi052 import PI052Policy, _fast_lin_ce  # noqa: E402
+from lerobot.policies.pi052.processor_pi052 import make_pi052_pre_post_processors  # noqa: E402
+
+
+def _fast_ce(logits, action_tokens, action_code_mask, predict_actions_t):
+    """Adapter: ``_fast_lin_ce`` is Liger-fused (hidden @ lm_head_weightᵀ).
+
+    Feeding an identity ``lm_head_weight`` makes the computed logits equal the
+    provided ``logits``, so these regression tests exercise the masking/gating
+    logic exactly as before the fused-CE refactor. Liger's Triton kernel is
+    GPU-only, so inputs are moved to CUDA and the loss is returned on CPU
+    (keeping grad flowing back to the CPU ``logits`` leaf).
+    """
+    if not torch.cuda.is_available():
+        pytest.skip("Liger fused CE requires CUDA")
+    vocab_size = logits.shape[-1]
+    eye = torch.eye(vocab_size, dtype=logits.dtype, device="cuda")
+    predict = predict_actions_t.cuda() if predict_actions_t is not None else None
+    loss = _fast_lin_ce(logits.cuda(), eye, action_tokens.cuda(), action_code_mask.cuda(), predict)
+    return loss.cpu()
+
+
+def test_fast_ce_supervises_only_discrete_action_codes():
+    """Wrapper tokens can be wrong without affecting the FAST action-code loss."""
+    vocab_size = 8
+    action_tokens = torch.tensor([[1, 2, 3, 4, 5, 0]])
+    action_code_mask = torch.tensor([[False, False, True, True, False, False]])
+
+    logits = torch.zeros(1, action_tokens.shape[1], vocab_size)
+    # Deliberately bad wrapper-token predictions. These should be ignored.
+    logits[0, 0, 7] = 10.0  # target would be token 2
+    logits[0, 3, 7] = 10.0  # target would be delimiter token 5
+    # Correct action-code predictions: hidden t predicts target t + 1.
+    logits[0, 1, 3] = 10.0
+    logits[0, 2, 4] = 10.0
+
+    loss = _fast_ce(logits, action_tokens, action_code_mask, predict_actions_t=None)
+    expected = F.cross_entropy(
+        torch.stack([logits[0, 1], logits[0, 2]]),
+        torch.tensor([3, 4]),
+        reduction="mean",
+    )
+
+    # Allow the fused GPU kernel's ~1e-7 difference on small losses.
+    assert torch.allclose(loss, expected, atol=1e-5, rtol=1e-3)
+
+
+def test_fast_ce_masks_non_action_samples():
+    """Recipe samples with predict_actions=False do not contribute FAST loss."""
+    vocab_size = 8
+    action_tokens = torch.tensor([[1, 2, 3, 4], [1, 2, 5, 6]])
+    action_code_mask = torch.tensor([[False, False, True, True], [False, False, True, True]])
+    predict_actions = torch.tensor([True, False])
+
+    logits = torch.zeros(2, action_tokens.shape[1], vocab_size)
+    logits[0, 1, 3] = 10.0
+    logits[0, 2, 4] = 10.0
+    # Bad predictions in the masked sample should not matter.
+    logits[1, 1, 7] = 10.0
+    logits[1, 2, 7] = 10.0
+
+    loss = _fast_ce(logits, action_tokens, action_code_mask, predict_actions)
+    expected = F.cross_entropy(
+        torch.stack([logits[0, 1], logits[0, 2]]),
+        torch.tensor([3, 4]),
+        reduction="mean",
+    )
+
+    # Allow the fused GPU kernel's ~1e-7 difference on small losses.
+    assert torch.allclose(loss, expected, atol=1e-5, rtol=1e-3)
+
+
+def test_fast_ce_returns_zero_when_no_action_code_positions_are_valid():
+    logits = torch.randn(2, 4, 8, requires_grad=True)
+    action_tokens = torch.tensor([[1, 2, 3, 4], [1, 2, 5, 6]])
+    action_code_mask = torch.zeros_like(action_tokens, dtype=torch.bool)
+
+    loss = _fast_ce(logits, action_tokens, action_code_mask, predict_actions_t=None)
+
+    assert loss.item() == 0
+    loss.backward()
+    assert logits.grad is not None
+
+
+def test_fast_ce_averages_each_action_sample_equally():
+    torch.manual_seed(0)
+    hidden = torch.randn(2, 5, 8)
+    lm_head_weight = torch.eye(8)
+    action_tokens = torch.tensor([[1, 2, 0, 0, 0], [1, 3, 4, 5, 6]])
+    action_code_mask = torch.tensor([[False, True, False, False, False], [False, True, True, True, True]])
+
+    loss = _fast_lin_ce(
+        hidden,
+        lm_head_weight,
+        action_tokens,
+        action_code_mask,
+        predict_actions_t=None,
+        reduction="mean",
+    )
+    per_sample = _fast_lin_ce(
+        hidden,
+        lm_head_weight,
+        action_tokens,
+        action_code_mask,
+        predict_actions_t=None,
+        reduction="none",
+    )
+
+    assert torch.allclose(loss, per_sample.mean())
+
+
+def test_pi052_rejects_fast_loss_without_recipe():
+    config = SimpleNamespace(recipe_path=None, enable_fast_action_loss=True)
+
+    with pytest.raises(ValueError, match="recipe_path"):
+        make_pi052_pre_post_processors(config)
+
+
+def test_pi052_rejects_missing_fast_batch_keys():
+    policy = PI052Policy.__new__(PI052Policy)
+    nn.Module.__init__(policy)
+    policy.config = SimpleNamespace(
+        enable_fast_action_loss=True,
+        fast_action_loss_weight=1.0,
+        flow_loss_weight=0.0,
+        text_loss_weight=1.0,
+    )
+    batch = {
+        "text_labels": torch.tensor([[1, 2]]),
+        "predict_actions": torch.tensor([True]),
+    }
+
+    with pytest.raises(ValueError, match="FAST action loss is enabled"):
+        policy.forward(batch)

tests/policies/pi052/test_pi052_fit_fast_tokenizer.py

@@ -0,0 +1,122 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import pytest
+
+from lerobot.policies.pi052.fit_fast_tokenizer import (
+    _apply_relative_actions,
+    _dataset_signature,
+    _is_global_leader,
+    _normalize_actions,
+    _select_episode_indices,
+    _validate_fast_reconstruction,
+)
+
+
+def test_fast_tokenizer_fit_uses_training_mean_std_normalization():
+    actions = np.array([[[1.0, 7.0], [3.0, 3.0]]], dtype=np.float32)
+    stats = {"mean": [2.0, 5.0], "std": [0.5, 2.0]}
+
+    normalized = _normalize_actions(actions, "MEAN_STD", stats)
+
+    np.testing.assert_allclose(normalized, [[[-2.0, 1.0], [2.0, -1.0]]])
+
+
+def test_fast_tokenizer_fit_quantiles_match_training_without_clipping():
+    actions = np.array([[[-1.0], [3.0]]], dtype=np.float32)
+    stats = {"q01": [0.0], "q99": [2.0]}
+
+    normalized = _normalize_actions(actions, "QUANTILES", stats)
+
+    np.testing.assert_allclose(normalized, [[[-2.0], [2.0]]])
+
+
+def test_fast_tokenizer_cache_signature_tracks_stats_and_episode_selection():
+    kwargs = {
+        "dataset_repo_id": "org/dataset",
+        "base_tokenizer_name": "physical-intelligence/fast",
+        "n_samples": 100,
+        "chunk_size": 20,
+        "normalization_mode": "QUANTILES",
+        "dataset_revision": "main",
+        "episodes": [1, 2, 3],
+        "exclude_episodes": [2],
+        "use_relative_actions": False,
+        "relative_action_mask": None,
+    }
+
+    first = _dataset_signature(**kwargs, action_stats={"q01": [0.0], "q99": [1.0]})
+    changed_stats = _dataset_signature(**kwargs, action_stats={"q01": [0.0], "q99": [2.0]})
+    changed_selection = _dataset_signature(
+        **{**kwargs, "exclude_episodes": [2, 3]},
+        action_stats={"q01": [0.0], "q99": [1.0]},
+    )
+
+    assert first != changed_stats
+    assert first != changed_selection
+
+
+def test_fast_tokenizer_uses_only_global_rank_zero(monkeypatch):
+    monkeypatch.setenv("RANK", "8")
+    monkeypatch.setenv("LOCAL_RANK", "0")
+    assert not _is_global_leader()
+
+    monkeypatch.setenv("RANK", "0")
+    assert _is_global_leader()
+
+
+def test_fast_tokenizer_episode_selection_applies_allowlist_and_exclusions():
+    selected = _select_episode_indices([0, 1, 2, 3], episodes=[1, 2, 3], exclude_episodes=[2])
+
+    assert selected == [1, 3]
+
+
+def test_fast_tokenizer_relative_actions_match_training_transform():
+    actions = np.array([[[2.0, 10.0], [3.0, 11.0]]], dtype=np.float32)
+    states = np.array([[1.0, 4.0]], dtype=np.float32)
+
+    relative = _apply_relative_actions(actions, states, [True, False])
+
+    np.testing.assert_allclose(relative, [[[1.0, 10.0], [2.0, 11.0]]])
+
+
+class _RoundTripTokenizer:
+    def __init__(self, offset: float = 0.0):
+        self.offset = offset
+
+    def __call__(self, actions):
+        return actions
+
+    def decode(self, tokens):
+        return tokens + self.offset
+
+
+def test_fast_tokenizer_reconstruction_validation_reports_error():
+    actions = np.arange(24, dtype=np.float32).reshape(2, 3, 4) / 24
+
+    report, decoded = _validate_fast_reconstruction(_RoundTripTokenizer(0.05), actions, 0.1, 0.1)
+
+    np.testing.assert_allclose(decoded, actions + 0.05)
+    assert report["reconstruction_rmse"] == pytest.approx(0.05)
+    assert report["max_dim_rmse"] == pytest.approx(0.05)
+
+
+def test_fast_tokenizer_reconstruction_validation_rejects_large_error():
+    actions = np.arange(24, dtype=np.float32).reshape(2, 3, 4) / 24
+
+    with pytest.raises(RuntimeError, match="exceeds the configured limit"):
+        _validate_fast_reconstruction(_RoundTripTokenizer(0.25), actions, 0.1, 0.2)

tests/policies/pi052/test_pi052_flex_attention.py

@@ -0,0 +1,65 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+
+import pytest
+import torch
+
+pytest.importorskip("transformers")
+
+import lerobot.policies.pi052.modeling_pi052 as modeling_pi052  # noqa: E402
+from lerobot.policies.pi052.configuration_pi052 import PI052Config  # noqa: E402
+
+
+def test_flex_backend_skips_non_cuda_without_initializing(monkeypatch):
+    monkeypatch.setattr(modeling_pi052, "_flex_fns", None)
+    monkeypatch.setattr(torch, "compile", lambda *args, **kwargs: pytest.fail("torch.compile was called"))
+    monkeypatch.setattr(
+        torch.cuda,
+        "get_device_properties",
+        lambda *args, **kwargs: pytest.fail("CUDA properties were queried"),
+    )
+
+    assert modeling_pi052._get_flex_fns(torch.device("cpu")) is None
+    assert modeling_pi052._get_flex_kernel_options(torch.device("cpu")) is None
+    assert modeling_pi052._flex_fns is None
+
+
+def test_flex_initialization_failure_falls_back(monkeypatch, caplog):
+    monkeypatch.setattr(modeling_pi052, "_flex_fns", None)
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+
+    def fail_compile(*args, **kwargs):
+        raise RuntimeError("compile failed")
+
+    monkeypatch.setattr(torch, "compile", fail_compile)
+
+    with caplog.at_level(logging.WARNING, logger=modeling_pi052.__name__):
+        assert modeling_pi052._get_flex_fns(torch.device("cuda", 0)) is None
+
+    assert modeling_pi052._flex_fns is False
+    assert "FlexAttention unavailable" in caplog.text
+
+
+def test_flex_rejects_single_repeat_configuration():
+    with pytest.raises(ValueError, match="use_flex_attention requires flow_num_repeats > 1"):
+        PI052Config(use_flex_attention=True, flow_num_repeats=1)
+
+
+def test_flex_accepts_amortized_repeat_configuration():
+    config = PI052Config(use_flex_attention=True, flow_num_repeats=5)
+    assert config.use_flex_attention

tests/policies/pi052/test_pi052_import.py

@@ -0,0 +1,27 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import subprocess
+import sys
+
+
+def test_pi052_config_import_does_not_load_model_or_dataset_processor():
+    code = """
+import sys
+from lerobot.policies import PI052Config
+assert PI052Config.__name__ == "PI052Config"
+assert "lerobot.policies.pi052.modeling_pi052" not in sys.modules
+assert "lerobot.policies.pi052.processor_pi052" not in sys.modules
+"""
+    subprocess.run([sys.executable, "-c", code], check=True)

tests/policies/pi052/test_pi052_joint_sequence.py

@@ -0,0 +1,148 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for PI052 joint-sequence (paper-style) subtask conditioning.
+
+Joint recipes train the subtask text and the action losses in one sequence,
+with the supervised subtask span attended causally. At inference the same
+layout is rebuilt around the *generated* subtask, so these tests pin:
+
+- the inference-side encoder produces the same token ids and target positions
+  as the training-time tokenizer step for the same messages;
+- OR-ing causal marks into a prefix reproduces the training-time attention
+  pattern (prompt cannot see the subtask; subtask is causal over itself);
+- the joint recipe file stays a valid message recipe;
+- the FAST id mapping with the default ``fast_skip_tokens`` stays clear of
+  PaliGemma's ``<loc>`` range so VQA and FAST supervision never collide.
+"""
+
+from pathlib import Path
+
+import torch
+
+from lerobot.configs.recipe import TrainingRecipe
+from lerobot.policies.pi052.text_processor_pi052 import (
+    PI052TextTokenizerStep,
+    encode_prompt_with_targets,
+)
+
+
+class _CharTokenizer:
+    """Char-level stub: 1 char = 1 token, so offsets are trivially aligned."""
+
+    pad_token_id = 0
+    eos_token = "\x1f"  # unit separator — a 1-char "EOS" for testing
+
+    def __call__(self, text, max_length=None, padding=None, return_tensors=None, **kwargs):
+        limit = max_length if max_length is not None else len(text)
+        ids = [ord(c) % 251 + 1 for c in text[:limit]]
+        offsets = [(i, i + 1) for i in range(len(ids))]
+        attention = [1] * len(ids)
+        if padding == "max_length" and max_length is not None and len(ids) < max_length:
+            pad = max_length - len(ids)
+            ids += [self.pad_token_id] * pad
+            offsets += [(0, 0)] * pad
+            attention += [0] * pad
+        return {
+            "input_ids": torch.tensor([ids], dtype=torch.long),
+            "attention_mask": torch.tensor([attention], dtype=torch.long),
+            "offset_mapping": torch.tensor([offsets], dtype=torch.long),
+        }
+
+
+_MESSAGES = [
+    {"role": "user", "content": "fold the towel"},
+    {"role": "assistant", "content": "grab the near corner"},
+]
+
+
+def test_encode_prompt_with_targets_matches_training_labels():
+    tokenizer = _CharTokenizer()
+
+    step = PI052TextTokenizerStep(max_length=120)
+    step._tokenizer = tokenizer
+    train_ids, train_attn, labels, predict_actions, _prompt = step._encode_messages(
+        tokenizer,
+        [dict(m) for m in _MESSAGES],
+        message_streams=["low_level", "low_level"],
+        target_indices=[1],
+        complementary={},
+    )
+    assert bool(predict_actions)
+
+    ids, attn, marks = encode_prompt_with_targets(tokenizer, [dict(m) for m in _MESSAGES], [1])
+
+    n = int(attn.sum())
+    assert n == int(train_attn.sum())
+    assert torch.equal(ids[0, :n], train_ids[:n])
+    # Causal marks at inference must cover exactly the supervised label span.
+    assert torch.equal(marks[0, :n], labels[:n] != -100)
+    assert marks.any(), "the assistant target span must be marked"
+    # The user turn must stay unmarked (bidirectional prompt).
+    user_len = len("User: fold the towel\n")
+    assert not marks[0, :user_len].any()
+
+
+def test_apply_causal_language_marks_reproduces_training_mask():
+    from lerobot.policies.pi05.modeling_pi05 import make_att_2d_masks
+    from lerobot.policies.pi052.modeling_pi052 import _apply_causal_language_marks
+
+    n_img, n_lang = 4, 8
+    prefix_len = n_img + n_lang
+    pad = torch.ones((1, prefix_len), dtype=torch.bool)
+    att = torch.zeros((1, prefix_len), dtype=torch.bool)
+    # Subtask span = language positions 5..7 (prefix positions 9..11).
+    marks = torch.zeros((1, n_lang), dtype=torch.bool)
+    marks[0, 5:8] = True
+
+    att_marked = _apply_causal_language_marks(att, marks)
+    att_2d = make_att_2d_masks(pad, att_marked)[0]
+
+    subtask = [n_img + 5, n_img + 6, n_img + 7]
+    # Prompt and images never see the subtask.
+    for q in range(n_img + 5):
+        for k in subtask:
+            assert not att_2d[q, k], f"prompt position {q} must not attend subtask position {k}"
+    # Subtask tokens see the full prompt and earlier subtask tokens only.
+    for qi, q in enumerate(subtask):
+        for k in range(n_img + 5):
+            assert att_2d[q, k]
+        for ki, k in enumerate(subtask):
+            assert bool(att_2d[q, k]) == (ki <= qi)
+
+
+def test_joint_recipe_is_a_valid_message_recipe():
+    recipe_path = Path(__file__).parents[3] / "src" / "lerobot" / "configs" / "recipes" / "subtask_joint.yaml"
+    recipe = TrainingRecipe.from_yaml(recipe_path)
+    assert recipe.messages is not None and len(recipe.messages) == 2
+    assert all(turn.stream == "low_level" for turn in recipe.messages)
+    assert not recipe.messages[0].target
+    assert recipe.messages[1].target
+    assert recipe.messages[1].if_present == "subtask"
+
+
+def test_default_fast_mapping_clears_loc_and_seg_ranges():
+    from lerobot.policies.pi052.configuration_pi052 import PI052Config
+    from lerobot.policies.pi052.modeling_pi052 import _FAST_ACTION_VOCAB_SIZE
+
+    skip = PI052Config.__dataclass_fields__["fast_skip_tokens"].default
+    assert skip == 1152
+
+    paligemma_vocab = 257152
+    fast_ids = paligemma_vocab - 1 - skip - torch.arange(_FAST_ACTION_VOCAB_SIZE)
+    # Below the <loc> range [256000, 257024) and the <seg> range [257024, 257152).
+    assert int(fast_ids.max()) < 256000
+    assert int(fast_ids.min()) >= 0

tests/policies/pi052/test_pi052_runtime_adapter.py

@@ -0,0 +1,85 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from types import SimpleNamespace
+
+from lerobot.policies.pi052.inference.pi052_adapter import PI052PolicyAdapter
+from lerobot.runtime import RuntimeState
+from lerobot.runtime.adapter import split_plan_and_say
+
+
+def test_pi052_adapter_builds_recipe_prompts_from_runtime_state():
+    adapter = PI052PolicyAdapter(policy=object())
+    state = RuntimeState(
+        task="clean the kitchen",
+        language_context={"memory": "cup moved", "plan": "pick then place"},
+        extra={"prior_subtask": "pick the cup"},
+    )
+
+    assert adapter.build_messages("subtask", state) == [{"role": "user", "content": "clean the kitchen"}]
+    assert adapter.build_messages("memory", state) == [
+        {"role": "user", "content": "clean the kitchen"},
+        {"role": "assistant", "content": "Previous memory: cup moved"},
+        {"role": "user", "content": "Completed subtask: pick the cup"},
+    ]
+    assert adapter.build_messages("interjection", state, user_text="wait") == [
+        {"role": "user", "content": "clean the kitchen"},
+        {"role": "assistant", "content": "Previous plan:\npick then place"},
+        {"role": "user", "content": "wait"},
+    ]
+
+
+def test_pi052_adapter_strips_say_markers_from_plan_text():
+    adapter = PI052PolicyAdapter(policy=object())
+    text = "Move to the sink. <say>heading to the sink</say>"
+
+    assert split_plan_and_say(text) == ("Move to the sink.", "heading to the sink")
+    assert adapter.plan_from_text(text) == "Move to the sink."
+
+
+def test_rollout_language_cli_smoke_does_not_load_model(monkeypatch):
+    """lerobot-rollout dispatches language flags to the adapter-based runtime."""
+    from lerobot.runtime import cli
+    from lerobot.scripts import lerobot_rollout
+
+    fake_policy = SimpleNamespace(config=SimpleNamespace(device="cpu", type="pi052"))
+
+    monkeypatch.setattr(
+        cli,
+        "_load_policy_and_preprocessor",
+        lambda policy_path, **kwargs: (fake_policy, None, None),
+    )
+    monkeypatch.setattr(cli, "_run_repl", lambda runtime, **kwargs: 0)
+
+    assert lerobot_rollout.main(["--policy.path=fake", "--no_robot", "--task=clean", "--max_ticks=0"]) == 0
+
+
+def test_rollout_language_dispatch_preserves_standard_molmoact2_path(monkeypatch):
+    """MolmoAct2 only opts into open prompting when a language flag is present."""
+    from lerobot.scripts import lerobot_rollout
+
+    standard = [
+        "--policy.path=lerobot/MolmoAct2-SO100_101-LeRobot",
+        "--robot.type=so101_follower",
+        "--task=pick up the cube",
+    ]
+    assert not lerobot_rollout._uses_language_runtime(standard)
+    assert lerobot_rollout._uses_language_runtime([*standard, "--direct_subtask"])
+    assert lerobot_rollout._uses_language_runtime(["--policy.path=lerobot/pi052_robocasa", "--sim"])
+
+    standard_calls = []
+    monkeypatch.setattr(lerobot_rollout, "register_third_party_plugins", lambda: None)
+    monkeypatch.setattr(lerobot_rollout, "rollout", lambda: standard_calls.append(True))
+    lerobot_rollout.main(standard)
+    assert standard_calls == [True]

tests/policies/pi052/test_pi052_sdpa_attention.py

@@ -0,0 +1,147 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Numerical-parity tests for the SDPA attention port.
+
+``pi05`` / ``pi052`` replaced the per-layer call from
+``modeling_gemma.eager_attention_forward`` with
+``sdpa_attention_forward`` (PyTorch SDPA + GQA repeat). The forward
+output must be bit-equivalent (within bf16 tolerance) on the masks
+this model actually uses — block-bidirectional with an arbitrary
+additive bias — otherwise we silently change training behaviour.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytest.importorskip("transformers")
+
+from transformers.models.gemma import modeling_gemma  # noqa: E402
+
+from lerobot.policies.pi052.modeling_pi052 import make_att_2d_masks  # noqa: E402
+from lerobot.policies.pi_gemma import sdpa_attention_forward  # noqa: E402
+from lerobot.utils.constants import OPENPI_ATTENTION_MASK_VALUE  # noqa: E402
+
+
+def _mock_self_attn(num_kv_groups: int, training: bool = False):
+    """Bare module surface that both forwards read."""
+    return SimpleNamespace(
+        num_key_value_groups=num_kv_groups,
+        training=training,
+    )
+
+
+def _build_inputs(
+    bsize: int,
+    num_heads: int,
+    num_kv_heads: int,
+    seq_len: int,
+    head_dim: int,
+    dtype: torch.dtype,
+    seed: int = 0,
+):
+    g = torch.Generator(device="cpu").manual_seed(seed)
+    q = torch.randn(bsize, num_heads, seq_len, head_dim, dtype=dtype, generator=g)
+    k = torch.randn(bsize, num_kv_heads, seq_len, head_dim, dtype=dtype, generator=g)
+    v = torch.randn(bsize, num_kv_heads, seq_len, head_dim, dtype=dtype, generator=g)
+    return q, k, v
+
+
+def _block_bidirectional_mask(
+    bsize: int, seq_len: int, block_sizes: list[int], dtype: torch.dtype
+) -> torch.Tensor:
+    """Mimic ``_prepare_attention_masks_4d`` on a block layout that
+    matches ``[images, language, suffix]`` from ``embed_prefix`` +
+    ``embed_suffix``: every block bidirectional internally, later
+    blocks visible to earlier ones via the cumulative-block rule.
+    """
+    assert sum(block_sizes) == seq_len
+    att_marks = []
+    for i, n in enumerate(block_sizes):
+        att_marks += [1 if i > 0 else 0] + [0] * (n - 1)
+    pad = torch.ones(bsize, seq_len, dtype=torch.bool)
+    att = torch.tensor(att_marks, dtype=torch.bool)[None].expand(bsize, seq_len)
+    att_2d = make_att_2d_masks(pad, att)
+    bias = torch.where(
+        att_2d[:, None, :, :],
+        torch.zeros((), dtype=dtype),
+        torch.tensor(OPENPI_ATTENTION_MASK_VALUE, dtype=dtype),
+    )
+    return bias
+
+
+@pytest.mark.parametrize(
+    "num_heads,num_kv_heads,head_dim",
+    [
+        (8, 1, 256),  # gemma_2b / paligemma config
+        (8, 8, 64),  # MHA control (no GQA repeat)
+    ],
+)
+def test_sdpa_parity_with_eager_block_bidirectional(num_heads, num_kv_heads, head_dim):
+    """SDPA forward output matches the eager softmax(QK^T)@V on the
+    block-bidirectional mask layout pi05 actually uses."""
+    bsize, seq_len = 2, 13
+    block_sizes = [4, 5, 4]  # images, language, suffix-style blocks
+    dtype = torch.float32  # cpu math kernel — keep fp32 for tight tol
+    scaling = head_dim**-0.5
+
+    q, k, v = _build_inputs(bsize, num_heads, num_kv_heads, seq_len, head_dim, dtype)
+    mask = _block_bidirectional_mask(bsize, seq_len, block_sizes, dtype)
+
+    module = _mock_self_attn(num_heads // num_kv_heads)
+
+    out_eager, _ = modeling_gemma.eager_attention_forward(module, q, k, v, mask, scaling)
+    out_sdpa, _ = sdpa_attention_forward(module, q, k, v, mask, scaling)
+    assert out_eager.shape == out_sdpa.shape
+    torch.testing.assert_close(out_sdpa, out_eager, atol=1e-5, rtol=1e-4)
+
+
+def test_sdpa_parity_bf16():
+    """bf16 path — looser tolerance, must still match eager."""
+    bsize, num_heads, num_kv_heads, seq_len, head_dim = 2, 8, 1, 17, 256
+    scaling = head_dim**-0.5
+    q, k, v = _build_inputs(bsize, num_heads, num_kv_heads, seq_len, head_dim, torch.bfloat16)
+    mask = _block_bidirectional_mask(bsize, seq_len, [5, 6, 6], torch.bfloat16)
+    module = _mock_self_attn(num_heads // num_kv_heads)
+
+    out_eager, _ = modeling_gemma.eager_attention_forward(module, q, k, v, mask, scaling)
+    out_sdpa, _ = sdpa_attention_forward(module, q, k, v, mask, scaling)
+    torch.testing.assert_close(out_sdpa, out_eager, atol=2e-2, rtol=2e-2)
+
+
+def test_sdpa_parity_backward():
+    """Gradients flow through SDPA and match the eager path within
+    bf16 tolerance — critical for any training-side parity claim."""
+    bsize, num_heads, num_kv_heads, seq_len, head_dim = 1, 4, 2, 9, 32
+    scaling = head_dim**-0.5
+    q, k, v = _build_inputs(bsize, num_heads, num_kv_heads, seq_len, head_dim, torch.float32)
+    q.requires_grad_(True)
+    k.requires_grad_(True)
+    v.requires_grad_(True)
+    mask = _block_bidirectional_mask(bsize, seq_len, [3, 3, 3], torch.float32)
+    module = _mock_self_attn(num_heads // num_kv_heads)
+
+    out_e, _ = modeling_gemma.eager_attention_forward(module, q, k, v, mask, scaling)
+    g_q_e, g_k_e, g_v_e = torch.autograd.grad(out_e.sum(), [q, k, v])
+
+    out_s, _ = sdpa_attention_forward(module, q, k, v, mask, scaling)
+    g_q_s, g_k_s, g_v_s = torch.autograd.grad(out_s.sum(), [q, k, v])
+
+    torch.testing.assert_close(g_q_s, g_q_e, atol=1e-5, rtol=1e-4)
+    torch.testing.assert_close(g_k_s, g_k_e, atol=1e-5, rtol=1e-4)
+    torch.testing.assert_close(g_v_s, g_v_e, atol=1e-5, rtol=1e-4)

tests/policies/pi052/test_pi052_text_processor.py

@@ -0,0 +1,223 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for PI052's text tokenizer.
+
+Covers ``say`` tool-call flattening (PaliGemma's flat prompt has no
+structured tool calls, so a ``say`` call must be serialized into a
+``<say>...</say>`` text marker) and EOS-termination supervision (the
+supervised target span must end with an EOS token so the LM head learns
+to stop instead of rambling to ``max_length`` at inference).
+"""
+
+import torch
+
+from lerobot.configs.recipe import MessageTurn, TrainingRecipe
+from lerobot.policies.pi052.text_processor_pi052 import (
+    PI052TextTokenizerStep,
+    _flatten_say_tool_calls,
+    _format_messages,
+)
+from lerobot.processor import PolicyProcessorPipeline
+from lerobot.processor.render_messages_processor import RenderMessagesStep
+from lerobot.types import TransitionKey
+from lerobot.utils.constants import (
+    OBS_LANGUAGE_ATTENTION_MASK,
+    OBS_LANGUAGE_TOKENS,
+    POLICY_PREPROCESSOR_DEFAULT_NAME,
+)
+
+
+def _say_call(text):
+    return {"type": "function", "function": {"name": "say", "arguments": {"text": text}}}
+
+
+def test_flatten_appends_say_marker_and_drops_tool_calls():
+    msg = {"role": "assistant", "content": "Heading to the cube.", "tool_calls": [_say_call("On it!")]}
+    out = _flatten_say_tool_calls(msg)
+    assert "tool_calls" not in out
+    assert out["content"] == "Heading to the cube.\n<say>On it!</say>"
+
+
+def test_flatten_marker_only_when_content_empty_or_none():
+    out = _flatten_say_tool_calls({"role": "assistant", "tool_calls": [_say_call("hi")]})
+    assert out["content"] == "<say>hi</say>"
+
+
+def test_flatten_accepts_json_string_arguments():
+    call = {"type": "function", "function": {"name": "say", "arguments": '{"text": "hello there"}'}}
+    out = _flatten_say_tool_calls({"role": "assistant", "content": "p", "tool_calls": [call]})
+    assert out["content"] == "p\n<say>hello there</say>"
+
+
+def test_flatten_leaves_messages_without_tool_calls_untouched():
+    msg = {"role": "assistant", "content": "just a plan"}
+    assert _flatten_say_tool_calls(msg) == msg
+
+
+def test_flatten_drops_non_say_tool_calls_but_keeps_content():
+    weather = {"type": "function", "function": {"name": "check_weather", "arguments": {}}}
+    out = _flatten_say_tool_calls({"role": "assistant", "content": "plan only", "tool_calls": [weather]})
+    assert out["content"] == "plan only"
+    assert "tool_calls" not in out
+
+
+def test_format_messages_appends_eos_to_target_turns_only():
+    msgs = [
+        {"role": "user", "content": "pick cube"},
+        {"role": "assistant", "content": "move to cube"},
+    ]
+    prompt, spans = _format_messages(msgs, target_indices=[1], eos_token="<eos>")
+    # EOS is appended to the supervised target (assistant) turn only.
+    assert prompt == "User: pick cube\nAssistant: move to cube<eos>\n"
+    # The user span is unchanged; the target span covers content + EOS.
+    assert prompt[spans[0][0] : spans[0][1]] == "pick cube"
+    assert prompt[spans[1][0] : spans[1][1]] == "move to cube<eos>"
+
+
+def test_format_messages_without_eos_args_is_unchanged():
+    """Inference callers omit target_indices / eos_token — no EOS baked in."""
+    prompt, spans = _format_messages([{"role": "user", "content": "hi"}])
+    assert prompt == "User: hi\n"
+    assert prompt[spans[0][0] : spans[0][1]] == "hi"
+
+
+def test_pi052_steps_roundtrip_through_standard_pipeline_loader(tmp_path):
+    recipe = TrainingRecipe(messages=[MessageTurn(role="user", content="${task}", stream="low_level")])
+    pipeline = PolicyProcessorPipeline(
+        steps=[
+            RenderMessagesStep(recipe),
+            PI052TextTokenizerStep(
+                tokenizer_name="custom-tokenizer",
+                max_length=77,
+                plan_dropout_prob=0.2,
+                dropout_seed=3,
+            ),
+        ],
+        name=POLICY_PREPROCESSOR_DEFAULT_NAME,
+    )
+    pipeline.save_pretrained(tmp_path)
+
+    loaded = PolicyProcessorPipeline.from_pretrained(
+        tmp_path, config_filename=f"{POLICY_PREPROCESSOR_DEFAULT_NAME}.json"
+    )
+
+    assert loaded.steps[0].recipe == recipe
+    assert loaded.steps[1].tokenizer_name == "custom-tokenizer"
+    assert loaded.steps[1].max_length == 77
+    assert loaded.steps[1].plan_dropout_prob == 0.2
+    assert loaded.steps[1].dropout_seed == 3
+
+
+def _eos_char_id() -> int:
+    """Token id _CharTokenizer assigns to its 1-char EOS."""
+    return ord("\x1f") % 251 + 1
+
+
+def test_pi052_text_tokenizer_supervises_eos_at_target_end():
+    """The appended EOS is the last supervised label on a target turn —
+    that's the signal that teaches the LM head to stop. The trailing
+    newline right after it stays unsupervised (-100)."""
+    step = PI052TextTokenizerStep(max_length=64)
+    step._tokenizer = _CharTokenizer()
+    transition = {
+        TransitionKey.OBSERVATION: {},
+        TransitionKey.COMPLEMENTARY_DATA: {
+            "messages": [
+                {"role": "user", "content": "pick cube"},
+                {"role": "assistant", "content": "move to cube"},
+            ],
+            "target_message_indices": [1],
+            "message_streams": ["high_level", "high_level"],
+            "index": torch.tensor(10),
+        },
+    }
+    out = step(transition)
+    ids = out[TransitionKey.OBSERVATION][OBS_LANGUAGE_TOKENS][0]
+    labels = out[TransitionKey.COMPLEMENTARY_DATA]["text_labels"][0]
+
+    supervised = (labels != -100).nonzero().flatten().tolist()
+    assert supervised, "target turn produced no supervised labels"
+    last = supervised[-1]
+    # The last supervised token is the appended EOS.
+    assert int(ids[last]) == _eos_char_id()
+    assert int(labels[last]) == _eos_char_id()
+    # The token right after the EOS (the trailing newline) is NOT supervised.
+    assert int(labels[last + 1]) == -100
+
+
+class _CharTokenizer:
+    pad_token_id = 0
+    eos_token = "\x1f"  # unit separator — a 1-char "EOS" for testing
+
+    def __call__(
+        self,
+        text,
+        max_length,
+        padding,
+        truncation,
+        return_tensors,
+        return_offsets_mapping,
+        padding_side,
+    ):
+        ids = [ord(c) % 251 + 1 for c in text[:max_length]]
+        offsets = [(i, i + 1) for i in range(len(ids))]
+        attention = [1] * len(ids)
+        if padding == "max_length" and len(ids) < max_length:
+            pad = max_length - len(ids)
+            ids += [self.pad_token_id] * pad
+            offsets += [(0, 0)] * pad
+            attention += [0] * pad
+        return {
+            "input_ids": torch.tensor([ids], dtype=torch.long),
+            "attention_mask": torch.tensor([attention], dtype=torch.long),
+            "offset_mapping": torch.tensor([offsets], dtype=torch.long),
+        }
+
+    def decode(self, token_ids, skip_special_tokens=False):
+        return "".join(chr(max(int(i) - 1, 0)) for i in token_ids if int(i) != self.pad_token_id)
+
+
+def test_pi052_text_tokenizer_handles_batched_rendered_messages():
+    step = PI052TextTokenizerStep(max_length=64)
+    step._tokenizer = _CharTokenizer()
+
+    transition = {
+        TransitionKey.OBSERVATION: {},
+        TransitionKey.COMPLEMENTARY_DATA: {
+            "messages": [
+                [
+                    {"role": "user", "content": "pick cube"},
+                    {"role": "assistant", "content": "move to cube"},
+                ],
+                [{"role": "user", "content": "open drawer"}],
+            ],
+            "target_message_indices": [[1], []],
+            "message_streams": [["high_level", "high_level"], ["low_level"]],
+            "index": torch.tensor([10, 11]),
+        },
+    }
+
+    out = step(transition)
+    obs = out[TransitionKey.OBSERVATION]
+    comp = out[TransitionKey.COMPLEMENTARY_DATA]
+
+    assert obs[OBS_LANGUAGE_TOKENS].shape == (2, 64)
+    assert obs[OBS_LANGUAGE_ATTENTION_MASK].shape == (2, 64)
+    assert comp["text_labels"].shape == (2, 64)
+    assert comp["predict_actions"].tolist() == [False, True]
+    assert (comp["text_labels"][0] != -100).any()
+    assert not (comp["text_labels"][1] != -100).any()

tests/policies/pi052/test_pi052_vision_checkpointing.py

@@ -0,0 +1,141 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from types import MethodType, SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+pytest.importorskip("transformers")
+
+from lerobot.policies.pi052.modeling_pi052 import PI05Pytorch
+
+
+class _MockVisionTower:
+    def __init__(self):
+        self.enable_kwargs = None
+        self.disable_calls = 0
+
+    def gradient_checkpointing_enable(self, **kwargs):
+        self.enable_kwargs = kwargs
+
+    def gradient_checkpointing_disable(self):
+        self.disable_calls += 1
+
+
+def _checkpoint_model():
+    tower = _MockVisionTower()
+    language_model = SimpleNamespace(gradient_checkpointing=False)
+    expert_model = SimpleNamespace(gradient_checkpointing=False)
+    model = PI05Pytorch.__new__(PI05Pytorch)
+    nn.Module.__init__(model)
+    model.gradient_checkpointing_enabled = False
+    model.paligemma_with_expert = SimpleNamespace(
+        paligemma=SimpleNamespace(model=SimpleNamespace(language_model=language_model, vision_tower=tower)),
+        gemma_expert=SimpleNamespace(model=expert_model),
+    )
+    return model, tower, language_model, expert_model
+
+
+def test_gradient_checkpointing_uses_vision_tower_layer_api():
+    model, tower, language_model, expert_model = _checkpoint_model()
+
+    PI05Pytorch.gradient_checkpointing_enable(model)
+
+    assert model.gradient_checkpointing_enabled
+    assert language_model.gradient_checkpointing
+    assert expert_model.gradient_checkpointing
+    assert tower.enable_kwargs == {"gradient_checkpointing_kwargs": {"use_reentrant": False}}
+
+    PI05Pytorch.gradient_checkpointing_disable(model)
+
+    assert not model.gradient_checkpointing_enabled
+    assert not language_model.gradient_checkpointing
+    assert not expert_model.gradient_checkpointing
+    assert tower.disable_calls == 1
+
+
+def test_siglip_layers_recompute_individually():
+    from transformers.models.siglip.configuration_siglip import SiglipVisionConfig
+    from transformers.models.siglip.modeling_siglip import SiglipVisionModel
+
+    config = SiglipVisionConfig(
+        hidden_size=16,
+        intermediate_size=32,
+        num_hidden_layers=2,
+        num_attention_heads=2,
+        num_channels=3,
+        image_size=16,
+        patch_size=8,
+    )
+    tower = SiglipVisionModel(config).train()
+    tower.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
+    calls = [0] * config.num_hidden_layers
+
+    for index, layer in enumerate(tower.vision_model.encoder.layers):
+        original_forward = layer.forward
+
+        def counted_forward(self, *args, _index=index, _forward=original_forward, **kwargs):
+            calls[_index] += 1
+            return _forward(*args, **kwargs)
+
+        layer.forward = MethodType(counted_forward, layer)
+
+    pixels = torch.randn(2, config.num_channels, config.image_size, config.image_size)
+    tower(pixels).last_hidden_state.sum().backward()
+
+    assert calls == [2] * config.num_hidden_layers
+
+
+def test_embed_prefix_does_not_wrap_the_whole_vision_tower_checkpoint():
+    model = PI05Pytorch.__new__(PI05Pytorch)
+    nn.Module.__init__(model)
+    model.config = SimpleNamespace()
+    model.gradient_checkpointing_enabled = True
+    model.train()
+
+    image_calls = []
+
+    def embed_image(image):
+        image_calls.append(image.shape)
+        return image[:, :1, 0, :2]
+
+    def embed_language_tokens(tokens):
+        return tokens.to(torch.float32).unsqueeze(-1).expand(*tokens.shape, 2)
+
+    model.paligemma_with_expert = SimpleNamespace(
+        embed_image=embed_image,
+        embed_language_tokens=embed_language_tokens,
+    )
+    outer_checkpoint_calls = []
+
+    def apply_checkpoint(func, value):
+        outer_checkpoint_calls.append(value.shape)
+        return func(value)
+
+    model._apply_checkpoint = apply_checkpoint
+
+    images = [torch.randn(2, 3, 4, 4), torch.randn(2, 3, 4, 4)]
+    image_masks = [torch.ones(2, dtype=torch.bool) for _ in images]
+    tokens = torch.ones(2, 3, dtype=torch.long)
+    token_masks = torch.ones_like(tokens, dtype=torch.bool)
+
+    embeddings, _, _ = model.embed_prefix(images, image_masks, tokens, token_masks)
+
+    assert image_calls == [image.shape for image in images]
+    assert outer_checkpoint_calls == [tokens.shape]
+    assert embeddings.shape == (2, 5, 2)

tests/policies/pi0_fast/test_pi0_fast_tokenizer_fit.py

@@ -0,0 +1,106 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from types import SimpleNamespace
+
+import pytest
+
+from lerobot.policies import factory
+from lerobot.policies.pi0_fast.configuration_pi0_fast import PI0FastConfig
+from lerobot.policies.pi052 import fit_fast_tokenizer as fit_module
+
+
+def test_pi0_fast_resolves_dataset_specific_tokenizer(monkeypatch, tmp_path):
+    config = PI0FastConfig(
+        auto_fit_fast_tokenizer=True,
+        action_tokenizer_name="base-tokenizer",
+        fast_tokenizer_cache_dir=str(tmp_path),
+        fast_tokenizer_fit_samples=17,
+        chunk_size=12,
+        n_action_steps=12,
+    )
+    received = {}
+
+    def fake_fit(**kwargs):
+        received.update(kwargs)
+        return "/cache/fitted-tokenizer"
+
+    monkeypatch.setattr(fit_module, "fit_fast_tokenizer", fake_fit)
+
+    assert fit_module.resolve_fast_tokenizer(config, "user/dataset") == "/cache/fitted-tokenizer"
+    assert received == {
+        "dataset_repo_id": "user/dataset",
+        "cache_dir": tmp_path,
+        "base_tokenizer_name": "base-tokenizer",
+        "n_samples": 17,
+        "chunk_size": 12,
+        "dataset_root": None,
+        "dataset_revision": None,
+        "episodes": None,
+        "exclude_episodes": None,
+        "normalization_mode": config.normalization_mapping["ACTION"],
+        "action_stats": None,
+        "use_relative_actions": False,
+        "relative_action_mask": None,
+    }
+
+
+def test_fast_fit_failure_is_not_silently_replaced(monkeypatch, tmp_path):
+    config = PI0FastConfig(auto_fit_fast_tokenizer=True, fast_tokenizer_cache_dir=str(tmp_path))
+    monkeypatch.setattr(
+        fit_module,
+        "fit_fast_tokenizer",
+        lambda **kwargs: (_ for _ in ()).throw(RuntimeError("fit failed")),
+    )
+
+    with pytest.raises(RuntimeError, match="fit failed"):
+        fit_module.resolve_fast_tokenizer(config, "user/dataset")
+
+
+def test_only_global_rank_zero_fits_shared_tokenizer(monkeypatch):
+    monkeypatch.setenv("RANK", "8")
+    monkeypatch.setenv("LOCAL_RANK", "0")
+    assert not fit_module._is_global_leader()
+
+    monkeypatch.setenv("RANK", "0")
+    assert fit_module._is_global_leader()
+
+
+def test_pretrained_pi0_fast_overrides_only_fitted_tokenizer(monkeypatch):
+    config = PI0FastConfig(auto_fit_fast_tokenizer=True)
+    calls = []
+
+    monkeypatch.setattr(
+        fit_module,
+        "resolve_fast_tokenizer",
+        lambda config, dataset_repo_id, *args: "/cache/fitted-tokenizer",
+    )
+
+    def fake_from_pretrained(cls, *args, **kwargs):
+        calls.append(kwargs)
+        return SimpleNamespace(steps=[])
+
+    monkeypatch.setattr(factory.PolicyProcessorPipeline, "from_pretrained", classmethod(fake_from_pretrained))
+
+    factory.make_pre_post_processors(
+        config,
+        pretrained_path="checkpoint",
+        dataset_repo_id="user/dataset",
+    )
+
+    assert calls[0]["overrides"] == {
+        "action_tokenizer_processor": {"action_tokenizer_name": "/cache/fitted-tokenizer"}
+    }

tests/policies/pi0_pi05/test_pi05.py

@@ -16,8 +16,12 @@
 
 """Test script to verify PI0.5 (pi05) support in PI0 policy"""
 
+from types import SimpleNamespace
+
 import pytest
 import torch
+from safetensors.torch import save_file
+from torch import nn
 
 pytest.importorskip("transformers")
 
@@ -31,6 +35,26 @@ from lerobot.utils.random_utils import set_seed
 from tests.utils import require_cuda, require_hf_token  # noqa: E402
 
 
+class _CheckpointPolicy(PI05Policy):
+    def __init__(self, config, **kwargs):
+        nn.Module.__init__(self)
+        self.config = config
+        self.loaded_state_dict = None
+
+    def load_state_dict(self, state_dict, strict=True, assign=False):
+        self.loaded_state_dict = state_dict
+        return [], []
+
+
+def test_from_pretrained_loads_existing_single_file_checkpoint(tmp_path):
+    save_file({"weight": torch.tensor([1.0])}, tmp_path / "model.safetensors")
+
+    policy = _CheckpointPolicy.from_pretrained(tmp_path, config=SimpleNamespace())
+
+    assert policy.loaded_state_dict is not None
+    torch.testing.assert_close(policy.loaded_state_dict["model.weight"], torch.tensor([1.0]))
+
+
 @require_cuda
 @require_hf_token
 def test_policy_instantiation():