refactor(pi052): use standard processor loading

src/lerobot/policies/factory.py

@@ -66,90 +66,6 @@ from .wall_x.configuration_wall_x import WallXConfig
 from .xvla.configuration_xvla import XVLAConfig
 
 
-def _restore_pi052_pretrained_state(
-    preprocessor: PolicyProcessorPipeline,
-    postprocessor: PolicyProcessorPipeline,
-    pretrained_path: str,
-) -> None:
-    """Restore checkpoint state into fresh PI052 pipelines that cannot JSON-roundtrip.
-
-    Steps are paired by position and registry name to prevent loading state into the wrong processor.
-    """
-    import json  # noqa: PLC0415
-    import logging  # noqa: PLC0415
-    from pathlib import Path  # noqa: PLC0415
-
-    from safetensors.torch import load_file  # noqa: PLC0415
-
-    log = logging.getLogger(__name__)
-
-    base = Path(pretrained_path)
-    if not base.exists():
-        # Resolve Hub processor configs and state files for the fresh PI052 pipelines.
-        try:
-            from huggingface_hub import snapshot_download  # noqa: PLC0415
-
-            base = Path(
-                snapshot_download(
-                    repo_id=str(pretrained_path),
-                    allow_patterns=["policy_preprocessor*", "policy_postprocessor*"],
-                )
-            )
-        except Exception as exc:  # noqa: BLE001
-            log.warning(
-                "PI052 state restore: %s is not a local dir and could not be resolved "
-                "as a hub repo (%s); normalizer stats left at fresh init",
-                pretrained_path,
-                exc,
-            )
-            return
-
-    for pipeline, config_filename in [
-        (preprocessor, f"{POLICY_PREPROCESSOR_DEFAULT_NAME}.json"),
-        (postprocessor, f"{POLICY_POSTPROCESSOR_DEFAULT_NAME}.json"),
-    ]:
-        config_path = base / config_filename
-        if not config_path.exists():
-            continue
-        saved = json.loads(config_path.read_text())
-
-        for idx, (saved_step, fresh_step) in enumerate(
-            zip(saved.get("steps", []), pipeline.steps, strict=False)
-        ):
-            state_file = saved_step.get("state_file")
-            if not state_file:
-                continue
-            saved_name = saved_step.get("registry_name")
-            fresh_name = getattr(type(fresh_step), "_registry_name", None)
-            if saved_name and fresh_name and saved_name != fresh_name:
-                log.warning(
-                    "PI052 state restore: %s step %d registry name mismatch "
-                    "(saved=%s, fresh=%s); skipping %s",
-                    config_filename,
-                    idx,
-                    saved_name,
-                    fresh_name,
-                    state_file,
-                )
-                continue
-            state_path = base / state_file
-            if not state_path.exists():
-                log.warning(
-                    "PI052 state restore: %s missing at %s; %s left at fresh init",
-                    state_file,
-                    base,
-                    fresh_name,
-                )
-                continue
-            fresh_step.load_state_dict(load_file(str(state_path)))
-            log.info(
-                "PI052 state restore: loaded %s into %s (step %d)",
-                state_file,
-                fresh_name,
-                idx,
-            )
-
-
 def _reconnect_relative_absolute_steps(
     preprocessor: PolicyProcessorPipeline, postprocessor: PolicyProcessorPipeline
 ) -> None:
@@ -395,33 +311,54 @@ def make_pre_post_processors(
         NotImplementedError: If a processor factory is not implemented for the given
             policy configuration type.
     """
-    if pretrained_path and getattr(policy_cfg, "type", None) == "pi052":
-        # Rebuild non-serializable PI052 steps, then restore their saved state.
-        from .pi052.processor_pi052 import make_pi052_pre_post_processors
-
-        preprocessor, postprocessor = make_pi052_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-            dataset_repo_id=kwargs.get("dataset_repo_id"),
-        )
-        _restore_pi052_pretrained_state(preprocessor, postprocessor, pretrained_path)
-        _reconnect_relative_absolute_steps(preprocessor, postprocessor)
-        return preprocessor, postprocessor
-
     if (
         pretrained_path
-        and getattr(policy_cfg, "type", None) == "pi0_fast"
+        and getattr(policy_cfg, "type", None) in {"pi0_fast", "pi052"}
         and getattr(policy_cfg, "auto_fit_fast_tokenizer", False)
         and kwargs.get("dataset_repo_id") is not None
     ):
-        from .pi0_fast.processor_pi0_fast import make_pi0_fast_pre_post_processors
+        if policy_cfg.type == "pi052":
+            from .pi052.processor_pi052 import make_pi052_pre_post_processors
 
-        return make_pi0_fast_pre_post_processors(
+            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 {}),
+        }
+
     if pretrained_path:
         if isinstance(policy_cfg, GrootConfig):
             from .groot.processor_groot import make_groot_pre_post_processors_from_pretrained

src/lerobot/policies/pi052/text_processor_pi052.py

@@ -12,10 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tokenize PI052's plain-text rendered messages and build text/action supervision masks.
+"""Tokenize PI052 messages and build text/action supervision masks."""
-
-PaliGemma is not chat-trained, so messages use explicit role delimiters instead of a chat template.
-"""
 
 from __future__ import annotations
 
@@ -37,13 +34,7 @@ logger = logging.getLogger(__name__)
 
 
 def discretize_state_str(state_row: Any) -> str:
-    """Discretize a single normalized state vector into 256 bins, space-joined.
+    """Format one normalized state row with PI0.5's 256-bin convention."""
-
-    Mirrors pi05's ``Pi05PrepareStateTokenizerProcessorStep`` (same bins /
-    convention) so pi052's low-level action prompt carries proprioception in
-    the exact format pi05 was trained on. Expects state already normalized by
-    the upstream ``NormalizerProcessorStep``.
-    """
     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())
@@ -73,16 +64,7 @@ def _content_to_text(content: Any) -> str:
 
 
 def _flatten_say_tool_calls(message: dict[str, Any]) -> dict[str, Any]:
-    """Serialize assistant ``say`` tool calls into a ``<say>...</say>`` marker.
+    """Move ``say`` tool calls into text markers that PaliGemma can learn."""
-
-    PaliGemma's flat text prompt has no notion of structured tool calls,
-    and ``_format_messages`` only reads ``role`` / ``content`` — so
-    without this a ``say`` tool call is dropped entirely and never
-    supervised. Rewriting it into the content text as a ``<say>...</say>``
-    marker lets the LM head learn to emit it; the runtime parses it back
-    via ``_split_plan_and_say``. Messages without ``say`` tool calls are
-    returned unchanged (the structured calls, if any, are still dropped).
-    """
     tool_calls = message.get("tool_calls")
     if not tool_calls:
         return message
@@ -115,15 +97,7 @@ def _flatten_say_tool_calls(message: dict[str, Any]) -> dict[str, Any]:
 
 
 def _strip_blocks(message: dict[str, Any]) -> dict[str, Any]:
-    """Normalise a message's content to a plain string.
+    """Flatten text blocks and drop image blocks handled by observation inputs."""
-
-    The recipe renderer can emit ``content`` as a string OR as a list
-    of HF-style multimodal blocks (``{type: text, text: ...}``,
-    ``{type: image, feature: ...}``). PaliGemma's text tokenizer can
-    only consume strings, so we flatten: drop image blocks (cameras
-    flow through ``observation.images.*`` separately) and join text
-    block texts.
-    """
     new = dict(message)
     new.pop("stream", None)
     new.pop("target", None)
@@ -166,24 +140,13 @@ def _sample_indices(value: Any, batch_size: int) -> list[int | None]:
     return [int(value)] * batch_size
 
 
-# Convert normalized Qwen2.5-VL coordinates to PaliGemma's resolution-independent <loc> range.
-
 _VQA_COORD_SCALE = 1000.0
 
 
 def register_paligemma_loc_tokens(tokenizer: Any) -> Any:
-    """Make PaliGemma's ``<locDDDD>`` ids match on raw text — single tokens.
+    """Register PaliGemma's reserved ``<locDDDD>`` strings as single tokens.
 
-    PaliGemma reserves vocab ids [256000, 257023] for ``<locDDDD>``
+    Without registration, the stock tokenizer splits each location into generic text pieces.
-    (detection / pointing) tokens, but the *stock* tokenizer does NOT
-    match them when encoding raw text — it BPE-splits ``<loc0162>`` into
-    7 pieces (``<``, ``loc``, ``0``, ``1``, ``6``, ``2``, ``>``). Training
-    the LM head on a ``<loc>`` target then supervises those 7 generic
-    BPE pieces instead of one detection-vocab id, the LM head learns to
-    emit the *character sequence*, and those pieces' logits dominate
-    other turns (the ``<loc>``-salad on subtasks). Registering the loc
-    tokens once makes them tokenize as their single ids (256000+idx),
-    leveraging PaliGemma's detection prior properly. Idempotent.
     """
     if "<loc0000>" in getattr(tokenizer, "added_tokens_encoder", {}):
         return tokenizer
