fix(smolvla2): make HighLevelSubtaskFwd actually fire at low hz + quiet startup log

Two runtime fixes that surfaced from on-robot testing.

(1) HighLevelSubtaskFwd was double-gated: HzTrigger fires every period
(e.g. every 5s at --high_level_hz=0.2) AND the step requires the
action queue to be empty. The queue-empty window is brief (~tens of
ms between drain and refill) and almost never coincides with the
low-hz timer, so HL effectively never fired and the subtask shown
in the runtime panel stayed on the dataset's frame-0 annotation.

Add HzTrigger.rearm() and have HighLevelSubtaskFwd call it when
skipping due to queue-non-empty — the trigger stays armed and tries
again on the next tick instead of waiting another full period.
LowLevelForward keeps the original "skip" semantics because chunk_hz
is meant as a true upper bound on chunk-generation rate.

(2) The "robot state at startup" warning in _build_robot_observation_provider
was meant to fire once but wasn't gated by _resize_logged like the
sibling "camera ... live=AxB" warning. Result: it spammed every
observation tick (~1-2s). Gate it on first_call (snapshot of
_resize_logged["done"]) so both logs fire once at session start.

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

src/lerobot/policies/pi05/configuration_pi05.py

@@ -96,6 +96,19 @@ class PI05Config(PreTrainedConfig):
     optimizer_foreach: bool | None = False
     optimizer_fused: bool | None = True
 
+    # LM-head LR multiplier. The PaliGemma `lm_head` projection (and its
+    # tied `embed_tokens`) is the surface the LM head's first-token
+    # distribution depends on. With ``knowledge_insulation`` blocking
+    # action→VLM gradients, the LM head only sees gradients on text-CE
+    # samples — which can be a small fraction of the mix (e.g. ~45% in
+    # ``subtask_mem.yaml``). Under aggressive cosine LR decay the head's
+    # first-token distribution can drift back toward PaliGemma's
+    # pretrained ``<loc>`` detection prior, despite teacher-forced CE
+    # staying near zero. Boosting just the LM-head LR (e.g. 5x) keeps
+    # the head pinned to fine-tuning targets without perturbing the
+    # backbone / vision tower / action expert. Default 1.0 = no change.
+    lm_head_lr_scale: float = 1.0
+
     # Scheduler settings: see openpi `CosineDecaySchedule`
     # Note: These will auto-scale if --steps < scheduler_decay_steps
     # For example, --steps=3000 will scale warmup to 100 and decay to 3000

src/lerobot/policies/pi05/modeling_pi05.py

@@ -1133,8 +1133,62 @@ class PI05Policy(PreTrainedPolicy):
 
         return fixed_state_dict
 
-    def get_optim_params(self) -> dict:
-        return self.parameters()
+    def get_optim_params(self):
+        """Return policy parameters, optionally split into LR-scaled groups.
+
+        When ``config.lm_head_lr_scale != 1.0``, the PaliGemma ``lm_head``
+        and its tied ``embed_tokens`` are placed in their own param
+        group with ``lr = base_lr * lm_head_lr_scale``. The cosine
+        scheduler multiplies both groups by the same lambda each step,
+        so the ratio is preserved across decay. Default ``1.0`` =
+        return ``self.parameters()`` (back-compat with existing checkpoints
+        and configs).
+        """
+        scale = float(getattr(self.config, "lm_head_lr_scale", 1.0))
+        if scale == 1.0:
+            return self.parameters()
+        head_params: list[torch.nn.Parameter] = []
+        other_params: list[torch.nn.Parameter] = []
+        # Both ``lm_head.weight`` and the tied ``embed_tokens.weight`` —
+        # boosting only the projection without the embedding pulls them
+        # apart and breaks the tie that PaliGemma was pre-trained with.
+        head_substrings = (
+            "paligemma_with_expert.paligemma.lm_head.",
+            "paligemma_with_expert.paligemma.model.language_model.embed_tokens.",
+        )
+        for name, p in self.named_parameters():
+            if not p.requires_grad:
+                continue
+            if any(s in name for s in head_substrings):
+                head_params.append(p)
+            else:
+                other_params.append(p)
+        base_lr = float(self.config.optimizer_lr)
+        groups: list[dict[str, object]] = []
+        if other_params:
+            groups.append({"params": other_params, "lr": base_lr, "name": "policy"})
+        if head_params:
+            groups.append(
+                {"params": head_params, "lr": base_lr * scale, "name": "lm_head"}
+            )
+        # Sanity: head_substrings must match at least one parameter, otherwise
+        # the scale silently does nothing — surface that fast.
+        if not head_params:
+            raise RuntimeError(
+                "lm_head_lr_scale != 1.0 but no parameters matched the LM-head "
+                "name patterns: "
+                f"{head_substrings!r}. Did the underlying PaliGemma module rename?"
+            )
+        logging.info(
+            "PI05Policy: LM-head LR scale = %.3g (base=%.3g, head=%.3g) over "
+            "%d head params + %d other params",
+            scale,
+            base_lr,
+            base_lr * scale,
+            len(head_params),
+            len(other_params),
+        )
+        return groups
 
     def reset(self):
         """Reset internal state - called when environment resets."""

