fix(smolvla2): only regenerate chunk when queue is fully drained

The previous refresh threshold (queue > chunk_size // 2) made each
new chunk *telescope* past the previous one: at queue=25, we kicked
off a new chunk forward from the current observation, but by the
time the new chunk's first action was actually dispatched, the
robot had executed the remaining 25 actions of the previous chunk
— so the new chunk was planned from an observation 25+ steps stale.

Canonical sense → think → act loop: execute the full chunk, then
re-observe and replan. Refresh only when the queue is empty. Every
step of every chunk still gets dispatched to the robot (no
behaviour change there), but each chunk is now planned from an
observation that's at most one chunk's worth of dispatch latency
old, not "previous chunk's worth of stale state on top of that".

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

examples/training/smolvla2_hirobot.slurm

@@ -0,0 +1,84 @@
+#!/bin/bash
+#SBATCH --job-name=smolvla2-hirobot
+#SBATCH --partition=hopper-prod
+#SBATCH --qos=high
+#SBATCH --time=48:00:00
+#SBATCH --ntasks=1
+#SBATCH --gpus-per-task=8
+
+# SmolVLA2 training on an annotated dataset, with image augmentation
+# and per-component prompt dropout enabled — the two regularisers
+# that move the model away from the "text_loss=6e-6 memorised one
+# epoch worth of frames" failure mode toward "learns concepts, not
+# pixels".
+#
+# What the regularisers do:
+#
+#   * --dataset.image_transforms.enable=true: applies torchvision
+#     v2 ColorJitter (brightness/contrast/saturation/hue),
+#     SharpnessJitter and RandomAffine per frame at training time.
+#     Set max_num_transforms to control how many are sampled per
+#     frame; defaults to 3 of the 6.
+#   * --policy.plan_dropout_prob / memory / subtask: at training,
+#     randomly drop the context messages that carry the named
+#     binding so the model is forced to handle missing/stale context.
+#     Mirrors Pi0.7's prompt-component dropout (§V.E).
+#
+# Expected effect: text_loss plateaus higher (~0.5-2.0 instead of
+# ~1e-5) and the model handles slight prompt/scene drift at
+# inference instead of collapsing to memorised fragments.
+
+set -euo pipefail
+
+cd "${LEROBOT_ROOT:-$HOME/lerobot}"
+
+export PATH="$HOME/miniconda3/bin:$HOME/.local/bin:$PATH"
+export LD_LIBRARY_PATH="$HOME/miniconda3/lib:${LD_LIBRARY_PATH:-}"
+export NCCL_TIMEOUT="${NCCL_TIMEOUT:-1800}"
+export HF_HUB_DOWNLOAD_TIMEOUT="${HF_HUB_DOWNLOAD_TIMEOUT:-120}"
+export WANDB_INIT_TIMEOUT="${WANDB_INIT_TIMEOUT:-300}"
+
+DATASET="${DATASET:-pepijn223/super_poulain_full_tool3}"
+POLICY_REPO_ID="${POLICY_REPO_ID:-pepijn223/smolvla2_hirobot_super_poulain_tool4}"
+JOB_NAME="${JOB_NAME:-smolvla2-hirobot-super-poulain-tool4}"
+NUM_PROCESSES="${NUM_PROCESSES:-8}"
+BATCH_SIZE="${BATCH_SIZE:-32}"
+STEPS="${STEPS:-10000}"
+RUN_ID="${SLURM_JOB_ID:-$(date +%Y%m%d_%H%M%S)}"
+OUTPUT_DIR="${OUTPUT_DIR:-/fsx/pepijn/outputs/train/smolvla2_hirobot_${RUN_ID}}"
+
+echo "Training smolvla2 on $DATASET"
+echo "  GPUs:         $NUM_PROCESSES"
+echo "  batch:        $BATCH_SIZE / GPU (global=$((NUM_PROCESSES * BATCH_SIZE)))"
+echo "  steps:        $STEPS"
+echo "  output:       $OUTPUT_DIR"
+echo "  augmentation: image_transforms ON, prompt dropout {plan:0.15 memory:0.15 subtask:0.20}"
+
+accelerate launch --multi_gpu --num_processes="$NUM_PROCESSES" \
+    -m lerobot.scripts.lerobot_train \
+    --policy.type=smolvla2 \
+    --policy.recipe_path=recipes/smolvla2_hirobot.yaml \
+    --dataset.repo_id="$DATASET" \
+    --dataset.revision=main \
+    --dataset.video_backend=pyav \
+    --dataset.image_transforms.enable=true \
+    --dataset.image_transforms.max_num_transforms=3 \
+    --dataset.image_transforms.random_order=true \
+    --policy.plan_dropout_prob=0.15 \
+    --policy.memory_dropout_prob=0.15 \
+    --policy.subtask_dropout_prob=0.20 \
+    --output_dir="$OUTPUT_DIR" \
+    --job_name="$JOB_NAME" \
+    --policy.repo_id="$POLICY_REPO_ID" \
+    --policy.compile_model=false \
+    --policy.device=cuda \
+    --policy.tokenizer_max_length=512 \
+    --steps="$STEPS" \
+    --policy.scheduler_decay_steps="$STEPS" \
+    --batch_size="$BATCH_SIZE" \
+    --wandb.enable=true \
+    --wandb.disable_artifact=true \
+    --wandb.project=hirobot \
+    --log_freq=100 \
+    --save_freq=1000 \
+    --num_workers=0

src/lerobot/policies/smolvla2/chat_processor_smolvla2.py

@@ -70,6 +70,22 @@ class SmolVLA2ChatTokenizerStep(ProcessorStep):
     padding: str = "longest"
     padding_side: str = "right"
     tools: list[dict[str, Any]] | None = None
+    # --- Per-component prompt dropout (Pi0.7 §V.E, plan follow-up
+    # ``feat/pi05-prompt-dropout``). At training, drop non-target
+    # messages whose content was substituted from the named recipe
+    # binding with the given probability. Forces the model to handle
+    # missing context at inference — directly attacks the memorisation
+    # collapse where ``current_subtask=""`` puts the prompt OOD. All
+    # default to 0.0 (no dropout) so behaviour is identical until
+    # explicitly opted in via the training config.
+    plan_dropout_prob: float = 0.0
+    memory_dropout_prob: float = 0.0
+    subtask_dropout_prob: float = 0.0
+    interjection_dropout_prob: float = 0.0
+    # Optional seed for the per-sample RNG. ``None`` ⇒ use
+    # ``sample_idx`` derived from the transition (when present), so
+    # dropout is reproducible across runs but varies per sample.
+    dropout_seed: int | None = None
 
     def __post_init__(self) -> None:
         # Lazy: don't load the tokenizer until the step actually runs,
@@ -101,19 +117,38 @@ class SmolVLA2ChatTokenizerStep(ProcessorStep):
 
         tokenizer = self._get_tokenizer()
 
+        # Pull a sample_idx for the dropout RNG. ``index`` is the
+        # canonical per-frame key on ``LeRobotDataset`` samples and
+        # flows through into ``COMPLEMENTARY_DATA`` unchanged. When
+        # absent (e.g. inference) we fall back to 0 which is harmless
+        # because the dropout probs are also 0 at inference time.
+        sample_idx_raw = comp.get("index")
+        if hasattr(sample_idx_raw, "item"):
+            try:
+                sample_idx_raw = sample_idx_raw.item()
+            except Exception:  # noqa: BLE001
+                pass
+
         if _is_batched_messages(messages):
+            indices_iter = (
+                sample_idx_raw
+                if isinstance(sample_idx_raw, (list, tuple))
+                else [sample_idx_raw] * len(messages)
+            )
             encoded = [
                 self._encode_messages(
                     tokenizer,
                     msg,
                     list(streams),
-                    sorted(int(i) for i in indices),
+                    sorted(int(i) for i in tgt_indices),
+                    sample_idx=int(s_idx) if s_idx is not None else None,
                 )
-                for msg, streams, indices in zip(
+                for msg, streams, tgt_indices, s_idx in zip(
                     messages,
                     comp.get("message_streams") or [[] for _ in messages],
                     comp.get("target_message_indices") or [[] for _ in messages],
-                    strict=True,
+                    indices_iter,
+                    strict=False,
                 )
             ]
         else:
@@ -123,6 +158,7 @@ class SmolVLA2ChatTokenizerStep(ProcessorStep):
                     messages,
                     list(comp.get("message_streams") or []),
                     sorted(int(i) for i in (comp.get("target_message_indices") or [])),
+                    sample_idx=int(sample_idx_raw) if sample_idx_raw is not None else None,
                 )
             ]
 
