diff --git a/src/lerobot/policies/factory.py b/src/lerobot/policies/factory.py
index 40d839ab1..8759e01d5 100644
--- a/src/lerobot/policies/factory.py
+++ b/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
diff --git a/src/lerobot/policies/pi052/text_processor_pi052.py b/src/lerobot/policies/pi052/text_processor_pi052.py
index d3036a692..45288a8e2 100644
--- a/src/lerobot/policies/pi052/text_processor_pi052.py
+++ b/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.
-
-PaliGemma is not chat-trained, so messages use explicit role delimiters instead of a chat template.
-"""
+"""Tokenize PI052 messages and build text/action supervision masks."""
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.
-
- 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``.
- """
+ """Format one normalized state row with PI0.5's 256-bin convention."""
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 ``...`` marker.
-
- 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 ``...``
- 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).
- """
+ """Move ``say`` tool calls into text markers that PaliGemma can learn."""
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.
-
- 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.
- """
+ """Flatten text blocks and drop image blocks handled by observation inputs."""
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 range.
-
_VQA_COORD_SCALE = 1000.0
def register_paligemma_loc_tokens(tokenizer: Any) -> Any:
- """Make PaliGemma's ```` ids match on raw text — single tokens.
+ """Register PaliGemma's reserved ```` strings as single tokens.
- PaliGemma reserves vocab ids [256000, 257023] for ````
- (detection / pointing) tokens, but the *stock* tokenizer does NOT
- match them when encoding raw text — it BPE-splits ```` into
- 7 pieces (``<``, ``loc``, ``0``, ``1``, ``6``, ``2``, ``>``). Training
- the LM head on a ```` 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 ````-salad on subtasks). Registering the loc
- tokens once makes them tokenize as their single ids (256000+idx),
- leveraging PaliGemma's detection prior properly. Idempotent.
+ Without registration, the stock tokenizer splits each location into generic text pieces.
"""
if "" 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 ```` text.
+ """Convert normalized bbox/keypoint answers to label-first PaliGemma locations.
- Input coordinates are in Qwen2.5-VL's 0–1000 normalized space (see
- 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 (`` 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 ```` regardless of the prompt. Putting the label
- first (``label ``) 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.
+ Label-first targets prevent location tokens from dominating every assistant turn; non-spatial answers return ``None``.
"""
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 ```` 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.
- """
+ """Rewrite spatial VQA target JSON as camera-independent ```` text."""
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
- 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``.
+ Supervised targets include EOS so generation learns when to stop.
"""
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.
-
- No chat template; concatenates messages as
- ``User: ... \\nAssistant: ...`` text.
- """
+ """Convert flat role-delimited messages into tokens and supervision masks."""
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 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.
-
- Targets are *never* dropped (they're the supervised payload).
- Re-maps target_indices to the new positions after drops.
- """
+ """Drop sampled context messages and remap the retained target positions."""
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"]
diff --git a/src/lerobot/processor/render_messages_processor.py b/src/lerobot/processor/render_messages_processor.py
index 0b5e4923f..2fca46e7e 100644
--- a/src/lerobot/processor/render_messages_processor.py
+++ b/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.
-
- 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.
- """
+ """Render language columns into recipe-defined messages and supervision metadata."""
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)
diff --git a/tests/policies/pi052/test_pi052_text_processor.py b/tests/policies/pi052/test_pi052_text_processor.py
index f56ed16e8..ae0a80c56 100644
--- a/tests/policies/pi052/test_pi052_text_processor.py
+++ b/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