src/lerobot/policies/pi052/modeling_pi052.py

@@ -930,12 +930,23 @@ class PI052Policy(PI05Policy):
         temperature: float = 0.0,
         top_p: float = 1.0,
         tokenizer: Any = None,
+        suppress_loc_tokens: bool = False,
     ) -> str:
         """Generate text continuation from a multimodal prefix.
 
         Mirrors ``SmolVLA2Policy.select_message`` so the same
         :class:`lerobot.policies.smolvla2.inference.SmolVLA2Runtime`
         can drive π0.5 v2 unchanged.
+
+        ``suppress_loc_tokens`` masks PaliGemma's reserved ``<locDDDD>``
+        ids ([256000, 257024)) to ``-inf`` before sampling. PaliGemma's
+        pretraining puts heavy first-token mass on these ids for any
+        ``Assistant:`` continuation; with a small fine-tuning text-CE
+        budget (or aggressive LR decay) the LM head can drift back
+        toward that prior even when teacher-forced argmax stays at
+        100%. Callsites that legitimately emit ``<loc>`` (VQA spatial
+        answers) must keep this ``False``; subtask / memory / plan
+        generation should pass ``True``.
         """
         self.eval()
 
@@ -1015,6 +1026,8 @@ class PI052Policy(PI05Policy):
             if special_ids and len(generated) < min_new_tokens:
                 for sid in special_ids:
                     logits_step[..., sid] = float("-inf")
+            if suppress_loc_tokens:
+                logits_step[..., 256000:257024] = float("-inf")
             next_ids = self._sample_next_token(logits_step, temperature, top_p)
             tok_id = int(next_ids[0].item())
             generated.append(tok_id)

src/lerobot/policies/smolvla2/inference/steps.py

@@ -463,8 +463,16 @@ class HighLevelSubtaskFwd(InferenceStep):
         # of 30/sec and the robot barely moves. Tying it to the same
         # "queue empty" condition as the chunk refresh produces a
         # clean sense → think → act cycle.
+        #
+        # Rearm the trigger when skipping so a low-hz schedule
+        # (e.g. ``--high_level_hz=0.2`` = once per 5 s) doesn't lose
+        # the slot: the trigger fires once on the timer but the brief
+        # queue-empty window almost never coincides, so without rearm
+        # HL would effectively never run.
         queue = state.get("action_queue") or []
         if len(queue) > 0:
+            if hasattr(self.trigger, "rearm"):
+                self.trigger.rearm()
             return None
         ctx = _msgs_for_subtask(state)
         observation = _maybe_observation(self.observation_provider)
@@ -484,6 +492,11 @@ class HighLevelSubtaskFwd(InferenceStep):
             min_new_tokens=int(state.get("text_gen_min_new_tokens") or 0),
             temperature=float(state.get("text_gen_temperature") or 0.0),
             top_p=float(state.get("text_gen_top_p") or 1.0),
+            # Subtasks never legitimately contain PaliGemma ``<loc>``
+            # tokens — suppress them so a checkpoint whose LM head
+            # has drifted toward the pretrained loc-prior falls back
+            # to its (still-correct) text mass.
+            suppress_loc_tokens=True,
         )
         # Diagnostics: surface what the model is *actually* producing
         # at chunk boundaries, even when the output gets rejected or
