From 6795b22b1e86eb81d3bc4b4f27d3f31de5886add Mon Sep 17 00:00:00 2001 From: Pepijn Date: Wed, 15 Jul 2026 16:07:14 +0200 Subject: [PATCH] refactor(factory): remove PI052 processor overrides --- src/lerobot/policies/factory.py | 46 +++---------------- .../pi052/test_pi052_text_processor.py | 22 --------- .../pi0_fast/test_pi0_fast_tokenizer_fit.py | 36 ++++++++++----- 3 files changed, 31 insertions(+), 73 deletions(-) diff --git a/src/lerobot/policies/factory.py b/src/lerobot/policies/factory.py index 8759e01d5..6235faad8 100644 --- a/src/lerobot/policies/factory.py +++ b/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): diff --git a/tests/policies/pi052/test_pi052_text_processor.py b/tests/policies/pi052/test_pi052_text_processor.py index ae0a80c56..8e1ea02c6 100644 --- a/tests/policies/pi052/test_pi052_text_processor.py +++ b/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 diff --git a/tests/policies/pi0_fast/test_pi0_fast_tokenizer_fit.py b/tests/policies/pi0_fast/test_pi0_fast_tokenizer_fit.py index d435c9f09..7f928a32b 100644 --- a/tests/policies/pi0_fast/test_pi0_fast_tokenizer_fit.py +++ b/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"} + }