@@ -198,26 +161,9 @@ def _loc_token(coord: float, scale: float = _VQA_COORD_SCALE) -> str:
 
 
 def _vqa_answer_to_loc(answer: dict[str, Any]) -> str | None:
-    """Convert a bbox / keypoint VQA answer dict to PaliGemma ``<loc>`` text.
+    """Convert normalized bbox/keypoint answers to label-first PaliGemma locations.
 
-    Input coordinates are in Qwen2.5-VL's 0–1000 normalized space (see
+    Label-first targets prevent location tokens from dominating every assistant turn; non-spatial answers return ``None``.
-    module-level note). y is emitted before x for each coordinate pair
-    (PaliGemma convention), with the integer indices in [0, 1023].
-
-    **Format: label first, locs after.** PaliGemma's pretraining puts
-    locs first (``<loc><loc> label``), but for our small-dataset VQA
-    blend that turns the LM head into a loc-emission attractor at every
-    ``Assistant:`` position — VQA targets share their first supervised
-    token with ~25% of all text samples, and the head collapses to
-    emitting ``<loc>`` regardless of the prompt. Putting the label
-    first (``label <locY><locX>``) means every text sample (subtask,
-    memory, VQA, …) starts the supervised target with a real word,
-    breaking the attractor. The model still learns the loc vocabulary
-    for the *spatial* portion of the answer; it just can't fire it as
-    the first generation step from a clean prompt.
-
-    Returns ``None`` for non-spatial answers (count / attribute /
-    spatial-relation) — those keep their JSON form.
     """
     point = answer.get("point")
     if isinstance(point, list | tuple) and len(point) == 2 and "point_format" in answer:
@@ -256,13 +202,7 @@ def _messages_vqa_to_loc(
     messages: list[dict[str, Any]],
     target_indices: list[int],
 ) -> list[dict[str, Any]]:
-    """Rewrite bbox / keypoint VQA *target* answers from JSON to ``<loc>`` text.
+    """Rewrite spatial VQA target JSON as camera-independent ``<loc>`` text."""
-
-    Each target turn whose content parses as a spatial VQA answer is
-    converted. Non-spatial answers and subtask / memory targets (plain
-    text → not JSON) are left untouched. Camera-independent: VQA coords
-    are 0–1000 normalized, so no observation lookup is needed.
-    """
     if not target_indices:
         return messages
     out = list(messages)
@@ -275,7 +215,7 @@ def _messages_vqa_to_loc(
         try:
             answer = json.loads(content)
         except (ValueError, TypeError):
-            continue  # subtask / memory targets are plain text — skip
+            continue
         if not isinstance(answer, dict):
             continue
         loc_text = _vqa_answer_to_loc(answer)
@@ -289,22 +229,9 @@ def _format_messages(
     target_indices: list[int] | None = None,
     eos_token: str | None = None,
 ) -> tuple[str, list[tuple[int, int]]]:
-    """Concatenate messages into the π0.5-style flat prompt.
+    """Build the flat PI0.5 prompt and each message's payload span.
 
-    When both ``target_indices`` and ``eos_token`` are given, the EOS
+    Supervised targets include EOS so generation learns when to stop.
-    string is appended to each supervised target turn's content and the
-    returned span covers it — so the label builder marks the EOS token
-    as a supervised label. That teaches the LM head where the answer
-    *ends*: without an EOS in the target span the model is never given a
-    stop signal and rambles to ``max_length`` at inference. Inference
-    callers omit both args (no EOS baked into the prompt — the model
-    generates it and ``select_message`` stops on it).
-
-    Returns:
-        prompt:       the full text the tokenizer will consume.
-        msg_spans:    list of ``(char_start, char_end)`` covering each
-                      message's supervised payload (content, plus the
-                      appended EOS for target turns) within ``prompt``.
     """
     targets = set(target_indices or [])
     parts: list[str] = []