@@ -526,8 +539,16 @@ class HighLevelSubtaskFwd(InferenceStep):
                 )
             return None
         if msg:
+            prev_subtask = state.get("current_subtask")
             changed = set_if_changed(state, "current_subtask", msg, label="subtask")
             if changed:
+                # Stash the just-completed subtask so ``MemoryUpdateFwd``
+                # can drop it into its prompt as ``Completed subtask:``
+                # — the recipe binds ``completed_subtask`` to
+                # ``nth_prev(style=subtask, offset=1)``, i.e. the subtask
+                # that was active *before* the change.
+                if prev_subtask:
+                    state["prior_subtask"] = prev_subtask
                 # Subtask change is a downstream trigger.
                 state.setdefault("events_this_tick", []).append("subtask_change")
                 state["subtask_repeat_count"] = 0
@@ -576,6 +597,7 @@ class MemoryUpdateFwd(InferenceStep):
             observation=observation,
             state=state,
             label="memory gen",
+            suppress_loc_tokens=True,
         )
         state["last_memory_raw"] = new_memory or ""
         if new_memory and _looks_like_gibberish(new_memory):
@@ -619,6 +641,7 @@ class UserInterjectionFwd(InferenceStep):
             observation=observation,
             state=state,
             label="plan/say gen",
+            suppress_loc_tokens=True,
         )
         if not out:
             # Don't log every empty completion — happens repeatedly on
@@ -851,21 +874,34 @@ def _msgs_for_subtask(state: dict[str, Any]) -> list[dict[str, Any]]:
 
 
 def _msgs_for_memory(state: dict[str, Any]) -> list[dict[str, Any]]:
-    """Memory-update prompt — boundary-frame tail of ``high_level_subtask``.
+    """Memory-update prompt — mirrors ``memory_update`` recipe layout.
 
-    Recipe layout on a boundary frame:
-        user:      "${task}\\nPlan: ${plan}\\nMemory: ${memory}"
-        assistant: "${subtask}"
+    Recipe layout (``subtask_mem.yaml``):
+
+        user:      "${task}"
+        assistant: "Previous memory: ${prior_memory}"     (if_present prior)
+        user:      "Completed subtask: ${completed}"       (if_present completed)
         assistant: → predicts new memory
 
-    Fired when the runtime detects a subtask transition; the
-    just-predicted subtask lives in ``state['current_subtask']``.
+    Fired by ``MemoryUpdateFwd`` on a ``subtask_change`` event:
+    ``state['current_memory']`` is the memory the policy last emitted
+    (= the ``prior_memory`` binding at training), and
+    ``state['prior_subtask']`` is the subtask that just got replaced
+    (= the ``completed_subtask`` binding at training).
     """
     msgs: list[dict[str, Any]] = [
-        {"role": "user", "content": _hirobot_user_head(state)},
+        {"role": "user", "content": state.get("task") or ""},
     ]
-    if state.get("current_subtask"):
-        msgs.append({"role": "assistant", "content": state["current_subtask"]})
+    prior_memory = state.get("current_memory")
+    if prior_memory:
+        msgs.append(
+            {"role": "assistant", "content": f"Previous memory: {prior_memory}"}
+        )
+    completed_subtask = state.get("prior_subtask")
+    if completed_subtask:
+        msgs.append(
+            {"role": "user", "content": f"Completed subtask: {completed_subtask}"}
+        )
     return msgs
 
 
@@ -925,6 +961,7 @@ def _generate_with_policy(
     min_new_tokens: int = 0,
     temperature: float = 0.0,
     top_p: float = 1.0,
+    suppress_loc_tokens: bool = False,
 ) -> str:
     """Drive ``policy.select_message`` with a chat batch (and optional obs).
 
@@ -959,13 +996,20 @@ def _generate_with_policy(
             for k, v in observation.items():
                 if isinstance(k, str) and k.startswith("observation.") and k not in batch:
                     batch[k] = v
-        return policy.select_message(
-            batch,
-            tokenizer=text_batch["tokenizer"],
-            min_new_tokens=min_new_tokens,
-            temperature=temperature,
-            top_p=top_p,
-        )
+        kwargs: dict[str, Any] = {
+            "tokenizer": text_batch["tokenizer"],
+            "min_new_tokens": min_new_tokens,
+            "temperature": temperature,
+            "top_p": top_p,
+        }
+        # Only pass ``suppress_loc_tokens`` to backbones that accept it
+        # (pi052). SmolVLA2's ``select_message`` does not, so we omit
+        # the kwarg there to avoid breaking the shared runtime.
+        import inspect  # noqa: PLC0415
+
+        if "suppress_loc_tokens" in inspect.signature(policy.select_message).parameters:
+            kwargs["suppress_loc_tokens"] = suppress_loc_tokens
+        return policy.select_message(batch, **kwargs)
     except Exception as exc:  # noqa: BLE001
         logger.warning("%s failed: %s", label, exc, exc_info=logger.isEnabledFor(logging.DEBUG))
         if state is not None:

src/lerobot/policies/smolvla2/inference/triggers.py

@@ -82,7 +82,15 @@ class Trigger(Protocol):
 
 @dataclass
 class HzTrigger:
-    """Fire at most ``hz`` times per second."""
+    """Fire at most ``hz`` times per second.
+
+    A step that gates further (e.g. ``HighLevelSubtaskFwd`` skipping
+    when the action queue is non-empty) and wants the trigger to
+    retry next tick instead of waiting a full period can call
+    :meth:`rearm` from inside ``run``. Without this, a low-hz trigger
+    (e.g. ``hz=0.2`` = once per 5 s) almost never coincides with the
+    brief queue-empty window and the step never fires at all.
+    """
 
     hz: float
     _last_seconds: float | None = field(default=None, init=False)
@@ -94,6 +102,15 @@ class HzTrigger:
             return True
         return False
 
+    def rearm(self) -> None:
+        """Mark the trigger as not having fired, so the next tick re-evaluates.
+
+        Used by a step that decided to skip after ``should_fire`` already
+        committed the firing — keeps the cadence honest without losing
+        the slot.
+        """
+        self._last_seconds = None
+
 
 @dataclass
 class EventTrigger:

src/lerobot/scripts/lerobot_smolvla2_runtime.py

@@ -772,6 +772,11 @@ def _build_robot_observation_provider(
                 import cv2 as _cv2  # noqa: PLC0415
                 import numpy as _np  # noqa: PLC0415
 
+                # Snapshot the gate state at the start of the call: the
+                # camera info and startup-state warnings are meant to fire
+                # exactly once (operator sanity check), so gate them on
+                # the *previous* value rather than the post-loop value.
+                first_call = not _resize_logged["done"]
                 for cam_key, (target_h, target_w) in target_image_shapes.items():
                     img = raw.get(cam_key)
                     if img is None or not isinstance(img, _np.ndarray):
@@ -779,7 +784,7 @@ def _build_robot_observation_provider(
                     if img.ndim != 3:
                         continue
                     cur_h, cur_w = img.shape[:2]
-                    if not _resize_logged["done"]:
+                    if first_call:
                         logger.warning(
                             "camera %s: live=%dx%d, training=%dx%d (resize=%s)",
                             cam_key,
@@ -793,13 +798,14 @@ def _build_robot_observation_provider(
                         continue
                     raw[cam_key] = _cv2.resize(img, (target_w, target_h), interpolation=_cv2.INTER_AREA)
                 _resize_logged["done"] = True
-                # Also print the state vector once so the operator
-                # can eyeball it against the dataset's stats. State
-                # OOD is a real failure mode for VLAs — the prefix
-                # carries state via the projection layer, and a
-                # neutral home pose can easily sit a couple σ off
-                # the supervised support region.
-                if "observation.state" in (ds_features or {}):
+                # Print the state vector once so the operator can eyeball
+                # it against the dataset's stats. State OOD is a real
+                # failure mode for VLAs — the prefix carries state via
+                # the projection layer, and a neutral home pose can
+                # easily sit a couple σ off the supervised support
+                # region. Gated on ``first_call`` so this doesn't spam
+                # every observation tick.
+                if first_call and "observation.state" in (ds_features or {}):
                     state_names = ds_features["observation.state"].get("names") or []
                     state_vals = [raw.get(n) for n in state_names]
                     logger.warning(