fix pi052 FAST training consistency

Align tokenizer fitting and loss reduction with the effective training dataset, and fail early when FAST supervision cannot be produced safely.

Co-authored-by: Cursor <cursoragent@cursor.com>

tests/policies/pi052/test_pi052_fast_action_loss.py

@@ -16,14 +16,18 @@
 
 """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 _fast_lin_ce  # noqa: E402
+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):
@@ -104,3 +108,55 @@ def test_fast_ce_returns_zero_when_no_action_code_positions_are_valid():
     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,92 @@
+#!/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.policies.pi052.fit_fast_tokenizer import (
+    _apply_relative_actions,
+    _dataset_signature,
+    _is_global_leader,
+    _normalize_actions,
+    _select_episode_indices,
+)
+
+
+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]]])

tests/policies/pi0_fast/test_pi0_fast_tokenizer_fit.py

@@ -47,6 +47,14 @@ def test_pi0_fast_resolves_dataset_specific_tokenizer(monkeypatch, 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,
     }
 
 
@@ -62,13 +70,13 @@ def test_fast_fit_failure_is_not_silently_replaced(monkeypatch, tmp_path):
         fit_module.resolve_fast_tokenizer(config, "user/dataset")
 
 
-def test_each_node_uses_its_local_rank_zero_as_fit_leader(monkeypatch):
+def test_only_global_rank_zero_fits_shared_tokenizer(monkeypatch):
     monkeypatch.setenv("RANK", "8")
     monkeypatch.setenv("LOCAL_RANK", "0")
-    assert fit_module._is_local_leader()
+    assert not fit_module._is_global_leader()
 
-    monkeypatch.setenv("LOCAL_RANK", "1")
-    assert not fit_module._is_local_leader()
+    monkeypatch.setenv("RANK", "0")
+    assert fit_module._is_global_leader()
 
 
 def test_pretrained_pi0_fast_overrides_only_fitted_tokenizer(monkeypatch):
@@ -78,7 +86,7 @@ def test_pretrained_pi0_fast_overrides_only_fitted_tokenizer(monkeypatch):
     monkeypatch.setattr(
         fit_module,
         "resolve_fast_tokenizer",
-        lambda config, dataset_repo_id: "/cache/fitted-tokenizer",
+        lambda config, dataset_repo_id, *args: "/cache/fitted-tokenizer",
     )
 
     def fake_from_pretrained(cls, *args, **kwargs):

tests/processor/test_tokenizer_processor.py

@@ -94,6 +94,7 @@ def test_action_tokenizer_config_preserves_token_mapping():
     processor.max_action_tokens = 384
     processor.fast_skip_tokens = 64
     processor.paligemma_tokenizer_name = "custom/paligemma"
+    processor.allow_truncation = False
     processor.action_tokenizer_name = "custom/fast"
     processor.action_tokenizer_input_object = None
 
@@ -102,10 +103,31 @@ def test_action_tokenizer_config_preserves_token_mapping():
         "max_action_tokens": 384,
         "fast_skip_tokens": 64,
         "paligemma_tokenizer_name": "custom/paligemma",
+        "allow_truncation": False,
         "action_tokenizer_name": "custom/fast",
     }
 
 
+def test_action_tokenizer_can_reject_truncated_sequences():
+    processor = object.__new__(ActionTokenizerProcessorStep)
+    processor.max_action_tokens = 4
+    processor.fast_skip_tokens = 128
+    processor.allow_truncation = False
+    processor.action_tokenizer = lambda _actions: [1, 2, 3]
+    processor._paligemma_tokenizer = type(
+        "Tokenizer",
+        (),
+        {
+            "vocab_size": 1000,
+            "bos_token_id": 2,
+            "encode": lambda _self, text, **_kwargs: [10, 11] if text == "Action: " else [12, 1],
+        },
+    )()
+
+    with pytest.raises(ValueError, match="max_action_tokens=4"):
+        processor._tokenize_action(torch.zeros(1, 2, 1))
+
+
 @pytest.fixture
 def mock_tokenizer():
     """Provide a mock tokenizer for testing."""