@@ -190,7 +226,15 @@ class SmolVLA2ChatTokenizerStep(ProcessorStep):
         messages: list[dict[str, Any]],
         message_streams: list[str | None],
         target_indices: list[int],
+        sample_idx: int | None = None,
     ) -> tuple[list[int], list[int], bool]:
+        # Apply per-component prompt dropout *before* tokenisation, so
+        # the dropped messages don't contribute tokens or label-mask
+        # positions at all. Re-maps ``target_indices`` to account for
+        # removed messages.
+        messages, target_indices = self._apply_prompt_dropout(
+            messages, target_indices, sample_idx
+        )
         text_messages = [_strip_lerobot_blocks(m) for m in messages]
 
         full_ids = tokenizer.apply_chat_template(
@@ -231,6 +275,62 @@ class SmolVLA2ChatTokenizerStep(ProcessorStep):
         )
         return [int(i) for i in full_ids], labels, predict_actions
 
+    def _apply_prompt_dropout(
+        self,
+        messages: list[dict[str, Any]],
+        target_indices: list[int],
+        sample_idx: int | None,
+    ) -> tuple[list[dict[str, Any]], list[int]]:
+        """Probabilistically drop non-target context messages.
+
+        Heuristic content sniffing — matches the prefix strings that
+        ``smolvla2_hirobot.yaml``'s recipes use when injecting plan /
+        memory / subtask / interjection content. Anything else is
+        kept unchanged. Target messages are never dropped (we still
+        need their tokens for supervision).
+
+        Returns ``(new_messages, new_target_indices)`` where the
+        indices are re-mapped to point at the same target turns in
+        the trimmed list.
+        """
+        probs = {
+            "plan": float(self.plan_dropout_prob or 0.0),
+            "memory": float(self.memory_dropout_prob or 0.0),
+            "subtask": float(self.subtask_dropout_prob or 0.0),
+            "interjection": float(self.interjection_dropout_prob or 0.0),
+        }
+        if not any(p > 0.0 for p in probs.values()):
+            return messages, target_indices
+
+        # Deterministic per-sample RNG so dropout is reproducible
+        # across runs (matters for debugging / repro) but varies
+        # frame-to-frame.
+        import random  # noqa: PLC0415
+
+        seed_int = self.dropout_seed if self.dropout_seed is not None else (sample_idx or 0)
+        rng = random.Random(int(seed_int) & 0xFFFFFFFF)
+
+        target_set = set(target_indices)
+        keep_flags: list[bool] = []
+        for i, msg in enumerate(messages):
+            if i in target_set:
+                keep_flags.append(True)
+                continue
+            kind = _classify_message_for_dropout(msg)
+            if kind and rng.random() < probs.get(kind, 0.0):
+                keep_flags.append(False)
+            else:
+                keep_flags.append(True)
+
+        new_messages = [m for m, keep in zip(messages, keep_flags) if keep]
+        # Re-map target_indices: each old index drops by the count of
+        # falsy flags before it.
+        new_target_indices: list[int] = []
+        for old_idx in target_indices:
+            dropped_before = sum(1 for k in keep_flags[:old_idx] if not k)
+            new_target_indices.append(old_idx - dropped_before)
+        return new_messages, sorted(new_target_indices)
+
     def transform_features(
         self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
     ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
@@ -296,6 +396,39 @@ def _is_batched_messages(messages: Any) -> bool:
     return isinstance(messages, list) and bool(messages) and isinstance(messages[0], list)
 
 
+def _classify_message_for_dropout(message: dict[str, Any]) -> str | None:
+    """Best-effort classification of which recipe binding contributed
+    to this message, used for per-component dropout.
+
+    The canonical recipe authors plan/memory/subtask injections with
+    distinctive prefix strings in the rendered content. Matching on
+    those prefixes is brittle if a future recipe author uses
+    different wording — but it's also localised to one place and
+    only affects the dropout fraction (never the actual semantics).
+    Returns ``None`` for messages we don't recognise; those are
+    always kept.
+    """
+    content = message.get("content")
+    if isinstance(content, list):
+        text_parts: list[str] = []
+        for block in content:
+            if isinstance(block, dict) and block.get("type") == "text":
+                t = block.get("text")
+                if isinstance(t, str):
+                    text_parts.append(t)
+        content = "\n".join(text_parts)
+    if not isinstance(content, str):
+        return None
+    head = content.lstrip().lower()
+    if head.startswith("plan:") or head.startswith("previous plan"):
+        return "plan"
+    if head.startswith("memory:") or head.startswith("previous memory"):
+        return "memory"
+    if head.startswith("current subtask") or head.startswith("completed subtask"):
+        return "subtask"
+    return None
+
+
 def _as_token_ids(value: Any) -> list[int]:
     if isinstance(value, dict) or (hasattr(value, "keys") and "input_ids" in value.keys()):
         value = value["input_ids"]

src/lerobot/policies/smolvla2/configuration_smolvla2.py

@@ -84,6 +84,24 @@ class SmolVLA2Config(SmolVLAConfig):
     effectively reduces SmolVLA2 back to SmolVLA's flow-only training,
     which is occasionally useful for ablations."""
 
+    # Per-component prompt dropout (Pi0.7 §V.E) ---------------------------
+    # At training, randomly drop non-target context messages whose
+    # content was substituted from the named recipe binding. Forces
+    # the model to handle missing context — directly attacks the
+    # memorisation collapse where a stale or missing plan/memory at
+    # inference puts the prompt out-of-distribution and the LM head
+    # falls back to dominant-mode fragments. All default to 0.0 so
+    # behaviour is identical until explicitly enabled.
+    plan_dropout_prob: float = 0.0
+    """Drop messages whose content starts with ``Plan:`` or ``Previous plan``
+    with this probability per sample."""
+    memory_dropout_prob: float = 0.0
+    """Drop messages whose content starts with ``Memory:`` or ``Previous memory``
+    with this probability per sample."""
+    subtask_dropout_prob: float = 0.0
+    """Drop messages whose content starts with ``Current subtask`` or
+    ``Completed subtask`` with this probability per sample."""
+
     def __post_init__(self) -> None:
         super().__post_init__()
         # Backbone needs gradients flowing through its text path when the

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

@@ -92,17 +92,27 @@ class LowLevelForward(InferenceStep):
         if self.policy is None or self.observation_provider is None:
             return None
         if not state.get("task"):
-            # No task yet → nothing useful to condition on.
             return None
+
+        # SmolVLA produces *action chunks* (typically 50 steps via
+        # flow-matching). Every step gets dispatched to the robot;
+        # popping one per dispatch tick is essentially free. Only
+        # generate a new chunk once the previous one has fully
+        # drained — this is the canonical "sense → think → act"
+        # loop. Refreshing while a chunk is still queued causes the
+        # new chunk to "telescope" past the old one (planned from an
+        # observation that's already 25+ steps stale by the time it
+        # starts dispatching).
+        queue = state.setdefault("action_queue", [])
+        if len(queue) > 0:
+            return None
+
         observation = self.observation_provider()
         if observation is None:
             return None
-        # SmolVLA's ``select_action`` expects the full preprocessed
-        # batch, including ``OBS_LANGUAGE_TOKENS`` /
-        # ``OBS_LANGUAGE_ATTENTION_MASK``. The observation provider
-        # only returns image / state features (the runtime drives
-        # messages itself), so build a low-level prompt from current
-        # runtime state and tokenize it inline.
+
+        # Same prompt construction as before — task + plan + memory,
+        # optional current subtask — then merge into the obs batch.
         ctx = _control_context_messages(state)
         if state.get("current_subtask"):
             ctx = ctx + [{"role": "assistant", "content": state["current_subtask"]}]
@@ -115,16 +125,38 @@ class LowLevelForward(InferenceStep):
         observation = dict(observation)
         observation[OBS_LANGUAGE_TOKENS] = text_batch["lang_tokens"]
         observation[OBS_LANGUAGE_ATTENTION_MASK] = text_batch["lang_masks"]
+
         try:
-            action = self.policy.select_action(observation)
+            # ``predict_action_chunk`` returns the *full* chunk shape
+            # ``(batch, n_action_steps, action_dim)``. Enqueue every
+            # step so DispatchAction at ctrl_hz can drain them
+            # smoothly until the next refresh.
+            chunk = self.policy.predict_action_chunk(observation)
         except Exception as exc:  # noqa: BLE001
-            logger.warning("select_action failed: %s", exc, exc_info=logger.isEnabledFor(logging.DEBUG))
-            push_log(state, f"  [warn] select_action failed: {type(exc).__name__}: {exc}")
+            logger.warning(
+                "predict_action_chunk failed: %s",
+                exc,
+                exc_info=logger.isEnabledFor(logging.DEBUG),
+            )
+            push_log(
+                state,
+                f"  [warn] predict_action_chunk failed: "
+                f"{type(exc).__name__}: {exc}",
+            )
             return None
-        # SmolVLA returns a single action; if the underlying policy
-        # streams chunks, split per-step here. For v1 we just enqueue
-        # the result.
-        state.setdefault("action_queue", []).append(action)
+
+        # ``chunk`` shape: ``(batch, n_action_steps, action_dim)``. Push
+        # each step as a ``(1, action_dim)`` tensor so the existing
+        # action executor's batch-squeeze logic works unchanged.
+        if chunk.ndim == 3:
+            chunk_iter = chunk[0]  # ``(n_action_steps, action_dim)``
+        elif chunk.ndim == 2:
+            chunk_iter = chunk
+        else:
+            chunk_iter = chunk.unsqueeze(0)
+        for step in chunk_iter:
+            queue.append(step.unsqueeze(0))
+        state["last_chunk_size"] = int(chunk_iter.shape[0])
         return None
 
 
@@ -147,6 +179,11 @@ class DispatchAction(InferenceStep):
         action = queue.popleft() if hasattr(queue, "popleft") else queue.pop(0)
         if self.robot_executor is not None:
             self.robot_executor(action)
+        # Track lifetime dispatch count so the REPL panel can show
+        # whether the action loop is actually doing useful work, even
+        # while the text head produces gibberish (the typical real-
+        # robot failure mode for a memorised model).
+        state["actions_dispatched"] = state.get("actions_dispatched", 0) + 1
         return None
 
 
@@ -170,7 +207,20 @@ def _build_text_batch(policy: Any, prompt_messages: list[dict[str, Any]]) -> dic
     if tokenizer.pad_token_id is None and tokenizer.eos_token_id is not None:
         tokenizer.pad_token = tokenizer.eos_token
 
-    text_messages = [_strip_recipe_keys(m) for m in prompt_messages]
+    # Reuse the *exact* normaliser that the training-time chat
+    # tokenizer step uses (``_strip_lerobot_blocks``). It handles all
+    # the cases the SmolVLM chat template expects:
+    #   * ``content: list[block]`` → keep text blocks, drop images
+    #   * ``content: None`` → ``[{type: text, text: ""}]``
+    #   * ``content: str`` / anything else → ``[{type: text, text: str(content)}]``
+    # Doing it any other way creates a training/inference mismatch in
+    # exactly the prompt shape the model was supervised on. Also
+    # strips ``stream`` / ``target`` recipe metadata.
+    from lerobot.policies.smolvla2.chat_processor_smolvla2 import (  # noqa: PLC0415
+        _strip_lerobot_blocks,
+    )
+
+    text_messages = [_strip_lerobot_blocks(m) for m in prompt_messages]
     encoded = tokenizer.apply_chat_template(
         text_messages,
         add_generation_prompt=True,
@@ -272,15 +322,25 @@ class HighLevelSubtaskFwd(InferenceStep):
             self.policy, ctx, observation=observation, state=state, label="subtask gen"
         )
         if msg and _looks_like_gibberish(msg):
-            push_log(state, f"  [info] subtask gen rejected (gibberish): {msg[:60]!r}")
+            # Bump a counter so the operator can see the model is
+            # struggling without spamming the log every tick. A first
+            # rejection still logs once so the failure is visible.
+            count = state.get("subtask_gibberish_count", 0) + 1
+            state["subtask_gibberish_count"] = count
+            if count == 1 or count % 30 == 0:
+                push_log(
+                    state,
+                    f"  [info] subtask gen rejected (gibberish ×{count}): {msg[:60]!r}",
+                )
             return None
         if msg:
             changed = set_if_changed(state, "current_subtask", msg, label="subtask")
             if changed:
                 # Subtask change is a downstream trigger.
                 state.setdefault("events_this_tick", []).append("subtask_change")
-        else:
-            push_log(state, "  [info] subtask gen produced no text this tick")
+        # Silently skip empty completions — common when the model
+        # warms up or generates only EOS; logging it every tick at
+        # ctrl_hz is just noise.
         return None
 
 
@@ -344,7 +404,9 @@ class UserInterjectionFwd(InferenceStep):
             self.policy, ctx, observation=observation, state=state, label="plan/say gen"
         )
         if not out:
-            push_log(state, "  [info] plan/say gen produced no text this tick")
+            # Don't log every empty completion — happens repeatedly on
+            # MPS during warm-up and floods the panel. The user can
+            # re-trigger by typing again.
             return None
         if _looks_like_gibberish(out):
             push_log(state, f"  [info] plan/say gen rejected (gibberish): {out[:60]!r}")
@@ -449,20 +511,22 @@ class DispatchToolCalls(InferenceStep):
 def _looks_like_gibberish(text: str) -> bool:
     """Heuristically detect generation that's clearly off the rails.
 
-    Memorised models can collapse to dominant-mode outputs (often the
-    JSON-token salad ``":":":":...`` from VQA training) when the prompt
-    drifts even slightly from training distribution. If we accept those
-    as new state, they pollute the next tick's prompt and cascade into
-    worse outputs. Reject anything that looks pathological:
+    Memorised models can collapse to dominant-mode outputs when the
+    prompt drifts even slightly from training distribution. Reject:
 
     * empty / whitespace-only
-    * mostly punctuation (``"``, ``:``, ``,``)
+    * too few alphabetic characters (mostly punctuation)
     * a single character repeated past the threshold
     * starts with ``":"`` and contains no letters
+    * too few unique tokens — e.g. ``"the"``, ``"the the the"``,
+      ``"Ass\\n::\\nthe"`` (the collapse seen on real-robot frames
+      where the model emits one or two memorised tokens repeatedly)
+    * chat-template fragment leakage (``Assistant:``, ``User:``,
+      ``Ass\\n``)
 
-    The thresholds are intentionally lenient — a real subtask like
-    ``"close the gripper"`` has ~70%+ alpha characters, while gibberish
-    like ``":":":"`` has ~0%.
+    Real subtasks look like ``"close the gripper to grasp the blue
+    cube"`` — multiple unique alphabetic tokens, no role-marker
+    fragments. Anything materially shorter than that is rejected.
     """
     if not text or not text.strip():
         return True
@@ -472,9 +536,22 @@ def _looks_like_gibberish(text: str) -> bool:
         return True
     if stripped.startswith('":') and stripped.count('"') > stripped.count(" "):
         return True
-    # Single repeating char: e.g. ``""""""``
+    # Single repeating char: e.g. ``""""""``.
     if len(set(stripped)) <= 2 and len(stripped) > 4:
         return True
+    # Chat-template fragment leakage — the model emits ``Ass``,
+    # ``Assistant:``, ``User:``, often with extra newlines/colons.
+    # Reject if the cleaned text is mostly role-marker shards.
+    cleaned = stripped.replace("\n", " ").replace(":", " ")
+    for marker in ("Assistant", "User", "Ass "):
+        if marker in cleaned and len(cleaned.split()) < 4:
+            return True
+    # Too few unique alphabetic tokens — model stuck on ``the`` or
+    # similar memorised single-token continuations.
+    tokens = [t for t in cleaned.split() if any(c.isalpha() for c in t)]
+    unique_alpha = {t.lower() for t in tokens}
+    if len(unique_alpha) < 3 and len(stripped) < 80:
+        return True
     return False
 
 

src/lerobot/policies/smolvla2/processor_smolvla2.py

@@ -83,6 +83,9 @@ def make_smolvla2_pre_post_processors(
             tokenizer_name=config.vlm_model_name,
             max_length=config.tokenizer_max_length,
             padding=config.pad_language_to,
+            plan_dropout_prob=getattr(config, "plan_dropout_prob", 0.0),
+            memory_dropout_prob=getattr(config, "memory_dropout_prob", 0.0),
+            subtask_dropout_prob=getattr(config, "subtask_dropout_prob", 0.0),
         ),
         DeviceProcessorStep(device=config.device),
         NormalizerProcessorStep(

src/lerobot/scripts/lerobot_smolvla2_runtime.py

@@ -180,14 +180,19 @@ def _parse_args(argv: list[str] | None = None) -> argparse.Namespace:
         ),
     )
     p.add_argument(
-        "--max_action_norm",
-        dest="max_action_norm",
-        type=float,
+        "--robot.max_relative_target",
+        dest="robot_max_relative_target",
+        type=str,
         default=None,
         help=(
-            "Safety clip: reject any individual action whose L2 norm "
-            "exceeds this value. Default ``None`` = no clipping. Useful "
-            "as a kill-switch when bringing up a new robot/task pair."
+            "Safety clip on per-motor relative motion, passed through to "
+            "``RobotConfig.max_relative_target``. Accepts either a float "
+            "(applied to every motor — e.g. ``5.0`` degrees) or a JSON "
+            "object mapping motor names to caps "
+            "(e.g. ``'{\"shoulder_pan\": 5, \"gripper\": 30}'``). The "
+            "robot driver clips each commanded position relative to the "
+            "current measured position before sending — same kill-switch "
+            "``lerobot-record`` uses. Default ``None`` = no clipping."
         ),
     )
     p.add_argument(
@@ -213,7 +218,16 @@ def _parse_args(argv: list[str] | None = None) -> argparse.Namespace:
         help="Pocket-tts voice name (or path to a .wav for cloning).",
     )
     p.add_argument(
-        "--chunk_hz", type=float, default=4.0, help="Action-chunk generation rate."
+        "--chunk_hz",
+        type=float,
+        default=1.0,
+        help=(
+            "Action-chunk generation rate (Hz). Default ``1.0`` — one "
+            "new chunk per second. Lower = less inference cost / "
+            "smoother behaviour but longer reaction time to changes. "
+            "Higher = fresher actions / more inference cost; cap at "
+            "~1/(forward-pass latency)."
+        ),
     )
     p.add_argument(
         "--ctrl_hz", type=float, default=50.0, help="Action dispatch rate."
@@ -427,21 +441,52 @@ def _build_robot(
     robot_port: str | None,
     robot_id: str | None,
     robot_cameras_json: str | None,
+    robot_max_relative_target: str | None,
 ):
     """Build and connect a robot from CLI args.
 
     Mirrors how ``lerobot-record`` builds a robot but takes the args
     flat from argparse instead of through draccus, so the runtime
-    keeps its plain ``--key=value`` CLI surface.
+    keeps its plain ``--key=value`` CLI surface. ``max_relative_target``
+    is passed through to the RobotConfig — the driver itself clips each
+    commanded joint position relative to the current measured one
+    before issuing it on the bus.
     """
+    import importlib  # noqa: PLC0415
     import json  # noqa: PLC0415
+    import pkgutil  # noqa: PLC0415
 
+    import lerobot.robots as _robots_pkg  # noqa: PLC0415
     from lerobot.robots import (  # noqa: PLC0415
         RobotConfig,
         make_robot_from_config,
     )
 
-    cls = RobotConfig.get_choice_class(robot_type)
+    # ``RobotConfig._choice_registry`` is populated lazily — each robot's
+    # ``config_<name>.py`` calls ``@RobotConfig.register_subclass`` at
+    # import time. ``lerobot.robots/__init__.py`` doesn't import the
+    # individual robot packages, so ``get_choice_class(robot_type)``
+    # raises ``KeyError`` until at least one robot module has been
+    # imported. Mirror what ``make_robot_from_config`` does internally:
+    # walk the robots package's submodules and import each so the
+    # decorator side-effect runs. Slow only on the first call (~200ms
+    # for ~10 dataclass modules); negligible for an autonomous run that
+    # then loops at ctrl_hz for minutes.
+    for _modinfo in pkgutil.iter_modules(_robots_pkg.__path__):
+        if _modinfo.name.startswith("_"):
+            continue
+        try:
+            importlib.import_module(f"lerobot.robots.{_modinfo.name}")
+        except Exception as exc:  # noqa: BLE001
+            logger.debug("could not import lerobot.robots.%s: %s", _modinfo.name, exc)
+
+    try:
+        cls = RobotConfig.get_choice_class(robot_type)
+    except KeyError as exc:
+        available = sorted(RobotConfig._choice_registry.keys())
+        raise ValueError(
+            f"Unknown robot type {robot_type!r}. Available choices: {available}"
+        ) from exc
     kwargs: dict[str, Any] = {}
     if robot_port:
         kwargs["port"] = robot_port
@@ -449,11 +494,67 @@ def _build_robot(
         kwargs["id"] = robot_id
     if robot_cameras_json:
         try:
-            kwargs["cameras"] = json.loads(robot_cameras_json)
+            cameras_raw = json.loads(robot_cameras_json)
         except json.JSONDecodeError as exc:
             raise ValueError(
                 f"--robot.cameras must be a JSON object, got {robot_cameras_json!r}: {exc}"
             ) from exc
+        # ``RobotConfig`` expects ``cameras: dict[str, CameraConfig]`` —
+        # each inner value must be an actual ``CameraConfig`` subclass
+        # instance, not a raw dict. Look up the matching subclass via
+        # ``CameraConfig.get_choice_class(<type>)`` (registered by
+        # ``@CameraConfig.register_subclass`` decorators on each camera
+        # backend's config) and instantiate it. Mirror the lazy-import
+        # pattern from above so the registry is populated.
+        import lerobot.cameras as _cameras_pkg  # noqa: PLC0415
+        from lerobot.cameras import CameraConfig  # noqa: PLC0415
+
+        for _modinfo in pkgutil.iter_modules(_cameras_pkg.__path__):
+            if _modinfo.name.startswith("_"):
+                continue
+            try:
+                importlib.import_module(f"lerobot.cameras.{_modinfo.name}")
+            except Exception as exc:  # noqa: BLE001
+                logger.debug("could not import lerobot.cameras.%s: %s", _modinfo.name, exc)
+
+        cameras: dict[str, Any] = {}
+        for cam_name, cam_dict in cameras_raw.items():
+            if not isinstance(cam_dict, dict):
+                raise ValueError(
+                    f"camera {cam_name!r} value must be a dict, got {cam_dict!r}"
+                )
+            cam_dict = dict(cam_dict)  # don't mutate caller's parsed JSON
+            cam_type = cam_dict.pop("type", None)
+            if cam_type is None:
+                raise ValueError(
+                    f"camera {cam_name!r} is missing a 'type' field "
+                    f"(e.g. 'opencv', 'intelrealsense')"
+                )
+            try:
+                cam_cls = CameraConfig.get_choice_class(cam_type)
+            except KeyError as exc:
+                available = sorted(CameraConfig._choice_registry.keys())
+                raise ValueError(
+                    f"camera {cam_name!r}: unknown type {cam_type!r}. "
+                    f"Available choices: {available}"
+                ) from exc
+            cameras[cam_name] = cam_cls(**cam_dict)
+        kwargs["cameras"] = cameras
+    if robot_max_relative_target:
+        # Accept either a bare float (uniform cap) or a JSON object
+        # (per-motor cap). Matches ``RobotConfig.max_relative_target``'s
+        # ``float | dict[str, float] | None`` shape.
+        s = robot_max_relative_target.strip()
+        try:
+            if s.startswith("{"):
+                kwargs["max_relative_target"] = json.loads(s)
+            else:
+                kwargs["max_relative_target"] = float(s)
+        except (json.JSONDecodeError, ValueError) as exc:
+            raise ValueError(
+                f"--robot.max_relative_target must be a float or JSON dict, "
+                f"got {robot_max_relative_target!r}: {exc}"
+            ) from exc
     cfg = cls(**kwargs)
     robot = make_robot_from_config(cfg)
     robot.connect()
@@ -466,17 +567,33 @@ def _build_robot_observation_provider(
     preprocessor: Any,
     device: str,
     task: str | None,
+    ds_features: dict[str, Any] | None,
 ) -> Callable[[], dict | None]:
     """Closure that reads from the robot, runs the policy preprocessor.
 
-    Each call: ``robot.get_observation()`` → wrap as a flat sample dict
-    → drop language columns (the runtime drives messages itself) →
-    preprocessor (rename, batch dim, normalise, device-place) → return
-    the observation batch ready for ``policy.select_action`` and
-    ``policy.select_message``.
+    Each call: ``robot.get_observation()`` (raw per-joint + per-camera
+    dict, possibly with scalar floats) → ``build_inference_frame``
+    (extract the keys the dataset declared, reshape per-joint floats
+    into a single ``observation.state`` vector, prefix camera keys
+    with ``observation.images.``, convert to tensors with batch dim
+    on device) → wrap in an ``EnvTransition`` (the preprocessor
+    pipeline is transition-shaped, keyed by ``TransitionKey``) →
+    preprocessor (rename, normalise) → unwrap and return the flat
+    observation batch ``policy.select_action`` / ``policy.select_message``
+    consume.
     """
     import torch  # noqa: PLC0415
 
+    from lerobot.policies.utils import (  # noqa: PLC0415
+        build_inference_frame,
+        prepare_observation_for_inference,
+    )
+
+    torch_device = torch.device(device) if isinstance(device, str) else device
+    robot_type = getattr(robot, "robot_type", None) or getattr(
+        getattr(robot, "config", None), "type", None
+    )
+
     def _provider() -> dict | None:
         try:
             raw = robot.get_observation()
@@ -484,30 +601,58 @@ def _build_robot_observation_provider(
             logger.warning("robot.get_observation failed: %s", exc)
             return None
 
-        sample: dict[str, Any] = dict(raw)
-        if task:
-            sample.setdefault("task", task)
-        # The render step expects either both language columns or
-        # neither — runtime supplies messages itself, so make sure
-        # nothing leaks through.
+        # Strip language-column leakage just in case (the runtime
+        # supplies messages itself).
         for k in ("language_persistent", "language_events"):
-            sample.pop(k, None)
+            raw.pop(k, None)
+
+        try:
+            if ds_features:
+                # Use the dataset's feature schema to pick the right
+                # raw keys and fold per-joint scalars into a single
+                # ``observation.state`` tensor. Then tensor-ise +
+                # device-place + add batch dim.
+                obs_tensors = build_inference_frame(
+                    raw, torch_device, ds_features=ds_features,
+                    task=task, robot_type=robot_type,
+                )
+            else:
+                # No dataset features available — fall back to the
+                # generic numpy-only path; only works when the robot
+                # already returns dataset-shaped keys.
+                obs_tensors = prepare_observation_for_inference(
+                    raw, torch_device, task=task, robot_type=robot_type,
+                )
+        except Exception as exc:  # noqa: BLE001
+            logger.warning("observation prep failed: %s", exc)
+            return None
 
         if preprocessor is not None:
+            # ``PolicyProcessorPipeline`` defaults its ``to_transition``
+            # to ``batch_to_transition``, which expects a *flat batch
+            # dict* keyed by ``observation.*`` / ``action`` / etc., and
+            # wraps it into an ``EnvTransition`` itself. Pre-wrapping
+            # here would just have ``batch_to_transition`` look for
+            # ``observation.*`` keys at top level, find none (they'd
+            # be nested under ``TransitionKey.OBSERVATION``), and
+            # produce an empty observation → ``ObservationProcessorStep``
+            # bails. Pass the flat dict straight in; ``to_output``
+            # gives us a flat dict back.
             try:
-                sample = preprocessor(sample)
+                processed = preprocessor(obs_tensors)
             except Exception as exc:  # noqa: BLE001
                 logger.warning("preprocessor failed on robot observation: %s", exc)
                 return None
+            obs_tensors = processed if isinstance(processed, dict) else {}
 
         observation = {
             k: v
-            for k, v in sample.items()
+            for k, v in obs_tensors.items()
             if isinstance(k, str) and k.startswith("observation.")
         }
         for k, v in list(observation.items()):
             if isinstance(v, torch.Tensor):
-                observation[k] = v.to(device)
+                observation[k] = v.to(torch_device)
         return observation
 
     return _provider
@@ -518,15 +663,15 @@ def _build_robot_action_executor(
     robot,
     postprocessor: Any,
     ds_meta: Any,
-    max_action_norm: float | None,
 ) -> Callable[[Any], None]:
     """Closure that postprocesses an action and dispatches to the robot.
 
     Mirrors ``lerobot-record``'s ``predict_action`` tail: postprocess
     (denormalise) → ``make_robot_action`` (tensor → ``{joint: value}``
-    dict) → ``robot.send_action(...)``. Optional safety clip on the
-    action's L2 norm acts as a kill switch when bringing up a new
-    robot/task pair.
+    dict) → ``robot.send_action(...)``. Safety clipping happens *inside*
+    ``robot.send_action`` via the driver's ``max_relative_target``
+    cap (passed in at ``RobotConfig`` construction time) — same place
+    ``lerobot-record`` enforces it.
     """
     import torch  # noqa: PLC0415
 
@@ -537,15 +682,6 @@ def _build_robot_action_executor(
             if postprocessor is not None:
                 action = postprocessor(action)
             if isinstance(action, torch.Tensor):
-                if max_action_norm is not None:
-                    norm = float(action.float().norm().item())
-                    if norm > max_action_norm:
-                        logger.warning(
-                            "action norm %.3f > max_action_norm=%.3f — "
-                            "rejecting tick",
-                            norm, max_action_norm,
-                        )
-                        return
                 if action.ndim > 1 and action.shape[0] == 1:
                     action = action.squeeze(0)
                 action_dict = make_robot_action(action, ds_meta.features)
@@ -601,12 +737,34 @@ def _run_autonomous(
         daemon=True,
     )
     thread.start()
+
+    redraw = _make_state_panel_renderer(runtime, mode_label="autonomous")
+    redraw()
     print(
-        "[smolvla2] autonomous loop running. Type interjections / "
-        "questions on stdin (Ctrl+C to stop).",
+        "  [autonomous] type interjections / '?' questions on stdin, "
+        "'stop' or Ctrl+C to quit",
         flush=True,
     )
 
+    # Background panel-redraw thread so state changes from the runtime
+    # loop (subtask refresh, plan update, etc.) are visible without the
+    # user typing anything. 2 Hz is plenty — generation runs at most
+    # ~1 Hz on MPS.
+    _panel_stop = threading.Event()
+
+    def _panel_loop() -> None:
+        while not _panel_stop.is_set():
+            try:
+                redraw()
+            except Exception:  # noqa: BLE001
+                pass
+            _panel_stop.wait(0.5)
+
+    panel_thread = threading.Thread(
+        target=_panel_loop, name="smolvla2-panel-redraw", daemon=True
+    )
+    panel_thread.start()
+
     try:
         while thread.is_alive():
             try:
@@ -630,6 +788,7 @@ def _run_autonomous(
     except KeyboardInterrupt:
         print("\n[smolvla2] interrupt — stopping", flush=True)
     finally:
+        _panel_stop.set()
         runtime.stop()
         # Give the loop a moment to drain.
         for _ in range(10):
@@ -645,6 +804,64 @@ def _run_autonomous(
     return 0
 
 
+def _make_state_panel_renderer(
+    runtime: Any,
+    *,
+    mode_label: str,
+) -> Callable[[list[str] | None], None]:
+    """Return a closure that prints the task/subtask/plan/memory panel.
+
+    Used by both ``_run_repl`` (dry-run, called per user input) and
+    ``_run_autonomous`` (real robot, called on a 2 Hz timer +
+    whenever the user types). Centralises the visual format so the
+    two modes look identical.
+    """
+    from rich.console import Console  # noqa: PLC0415
+
+    console = Console(highlight=False)
+
+    def _redraw(robot_lines: list[str] | None = None) -> None:
+        console.clear()
+        console.rule(f"[bold]SmolVLA2[/] · {mode_label}", style="cyan")
+        st = runtime.state
+        for key, label in (
+            ("task", "task"),
+            ("current_subtask", "subtask"),
+            ("current_plan", "plan"),
+            ("current_memory", "memory"),
+        ):
+            value = st.get(key)
+            if value:
+                console.print(f"  [bold cyan]{label:<8}[/] {value}")
+            else:
+                console.print(f"  [dim]{label:<8} (not set)[/]")
+        queue_len = (
+            len(st["action_queue"])
+            if isinstance(st.get("action_queue"), (list, tuple))
+            or hasattr(st.get("action_queue"), "__len__")
+            else 0
+        )
+        pending = len(st.get("tool_calls_pending") or [])
+        dispatched = int(st.get("actions_dispatched") or 0)
+        console.print(
+            f"  [dim]queued actions: {queue_len}    "
+            f"dispatched: {dispatched}    "
+            f"pending tool calls: {pending}[/]"
+        )
+        console.rule(style="cyan")
+        if robot_lines:
+            for line in robot_lines:
+                console.print(f"  [magenta]{line.strip()}[/]")
+            console.print()
+        if not st.get("task"):
+            console.print(
+                "  [dim]Type the task to begin. Lines ending in '?' are VQA, "
+                "anything else is an interjection. Type 'stop' to exit.[/]"
+            )
+
+    return _redraw
+
+
 def _build_tools(no_tts: bool, tts_voice: str) -> dict[str, Any]:
     """Instantiate the tools declared on this dataset/policy."""
     if no_tts:
@@ -745,18 +962,19 @@ def main(argv: list[str] | None = None) -> int:
             robot_port=args.robot_port,
             robot_id=args.robot_id,
             robot_cameras_json=args.robot_cameras,
+            robot_max_relative_target=args.robot_max_relative_target,
         )
         observation_provider = _build_robot_observation_provider(
             robot=robot,
             preprocessor=preprocessor,
             device=str(getattr(policy.config, "device", "cpu")),
             task=args.task,
+            ds_features=ds_meta.features if ds_meta is not None else None,
         )
         robot_executor = _build_robot_action_executor(
             robot=robot,
             postprocessor=postprocessor,
             ds_meta=ds_meta,
-            max_action_norm=args.max_action_norm,
         )
     elif args.dataset_repo_id is not None:
         print(
@@ -839,48 +1057,11 @@ def _run_repl(runtime: Any, *, initial_task: str | None, max_ticks: int | None)
         )
         return 2
 
+    _redraw = _make_state_panel_renderer(runtime, mode_label="dry-run")
+    # Keep a local ``console`` just for the styled input prompt; the
+    # state panel is owned by the shared renderer.
     console = Console(highlight=False)
 
-    def _redraw(robot_lines: list[str] | None = None) -> None:
-        # ANSI clear screen + home cursor. Falls back gracefully on
-        # dumb terminals — they just see scrolled output, which is
-        # fine.
-        console.clear()
-        console.rule("[bold]SmolVLA2[/] · dry-run", style="cyan")
-        st = runtime.state
-        for key, label in (
-            ("task", "task"),
-            ("current_subtask", "subtask"),
-            ("current_plan", "plan"),
-            ("current_memory", "memory"),
-        ):
-            value = st.get(key)
-            if value:
-                console.print(f"  [bold cyan]{label:<8}[/] {value}")
-            else:
-                console.print(f"  [dim]{label:<8} (not set)[/]")
-        queue_len = (
-            len(st["action_queue"])
-            if isinstance(st.get("action_queue"), (list, tuple))
-            or hasattr(st.get("action_queue"), "__len__")
-            else 0
-        )
-        pending = len(st.get("tool_calls_pending") or [])
-        console.print(
-            f"  [dim]queued actions: {queue_len}    pending tool calls: {pending}[/]"
-        )
-        console.rule(style="cyan")
-        if robot_lines:
-            for line in robot_lines:
-                console.print(f"  [magenta]{line.strip()}[/]")
-            console.print()
-        # Help line under the divider when nothing is set yet.
-        if not st.get("task"):
-            console.print(
-                "  [dim]Type the task to begin. Lines ending in '?' are VQA, "
-                "anything else is an interjection. Type 'stop' to exit.[/]"
-            )
-
     last_logs: list[str] = []
     _redraw()
     if initial_task is None: