refactor(factory): remove PI052 processor overrides

src/lerobot/policies/factory.py

@@ -317,47 +317,15 @@ def make_pre_post_processors(
         and getattr(policy_cfg, "auto_fit_fast_tokenizer", False)
         and kwargs.get("dataset_repo_id") is not None
     ):
-        if policy_cfg.type == "pi052":
-            from .pi052.processor_pi052 import make_pi052_pre_post_processors
+        from .pi052.fit_fast_tokenizer import resolve_fast_tokenizer
 
-            factory = make_pi052_pre_post_processors
-        else:
-            from .pi0_fast.processor_pi0_fast import make_pi0_fast_pre_post_processors
-
-            factory = make_pi0_fast_pre_post_processors
-
-        return factory(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-            dataset_repo_id=kwargs.get("dataset_repo_id"),
-        )
-
-    if (
-        pretrained_path
-        and getattr(policy_cfg, "type", None) == "pi052"
-        and getattr(policy_cfg, "recipe_path", None)
-    ):
-        from .pi052.processor_pi052 import _load_recipe
-
-        pi052_overrides = {
-            "render_messages_processor": {"recipe": _load_recipe(policy_cfg.recipe_path)},
-            "pi052_text_tokenizer": {
-                "tokenizer_name": "google/paligemma-3b-pt-224",
-                "max_length": policy_cfg.tokenizer_max_length,
-                "plan_dropout_prob": policy_cfg.plan_dropout_prob,
-                "memory_dropout_prob": policy_cfg.memory_dropout_prob,
-                "subtask_dropout_prob": policy_cfg.subtask_dropout_prob,
-            },
-            "action_tokenizer_processor": {
-                "action_tokenizer_name": policy_cfg.action_tokenizer_name,
-                "max_action_tokens": policy_cfg.max_action_tokens,
-                "fast_skip_tokens": policy_cfg.fast_skip_tokens,
-            },
-        }
-        kwargs["preprocessor_overrides"] = {
-            **pi052_overrides,
-            **(kwargs.get("preprocessor_overrides") or {}),
+        overrides = dict(kwargs.get("preprocessor_overrides") or {})
+        action_tokenizer_override = {
+            **overrides.get("action_tokenizer_processor", {}),
+            "action_tokenizer_name": resolve_fast_tokenizer(policy_cfg, kwargs.get("dataset_repo_id")),
         }
+        overrides["action_tokenizer_processor"] = action_tokenizer_override
+        kwargs["preprocessor_overrides"] = overrides
 
     if pretrained_path:
         if isinstance(policy_cfg, GrootConfig):

tests/policies/pi052/test_pi052_text_processor.py

@@ -23,13 +23,9 @@ supervised target span must end with an EOS token so the LM head learns
 to stop instead of rambling to ``max_length`` at inference).
 """
 
-from types import SimpleNamespace
-
 import torch
 
 from lerobot.configs.recipe import MessageTurn, TrainingRecipe
-from lerobot.policies import factory
-from lerobot.policies.pi052.configuration_pi052 import PI052Config
 from lerobot.policies.pi052.text_processor_pi052 import (
     PI052TextTokenizerStep,
     _flatten_say_tool_calls,
@@ -126,24 +122,6 @@ def test_pi052_steps_roundtrip_through_standard_pipeline_loader(tmp_path):
     assert loaded.steps[1].dropout_seed == 3
 
 
-def test_pi052_legacy_checkpoint_uses_standard_loader_with_rebuild_overrides(monkeypatch):
-    calls = []
-
-    def fake_from_pretrained(cls, *args, **kwargs):
-        calls.append(kwargs)
-        return SimpleNamespace(steps=[])
-
-    monkeypatch.setattr(PolicyProcessorPipeline, "from_pretrained", classmethod(fake_from_pretrained))
-    config = PI052Config(recipe_path="recipes/subtask_mem.yaml", auto_fit_fast_tokenizer=False)
-
-    factory.make_pre_post_processors(config, pretrained_path="checkpoint")
-
-    overrides = calls[0]["overrides"]
-    assert isinstance(overrides["render_messages_processor"]["recipe"], TrainingRecipe)
-    assert overrides["pi052_text_tokenizer"]["max_length"] == config.tokenizer_max_length
-    assert overrides["action_tokenizer_processor"]["action_tokenizer_name"] == config.action_tokenizer_name
-
-
 def _eos_char_id() -> int:
     """Token id _CharTokenizer assigns to its 1-char EOS."""
     return ord("\x1f") % 251 + 1

tests/policies/pi0_fast/test_pi0_fast_tokenizer_fit.py

@@ -14,8 +14,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from types import SimpleNamespace
+
 from lerobot.policies import factory
-from lerobot.policies.pi0_fast import processor_pi0_fast
 from lerobot.policies.pi0_fast.configuration_pi0_fast import PI0FastConfig
 from lerobot.policies.pi052 import fit_fast_tokenizer as fit_module
 
@@ -47,17 +48,28 @@ def test_pi0_fast_resolves_dataset_specific_tokenizer(monkeypatch, tmp_path):
     }
 
 
-def test_pretrained_pi0_fast_rebuilds_processor_only_during_dataset_fit(monkeypatch):
+def test_pretrained_pi0_fast_overrides_only_fitted_tokenizer(monkeypatch):
     config = PI0FastConfig(auto_fit_fast_tokenizer=True)
-    expected = (object(), object())
+    calls = []
 
-    monkeypatch.setattr(processor_pi0_fast, "make_pi0_fast_pre_post_processors", lambda **_: expected)
-
-    assert (
-        factory.make_pre_post_processors(
-            config,
-            pretrained_path="checkpoint",
-            dataset_repo_id="user/dataset",
-        )
-        == expected
+    monkeypatch.setattr(
+        fit_module,
+        "resolve_fast_tokenizer",
+        lambda config, dataset_repo_id: "/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"}
+    }