@@ -313,12 +240,8 @@ def _format_messages(
     for i, m in enumerate(messages):
         role = m.get("role", "user")
         content = m.get("content", "") or ""
-        # Supervise the explicit role format used again during generation.
         header = f"{role.capitalize()}: "
-        # Include EOS only in supervised target spans so generation learns to stop.
         body = content + eos_token if (eos_token and i in targets) else content
-        # span covers the content (+ EOS) portion only — never the role
-        # tag — so labels are computed over the supervised payload.
         full = header + body + "\n"
         start = cursor + len(header)
         end = start + len(body)
@@ -331,11 +254,7 @@ def _format_messages(
 @dataclass
 @ProcessorStepRegistry.register(name="pi052_text_tokenizer")
 class PI052TextTokenizerStep(ProcessorStep):
-    """Render messages → token ids + label mask + predict_actions flag.
+    """Convert flat role-delimited messages into tokens and supervision masks."""
-
-    No chat template; concatenates messages as
-    ``User: ... \\nAssistant: ...`` text.
-    """
 
     tokenizer_name: str = "google/paligemma-3b-pt-224"
     max_length: int = 200
@@ -350,6 +269,19 @@ class PI052TextTokenizerStep(ProcessorStep):
     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
@@ -364,13 +296,10 @@ class PI052TextTokenizerStep(ProcessorStep):
         messages = complementary.get("messages") or []
 
         if not messages:
-            # Preserve the transition for the plain PI0.5 prompt fallback.
             return transition
 
         tokenizer = self._ensure_tokenizer()
-        # Add normalized proprioception to low-level prompts, matching PI0.5.
         state_all = (transition.get(TransitionKey.OBSERVATION) or {}).get(OBS_STATE)
-        # Normalized VQA coordinates need no camera lookup.
         if _is_batched_messages(messages):
             indices_iter = _sample_indices(complementary.get("index"), len(messages))
             encoded = [
@@ -429,7 +358,6 @@ class PI052TextTokenizerStep(ProcessorStep):
         sample_idx: int | None = None,
         state_row: Any = None,
     ) -> tuple[Tensor, Tensor, Tensor, Tensor, str]:
-        # Remap target indices after optional context dropout.
         if (
             self.plan_dropout_prob
             or self.memory_dropout_prob
@@ -443,12 +371,10 @@ class PI052TextTokenizerStep(ProcessorStep):
                 sample_idx=sample_idx,
             )
 
-        # Rewrite normalized VQA answers as PaliGemma <loc> text.
         messages = _messages_vqa_to_loc(messages, target_indices)
 
-        # Flatten ``say`` calls because PaliGemma receives plain text.
         messages = [_strip_blocks(_flatten_say_tool_calls(m)) for m in messages]
-        # Add state only to low-level action prompts; keep higher-level streams state-free.
+        # 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):
@@ -456,8 +382,6 @@ class PI052TextTokenizerStep(ProcessorStep):
                     base = _content_to_text(m.get("content", ""))
                     m["content"] = f"{base}, State: {state_str};"
                     break
-        # Append EOS to supervised target turns so the LM head learns to
-        # stop (the span covers it → it becomes a supervised label).
         prompt, spans = _format_messages(messages, target_indices, getattr(tokenizer, "eos_token", None))
 
         encoded = tokenizer(
@@ -472,10 +396,8 @@ class PI052TextTokenizerStep(ProcessorStep):
 
         input_ids = encoded["input_ids"][0]
         attention_mask = encoded["attention_mask"][0].bool()
-        offsets = encoded["offset_mapping"][0]  # (seq, 2), char (start,end)
+        offsets = encoded["offset_mapping"][0]
 
-        # Build label mask: -100 everywhere except over supervised
-        # target message char ranges.
         labels = torch.full_like(input_ids, fill_value=-100)
         for idx in target_indices:
             if idx >= len(spans):
@@ -489,7 +411,6 @@ class PI052TextTokenizerStep(ProcessorStep):
                     continue
                 labels[token_pos] = input_ids[token_pos]
 
-        # Scan all streams because low-level flow may intentionally have no text target.
         predict_actions = torch.tensor(
             bool(any(s == "low_level" for s in message_streams)),
             dtype=torch.bool,
@@ -503,16 +424,11 @@ class PI052TextTokenizerStep(ProcessorStep):
         complementary: dict[str, Any],
         sample_idx: int | None = None,
     ) -> tuple[list[dict[str, Any]], list[int]]:
-        """Drop messages classified as plan/memory/subtask context.
+        """Drop sampled context messages and remap the retained target positions."""
-
-        Targets are *never* dropped (they're the supervised payload).
-        Re-maps target_indices to the new positions after drops.
-        """
         import random  # noqa: PLC0415
 
         seed = self.dropout_seed
         if seed is None:
-            # Use the canonical row index to avoid identical dropout across an epoch.
             seed_src = sample_idx if sample_idx is not None else complementary.get("index", 0)
             try:
                 if hasattr(seed_src, "item"):
@@ -538,7 +454,6 @@ class PI052TextTokenizerStep(ProcessorStep):
                 continue
             keep_indices.append(idx)
 
-        # Build remap and apply
         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]
@@ -551,7 +466,7 @@ class PI052TextTokenizerStep(ProcessorStep):
 
 
 def _classify_for_dropout(message: dict[str, Any]) -> str | None:
-    """Heuristic content-prefix classifier (plan / memory / subtask)."""
+    """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"]

src/lerobot/processor/render_messages_processor.py

@@ -16,7 +16,7 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass
+from dataclasses import asdict, dataclass
 from typing import Any
 
 from lerobot.configs import PipelineFeatureType, PolicyFeature
@@ -32,17 +32,18 @@ from .pipeline import ProcessorStep, ProcessorStepRegistry
 @dataclass
 @ProcessorStepRegistry.register(name="render_messages_processor")
 class RenderMessagesStep(ProcessorStep):
-    """Processor step that turns raw language columns into rendered chat messages.
+    """Render language columns into recipe-defined messages and supervision metadata."""
-
-    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 {}
@@ -169,7 +170,7 @@ def _batch_value(value: Any, index: int) -> Any:
         return None
     if isinstance(value, list):
         return value[index]
-    if hasattr(value, "ndim") and getattr(value, "ndim") > 0:
+    if hasattr(value, "ndim") and value.ndim > 0:
         return _scalar(value[index])
     return _scalar(value)
 

tests/policies/pi052/test_pi052_text_processor.py

@@ -23,15 +23,26 @@ 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,
     _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
+from lerobot.utils.constants import (
+    OBS_LANGUAGE_ATTENTION_MASK,
+    OBS_LANGUAGE_TOKENS,
+    POLICY_PREPROCESSOR_DEFAULT_NAME,
+)
 
 
 def _say_call(text):
@@ -88,6 +99,51 @@ def test_format_messages_without_eos_args_is_unchanged():
     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 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