fix(smolvla2): per-recipe inference prompts to match training shape

The four high-level steps shared one generic
``_control_context_messages`` that jammed task + plan + memory +
completed_subtask into a single user message. The recipes in
``smolvla2_hirobot.yaml`` each have a *specific* multi-message layout
(``memory_update``: ``user(task) → assistant(prev memory) →
user(completed subtask)``; ``high_level_subtask``: ``user(task+plan+
memory) → user(current subtask)``; ``user_interjection_response``:
``user(task) → assistant(prev plan) → user(interjection)``). After
``apply_chat_template`` those layouts produce different prompts than
the runtime's flattened single-user-turn version, and the model fell
back to its dominant training mode (VQA JSON output) — generating
``":":":":":":...`` repetition.

Add four per-recipe prompt builders (``_msgs_for_subtask``,
``_msgs_for_memory``, ``_msgs_for_interjection``, ``_msgs_for_vqa``),
each mirroring its sub-recipe's exact message structure including
the ``if_present`` skips. Wire each high-level step to its matching
builder. Inference prompts now line up with what the model saw in
training, so generation should produce coherent text instead of
repeated tokens.

Generic ``_control_context_messages`` is kept (still used by tests
and the no-recipe fallback path).

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

pyproject.toml

@@ -213,7 +213,7 @@ annotations = [
 # are isolated so adding a new tool doesn't bloat the base install.
 # Currently only `say` (Kyutai pocket-tts; CPU-only, ~100M params).
 tools = [
-    "pocket-tts>=0.1.0,<1.0.0",
+    "pocket-tts>=1.0.0,<3.0.0",
     "scipy>=1.11.0,<2.0.0",  # SayTool.output_dir uses scipy.io.wavfile
 ]
 

src/lerobot/configs/recipes/smolvla2_hirobot.yaml

@@ -62,14 +62,14 @@ blend:
   ask_vqa_top:
     weight: 0.10
     bindings:
-      vqa_query: "emitted_at(t, style=vqa, role=user, camera=observation.images.top)"
-      vqa: "emitted_at(t, style=vqa, role=assistant, camera=observation.images.top)"
+      vqa_query: "emitted_at(t, style=vqa, role=user, camera=observation.images.front)"
+      vqa: "emitted_at(t, style=vqa, role=assistant, camera=observation.images.front)"
     messages:
       - role: user
         stream: high_level
         if_present: vqa_query
         content:
-          - {type: image, feature: observation.images.top}
+          - {type: image, feature: observation.images.front}
           - {type: text, text: "${vqa_query}"}
       - {role: assistant, content: "${vqa}", stream: high_level, target: true, if_present: vqa}
 

src/lerobot/datasets/__init__.py

@@ -34,7 +34,6 @@ from .dataset_tools import (
     remove_feature,
     split_dataset,
 )
-from .factory import make_dataset, resolve_delta_timestamps
 from .image_writer import safe_stop_image_writer
 from .io_utils import load_episodes, write_stats
 from .language import (
@@ -53,6 +52,19 @@ from .streaming_dataset import StreamingLeRobotDataset
 from .utils import DEFAULT_EPISODES_PATH, create_lerobot_dataset_card
 from .video_utils import VideoEncodingManager
 
+
+def make_dataset(*args, **kwargs):
+    from .factory import make_dataset as _make_dataset
+
+    return _make_dataset(*args, **kwargs)
+
+
+def resolve_delta_timestamps(*args, **kwargs):
+    from .factory import resolve_delta_timestamps as _resolve_delta_timestamps
+
+    return _resolve_delta_timestamps(*args, **kwargs)
+
+
 # NOTE: Low-level I/O functions (cast_stats_to_numpy, get_parquet_file_size_in_mb, etc.)
 # and legacy migration constants are intentionally NOT re-exported here.
 # Import directly: ``from lerobot.datasets.io_utils import ...``

src/lerobot/datasets/dataset_reader.py

@@ -126,10 +126,53 @@ class DatasetReader:
     def _load_hf_dataset(self) -> datasets.Dataset:
         """hf_dataset contains all the observations, states, actions, rewards, etc."""
         features = get_hf_features_from_features(self._meta.features)
+        # Datasets annotated with the PR1 language columns may have been
+        # written without registering those columns in ``meta/info.json``
+        # (e.g. they predate ``CODEBASE_VERSION="v3.1"`` and were
+        # back-filled by ``lerobot-annotate``). Probe a single parquet
+        # shard and graft the column features on so the strict
+        # ``Dataset.from_parquet`` cast doesn't fail with
+        # ``column names don't match``.
+        features = self._extend_features_with_language_columns(features)
         hf_dataset = load_nested_dataset(self.root / "data", features=features, episodes=self.episodes)
         hf_dataset.set_transform(hf_transform_to_torch)
         return hf_dataset
 
+    def _extend_features_with_language_columns(
+        self, features: datasets.Features
+    ) -> datasets.Features:
+        """Add ``language_persistent`` / ``language_events`` to ``features``
+        when the underlying parquet shards declare them but the metadata
+        doesn't. No-op when neither column is present or both are
+        already registered.
+        """
+        # Find any one parquet to peek at; bail if there are none yet
+        # (the dataset will fail later for an unrelated reason and we
+        # want that error to surface as-is).
+        try:
+            sample = next((self.root / "data").glob("*/*.parquet"))
+        except StopIteration:
+            return features
+
+        from pyarrow import parquet as _pq  # noqa: PLC0415
+
+        schema_names = set(_pq.read_schema(sample).names)
+        from .language import (  # noqa: PLC0415
+            LANGUAGE_EVENTS,
+            LANGUAGE_PERSISTENT,
+            language_events_column_feature,
+            language_persistent_column_feature,
+        )
+
+        extra: dict[str, object] = {}
+        if LANGUAGE_PERSISTENT in schema_names and LANGUAGE_PERSISTENT not in features:
+            extra[LANGUAGE_PERSISTENT] = language_persistent_column_feature()
+        if LANGUAGE_EVENTS in schema_names and LANGUAGE_EVENTS not in features:
+            extra[LANGUAGE_EVENTS] = language_events_column_feature()
+        if not extra:
+            return features
+        return datasets.Features({**features, **extra})
+
     def _check_cached_episodes_sufficient(self) -> bool:
         """Check if the cached dataset contains all requested episodes and their video files."""
         if self.hf_dataset is None or len(self.hf_dataset) == 0:

src/lerobot/datasets/language.py

@@ -64,8 +64,22 @@ def _json_arrow_type() -> pa.DataType:
 
 
 def _json_feature() -> object:
-    """Return the HF ``datasets`` JSON feature, falling back to a string value."""
-    return datasets.Json() if hasattr(datasets, "Json") else datasets.Value("string")
+    """Return the HF feature used for tool-call payloads.
+
+    Older ``datasets`` versions do not expose ``datasets.Json``. The
+    annotation pipeline currently emits the canonical ``say`` tool call
+    shape, so use that explicit struct instead of falling back to a string
+    that cannot cast structured parquet values.
+    """
+    if hasattr(datasets, "Json"):
+        return datasets.Json()
+    return {
+        "type": datasets.Value("string"),
+        "function": {
+            "name": datasets.Value("string"),
+            "arguments": {"text": datasets.Value("string")},
+        },
+    }
 
 
 def language_persistent_row_arrow_type() -> pa.StructType:

src/lerobot/policies/__init__.py

@@ -26,6 +26,7 @@ from .sac.configuration_sac import SACConfig as SACConfig
 from .sac.reward_model.configuration_classifier import RewardClassifierConfig as RewardClassifierConfig
 from .sarm.configuration_sarm import SARMConfig as SARMConfig
 from .smolvla.configuration_smolvla import SmolVLAConfig as SmolVLAConfig
+from .smolvla2.configuration_smolvla2 import SmolVLA2Config as SmolVLA2Config
 from .tdmpc.configuration_tdmpc import TDMPCConfig as TDMPCConfig
 from .utils import make_robot_action, prepare_observation_for_inference
 from .vqbet.configuration_vqbet import VQBeTConfig as VQBeTConfig
@@ -49,6 +50,7 @@ __all__ = [
     "SACConfig",
     "SARMConfig",
     "SmolVLAConfig",
+    "SmolVLA2Config",
     "TDMPCConfig",
     "VQBeTConfig",
     "WallXConfig",

src/lerobot/policies/smolvla2/chat_processor_smolvla2.py

@@ -99,85 +99,67 @@ class SmolVLA2ChatTokenizerStep(ProcessorStep):
             # falls back to whatever ``task`` is in the transition.
             return transition
 
-        message_streams: list[str | None] = list(comp.get("message_streams") or [])
-        target_indices: list[int] = sorted(
-            int(i) for i in (comp.get("target_message_indices") or [])
-        )
-
         tokenizer = self._get_tokenizer()
-        text_messages = [_strip_lerobot_blocks(m) for m in messages]
 
-        # Tokenize the full chat once.
-        full_ids = tokenizer.apply_chat_template(
-            text_messages,
-            tools=self.tools,
-            add_generation_prompt=False,
-            tokenize=True,
-            return_tensors=None,
+        if _is_batched_messages(messages):
+            encoded = [
+                self._encode_messages(
+                    tokenizer,
+                    msg,
+                    list(streams),
+                    sorted(int(i) for i in indices),
+                )
+                for msg, streams, indices in zip(
+                    messages,
+                    comp.get("message_streams") or [[] for _ in messages],
+                    comp.get("target_message_indices") or [[] for _ in messages],
+                    strict=True,
+                )
+            ]
+        else:
+            encoded = [
+                self._encode_messages(
+                    tokenizer,
+                    messages,
+                    list(comp.get("message_streams") or []),
+                    sorted(int(i) for i in (comp.get("target_message_indices") or [])),
+                )
+            ]
+
+        pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0
+        target_length = self.max_length if self.padding == "max_length" else max(
+            len(ids) for ids, _, _ in encoded
         )
-        if isinstance(full_ids, list) and full_ids and isinstance(full_ids[0], list):
-            full_ids = full_ids[0]
+        target_length = min(target_length, self.max_length)
 
-        # Build the label mask by re-rendering progressively up to each
-        # target message and reading off the prefix length. This is the
-        # robust way to get exact token boundaries: we use the same
-        # tokenizer, the same ``tools=`` argument, and the same chat
-        # template — so the prefix tokens are guaranteed to be a prefix
-        # of the full sequence.
-        labels = [-100] * len(full_ids)
-        for tgt in target_indices:
-            prefix_ids = tokenizer.apply_chat_template(
-                text_messages[:tgt],
-                tools=self.tools,
-                add_generation_prompt=False,
-                tokenize=True,
-                return_tensors=None,
-            )
-            full_through_target = tokenizer.apply_chat_template(
-                text_messages[: tgt + 1],
-                tools=self.tools,
-                add_generation_prompt=False,
-                tokenize=True,
-                return_tensors=None,
-            )
-            if isinstance(prefix_ids, list) and prefix_ids and isinstance(prefix_ids[0], list):
-                prefix_ids = prefix_ids[0]
-            if (
-                isinstance(full_through_target, list)
-                and full_through_target
-                and isinstance(full_through_target[0], list)
-            ):
-                full_through_target = full_through_target[0]
-            start = len(prefix_ids)
-            end = min(len(full_through_target), len(full_ids))
-            for pos in range(start, end):
-                labels[pos] = int(full_ids[pos])
+        ids_batch = []
+        attn_batch = []
+        labels_batch = []
+        predict_actions = []
+        for ids, labels, predict_action in encoded:
+            ids = ids[:target_length]
+            labels = labels[:target_length]
+            attn = [1] * len(ids)
+            if len(ids) < target_length:
+                n_pad = target_length - len(ids)
+                ids = ids + [pad_id] * n_pad
+                labels = labels + [-100] * n_pad
+                attn = attn + [0] * n_pad
+            ids_batch.append(ids)
+            attn_batch.append(attn)
+            labels_batch.append(labels)
+            predict_actions.append(predict_action)
 
-        # Truncate / pad to ``max_length`` so batches collate cleanly.
-        # The SmolVLA pipeline downstream relies on a fixed length
-        # behaviour ("longest" or "max_length") — we mirror it here.
-        if len(full_ids) > self.max_length:
-            full_ids = full_ids[: self.max_length]
-            labels = labels[: self.max_length]
-        attn = [1] * len(full_ids)
-        if self.padding == "max_length" and len(full_ids) < self.max_length:
-            pad_id = (
-                tokenizer.pad_token_id
-                if tokenizer.pad_token_id is not None
-                else 0
-            )
-            n_pad = self.max_length - len(full_ids)
-            full_ids = full_ids + [pad_id] * n_pad
-            labels = labels + [-100] * n_pad
-            attn = attn + [0] * n_pad
+        ids_t = torch.tensor(ids_batch, dtype=torch.long)
+        attn_t = torch.tensor(attn_batch, dtype=torch.bool)
+        labels_t = torch.tensor(labels_batch, dtype=torch.long)
+        predict_actions_t = torch.tensor(predict_actions, dtype=torch.bool)
 
-        ids_t = torch.tensor(full_ids, dtype=torch.long)
-        attn_t = torch.tensor(attn, dtype=torch.bool)
-        labels_t = torch.tensor(labels, dtype=torch.long)
-        predict_actions = any(
-            i < len(message_streams) and message_streams[i] == "low_level"
-            for i in target_indices
-        )
+        if not _is_batched_messages(messages):
+            ids_t = ids_t.squeeze(0)
+            attn_t = attn_t.squeeze(0)
+            labels_t = labels_t.squeeze(0)
+            predict_actions_t = predict_actions_t.squeeze(0)
 
         new_complementary = dict(comp)
         # Drop the per-recipe sidecar keys; everything downstream needs
@@ -194,7 +176,7 @@ class SmolVLA2ChatTokenizerStep(ProcessorStep):
         observation[OBS_LANGUAGE_TOKENS] = ids_t
         observation[OBS_LANGUAGE_ATTENTION_MASK] = attn_t
         new_complementary["text_labels"] = labels_t
-        new_complementary["predict_actions"] = torch.tensor(predict_actions, dtype=torch.bool)
+        new_complementary["predict_actions"] = predict_actions_t
         new_complementary.pop("task", None)
 
         new_transition = dict(transition)
@@ -202,6 +184,53 @@ class SmolVLA2ChatTokenizerStep(ProcessorStep):
         new_transition[TransitionKey.OBSERVATION] = observation
         return new_transition
 
+    def _encode_messages(
+        self,
+        tokenizer: Any,
+        messages: list[dict[str, Any]],
+        message_streams: list[str | None],
+        target_indices: list[int],
+    ) -> tuple[list[int], list[int], bool]:
+        text_messages = [_strip_lerobot_blocks(m) for m in messages]
+
+        full_ids = tokenizer.apply_chat_template(
+            text_messages,
+            tools=self.tools,
+            add_generation_prompt=False,
+            tokenize=True,
+            return_tensors=None,
+        )
+        full_ids = _as_token_ids(full_ids)
+
+        labels = [-100] * len(full_ids)
+        for tgt in target_indices:
+            prefix_ids = tokenizer.apply_chat_template(
+                text_messages[:tgt],
+                tools=self.tools,
+                add_generation_prompt=False,
+                tokenize=True,
+                return_tensors=None,
+            )
+            full_through_target = tokenizer.apply_chat_template(
+                text_messages[: tgt + 1],
+                tools=self.tools,
+                add_generation_prompt=False,
+                tokenize=True,
+                return_tensors=None,
+            )
+            prefix_ids = _as_token_ids(prefix_ids)
+            full_through_target = _as_token_ids(full_through_target)
+            start = len(prefix_ids)
+            end = min(len(full_through_target), len(full_ids))
+            for pos in range(start, end):
+                labels[pos] = int(full_ids[pos])
+
+        predict_actions = any(
+            i < len(message_streams) and message_streams[i] == "low_level"
+            for i in target_indices
+        )
+        return [int(i) for i in full_ids], labels, predict_actions
+
     def transform_features(
         self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
     ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
@@ -247,15 +276,11 @@ def _strip_lerobot_blocks(message: dict[str, Any]) -> dict[str, Any]:
                 continue
             if block.get("type") == "text":
                 text_parts.append({"type": "text", "text": str(block.get("text", ""))})
-        # If only one text block survives, flatten to a string for
-        # template friendliness; some chat templates choke on a single-
-        # element list.
-        if len(text_parts) == 1:
-            new["content"] = text_parts[0]["text"]
-        elif text_parts:
-            new["content"] = text_parts
-        else:
-            new["content"] = ""
+        new["content"] = text_parts or [{"type": "text", "text": ""}]
+    elif content is None:
+        new["content"] = [{"type": "text", "text": ""}]
+    else:
+        new["content"] = [{"type": "text", "text": str(content)}]
     if "tool_calls" in new and not new["tool_calls"]:
         # Drop empty tool_calls — some templates render them as a
         # spurious empty marker.
@@ -267,5 +292,19 @@ def _strip_lerobot_blocks(message: dict[str, Any]) -> dict[str, Any]:
     return new
 
 
+def _is_batched_messages(messages: Any) -> bool:
+    return isinstance(messages, list) and bool(messages) and isinstance(messages[0], list)
+
+
+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"]
+    if hasattr(value, "tolist"):
+        value = value.tolist()
+    if isinstance(value, list) and value and isinstance(value[0], list):
+        value = value[0]
+    return [int(i) for i in value]
+
+
 # Re-export for tests / introspection
 strip_lerobot_blocks = _strip_lerobot_blocks

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

@@ -40,6 +40,7 @@ from .steps import (
     UserInterjectionFwd,
 )
 from .triggers import EventTrigger, HzTrigger, Tick, TickClock, Trigger
+from .ui import make_state_panel, print_robot_lines, print_user_line
 
 __all__ = [
     # runtime
@@ -65,4 +66,8 @@ __all__ = [
     "UserInterjectionFwd",
     "AskVQAFwd",
     "DispatchToolCalls",
+    # UI
+    "make_state_panel",
+    "print_robot_lines",
+    "print_user_line",
 ]

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

@@ -82,10 +82,20 @@ class SmolVLA2Runtime:
             HighLevelSubtaskFwd(
                 trigger=HzTrigger(self.high_level_hz),
                 policy=self.policy,
+                observation_provider=self.observation_provider,
+            ),
+            MemoryUpdateFwd(
+                policy=self.policy,
+                observation_provider=self.observation_provider,
+            ),
+            UserInterjectionFwd(
+                policy=self.policy,
+                observation_provider=self.observation_provider,
+            ),
+            AskVQAFwd(
+                policy=self.policy,
+                observation_provider=self.observation_provider,
             ),
-            MemoryUpdateFwd(policy=self.policy),
-            UserInterjectionFwd(policy=self.policy),
-            AskVQAFwd(policy=self.policy),
             DispatchToolCalls(tools=self.tools),
         ]
         self.state = initial_runtime_state()
@@ -127,6 +137,39 @@ class SmolVLA2Runtime:
 
         self._on_shutdown()
 
+    # ------------------------------------------------------------------
+    # REPL helper: drive one full pipeline pass and return its logs
+    # ------------------------------------------------------------------
+
+    def step_once(self) -> list[str]:
+        """Run one tick of the pipeline and return the log lines.
+
+        Used by the interactive REPL: instead of a background thread,
+        the CLI drives ticks synchronously after each user input. Logs
+        are returned (not printed) so the caller can route them into
+        the rich-Live chat scrollback.
+        """
+        from .triggers import Tick  # noqa: PLC0415
+
+        # Synthesize a tick. We don't need the real wall-clock pacing
+        # here — the REPL drives the runtime, not vice versa — but
+        # ``HzTrigger`` uses ``tick.monotonic_seconds`` to gate, so we
+        # bump it generously so every Hz-triggered step considers
+        # itself due.
+        import time as _time  # noqa: PLC0415
+
+        prev_index = self.state.get("_tick").index if isinstance(self.state.get("_tick"), Tick) else 0
+        self.state["_tick"] = Tick(index=prev_index + 1, monotonic_seconds=_time.monotonic())
+        self.state["log_lines"] = []
+        # ``events_this_tick`` is set up by the caller before
+        # ``step_once`` (the REPL pushes user-driven events first).
+        self.state.setdefault("events_this_tick", [])
+
+        for step in self.pipeline:
+            self.state = step(self.state)
+
+        return list(self.state.get("log_lines") or [])
+
     # ------------------------------------------------------------------
     # I/O
     # ------------------------------------------------------------------

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

@@ -91,10 +91,36 @@ class LowLevelForward(InferenceStep):
     def run(self, state: dict[str, Any]) -> dict[str, Any] | None:
         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
         observation = self.observation_provider()
         if observation is None:
             return None
-        action = self.policy.select_action(observation)
+        # 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.
+        ctx = _control_context_messages(state)
+        if state.get("current_subtask"):
+            ctx = ctx + [{"role": "assistant", "content": state["current_subtask"]}]
+        text_batch = _build_text_batch(self.policy, ctx)
+        from lerobot.utils.constants import (  # noqa: PLC0415
+            OBS_LANGUAGE_ATTENTION_MASK,
+            OBS_LANGUAGE_TOKENS,
+        )
+
+        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)
+        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}")
+            return None
         # SmolVLA returns a single action; if the underlying policy
         # streams chunks, split per-step here. For v1 we just enqueue
         # the result.
@@ -145,17 +171,68 @@ def _build_text_batch(policy: Any, prompt_messages: list[dict[str, Any]]) -> dic
         tokenizer.pad_token = tokenizer.eos_token
 
     text_messages = [_strip_recipe_keys(m) for m in prompt_messages]
-    ids = tokenizer.apply_chat_template(
+    encoded = tokenizer.apply_chat_template(
         text_messages,
         add_generation_prompt=True,
         tokenize=True,
         return_tensors="pt",
     )
+    # ``apply_chat_template`` can return any of:
+    #   - a Tensor of shape ``(seq,)`` or ``(1, seq)`` (older transformers),
+    #   - a list[int] / list[list[int]] (when ``return_tensors`` is ignored),
+    #   - a ``BatchEncoding`` dict-like with ``input_ids`` / ``attention_mask``
+    #     (newer transformers, especially via processor.apply_chat_template
+    #     forwarding through here).
+    # Normalise to ``ids: Tensor[1, seq]`` and grab the encoder's
+    # attention mask when available so we don't have to re-derive it
+    # from ``pad_token_id`` (which can be ``None`` for SmolVLM).
+    attn: Any = None
+    if hasattr(encoded, "input_ids"):
+        ids = encoded.input_ids
+        attn = getattr(encoded, "attention_mask", None)
+    elif isinstance(encoded, dict) and "input_ids" in encoded:
+        ids = encoded["input_ids"]
+        attn = encoded.get("attention_mask")
+    else:
+        ids = encoded
     if isinstance(ids, list):
-        ids = ids[0] if ids else []
+        if ids and isinstance(ids[0], list):
+            ids = ids[0]
+        import torch  # noqa: PLC0415
+
+        ids = torch.tensor(ids, dtype=torch.long)
     if hasattr(ids, "ndim") and ids.ndim == 1:
         ids = ids.unsqueeze(0)
-    attn = (ids != tokenizer.pad_token_id) if tokenizer.pad_token_id is not None else None
+    if attn is None and tokenizer.pad_token_id is not None:
+        attn = ids != tokenizer.pad_token_id
+    elif isinstance(attn, list):
+        import torch  # noqa: PLC0415
+
+        attn = torch.tensor(attn, dtype=torch.long)
+        if attn.ndim == 1:
+            attn = attn.unsqueeze(0)
+    # SmolVLA's ``eager_attention_forward`` does
+    # ``torch.where(attention_mask[..., None, :, :], ...)`` which
+    # requires a *bool* condition tensor; ``BatchEncoding``'s
+    # attention_mask is typically Long (0/1). Cast so the prefix
+    # forward doesn't blow up with ``where expected condition to be a
+    # boolean tensor, but got a tensor with dtype Long``.
+    if attn is not None and hasattr(attn, "dtype"):
+        import torch as _torch  # noqa: PLC0415
+
+        if attn.dtype != _torch.bool:
+            attn = attn.bool()
+    # Move tokens onto the policy's device — otherwise prefix embedding
+    # raises a device-mismatch on every forward (CPU tensor vs MPS / CUDA
+    # model), which the caller's broad except would swallow silently.
+    device = getattr(getattr(policy, "config", None), "device", None)
+    if device is not None:
+        try:
+            ids = ids.to(device)
+            if attn is not None and hasattr(attn, "to"):
+                attn = attn.to(device)
+        except Exception as exc:  # noqa: BLE001
+            logger.debug("could not move lang tokens to %s: %s", device, exc)
     return {"lang_tokens": ids, "lang_masks": attn, "tokenizer": tokenizer}
 
 
@@ -168,37 +245,68 @@ def _strip_recipe_keys(m: dict[str, Any]) -> dict[str, Any]:
 
 @dataclass
 class HighLevelSubtaskFwd(InferenceStep):
-    """At ~1 Hz, ask the policy for the next subtask."""
+    """At ~1 Hz, ask the policy for the next subtask.
+
+    Mirrors the ``high_level_subtask`` recipe layout exactly:
+
+        user:   "${task}\\nPlan: ${plan}\\nMemory: ${memory}"
+        user:   "Current subtask: ${subtask}"        (if subtask present)
+        ↓ generate ↓
+        assistant: <next subtask>
+    """
 
     policy: Any = None
+    observation_provider: Any = None
+    """Same shape as ``LowLevelForward.observation_provider``. When
+    set, the resulting observation is merged into ``select_message``'s
+    batch so text generation runs against real video + state."""
+
     trigger: Trigger = field(default_factory=lambda: HzTrigger(hz=1.0))
 
     def run(self, state: dict[str, Any]) -> dict[str, Any] | None:
         if self.policy is None or not state.get("task"):
             return None
-        ctx = _control_context_messages(state)
-        msg = _generate_with_policy(self.policy, ctx)
+        ctx = _msgs_for_subtask(state)
+        observation = _maybe_observation(self.observation_provider)
+        msg = _generate_with_policy(
+            self.policy, ctx, observation=observation, state=state, label="subtask gen"
+        )
         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")
         return None
 
 
 @dataclass
 class MemoryUpdateFwd(InferenceStep):
-    """On subtask boundary, refresh the compressed memory."""
+    """On subtask boundary, refresh the compressed memory.
+
+    Mirrors the ``memory_update`` recipe layout exactly:
+
+        user:      "${task}"
+        assistant: "Previous memory: ${prior_memory}"   (if prior memory)
+        user:      "Completed subtask: ${completed_subtask}"  (if subtask)
+        ↓ generate ↓
+        assistant: <new memory>
+    """
 
     policy: Any = None
+    observation_provider: Any = None
     trigger: Trigger = field(default_factory=lambda: EventTrigger("subtask_change"))
 
     def run(self, state: dict[str, Any]) -> dict[str, Any] | None:
         # Don't consume the event — multiple steps may want to react.
         if self.policy is None:
             return None
-        ctx = _control_context_messages(state, include_completed=True)
-        new_memory = _generate_with_policy(self.policy, ctx)
+        ctx = _msgs_for_memory(state)
+        observation = _maybe_observation(self.observation_provider)
+        new_memory = _generate_with_policy(
+            self.policy, ctx, observation=observation, state=state, label="memory gen"
+        )
         if new_memory:
             set_if_changed(state, "current_memory", new_memory, label="memory")
         return None
@@ -206,20 +314,31 @@ class MemoryUpdateFwd(InferenceStep):
 
 @dataclass
 class UserInterjectionFwd(InferenceStep):
-    """On stdin interjection, refresh the plan + emit a paired ``say``."""
+    """On stdin interjection, refresh the plan + emit a paired ``say``.
+
+    Mirrors the ``user_interjection_response`` recipe layout exactly:
+
+        user:      "${task}"
+        assistant: "Previous plan:\\n${prior_plan}"   (if prior plan)
+        user:      "${interjection}"                  (the new utterance)
+        ↓ generate ↓
+        assistant: <plan + <say>...</say>>
+    """
 
     policy: Any = None
+    observation_provider: Any = None
     trigger: Trigger = field(default_factory=lambda: EventTrigger("user_interjection"))
 
     def run(self, state: dict[str, Any]) -> dict[str, Any] | None:
         if self.policy is None or not take_event(state, "user_interjection"):
             return None
-        ctx = _control_context_messages(
-            state,
-            extra_user=state.get("recent_interjection"),
+        ctx = _msgs_for_interjection(state)
+        observation = _maybe_observation(self.observation_provider)
+        out = _generate_with_policy(
+            self.policy, ctx, observation=observation, state=state, label="plan/say gen"
         )
-        out = _generate_with_policy(self.policy, ctx)
         if not out:
+            push_log(state, "  [info] plan/say gen produced no text this tick")
             return None
         # Heuristic split: model is trained to emit one assistant turn
         # carrying both plan text AND a `say` tool call. Look for a
@@ -244,9 +363,20 @@ class UserInterjectionFwd(InferenceStep):
 
 @dataclass
 class AskVQAFwd(InferenceStep):
-    """On stdin question, answer a frame-grounded VQA."""
+    """On stdin question, answer a frame-grounded VQA.
+
+    Mirrors the ``ask_vqa_*`` recipe layout exactly: a single user
+    turn carrying just the VQA question, plus the camera image block
+    in training (we drop the image at inference because the dataset's
+    image preprocessing doesn't match SmolVLM's vision tower input).
+
+        user:   <question>
+        ↓ generate ↓
+        assistant: <vqa answer>
+    """
 
     policy: Any = None
+    observation_provider: Any = None
     trigger: Trigger = field(default_factory=lambda: EventTrigger("user_vqa_query"))
 
     def run(self, state: dict[str, Any]) -> dict[str, Any] | None:
@@ -255,8 +385,11 @@ class AskVQAFwd(InferenceStep):
         question = state.get("recent_vqa_query")
         if not question:
             return None
-        ctx = _control_context_messages(state, extra_user=question)
-        answer = _generate_with_policy(self.policy, ctx)
+        ctx = _msgs_for_vqa(question)
+        observation = _maybe_observation(self.observation_provider)
+        answer = _generate_with_policy(
+            self.policy, ctx, observation=observation, state=state, label="vqa gen"
+        )
         if answer:
             push_log(state, f"  vqa: {answer}")
         state["recent_vqa_query"] = None
@@ -326,38 +459,136 @@ def _control_context_messages(
     return msgs
 
 
-def _generate_with_policy(policy: Any, messages: list[dict[str, Any]]) -> str:
-    """Drive ``policy.select_message`` with a minimal text-only batch.
+# ---------------------------------------------------------------------------
+# Per-recipe prompt builders. Each one mirrors a single sub-recipe's
+# message layout in ``smolvla2_hirobot.yaml`` so the chat-templated
+# prompt at inference matches what the model saw during training.
+# Generic ``_control_context_messages`` is kept around as a fallback
+# for ad-hoc callers but the four high-level steps now use these.
+# ---------------------------------------------------------------------------
 
-    Best-effort: the runtime today doesn't construct a full
-    observation batch with images / state for text generation; the
-    text-head was trained over images + lang + state, so generations
-    here may differ in distribution from training. This is acceptable
-    for a v1 REPL; a follow-up will plug in the real observation.
+
+def _msgs_for_subtask(state: dict[str, Any]) -> list[dict[str, Any]]:
+    """``high_level_subtask`` recipe layout."""
+    head_parts = [state.get("task") or ""]
+    if state.get("current_plan"):
+        head_parts.append(f"Plan: {state['current_plan']}")
+    if state.get("current_memory"):
+        head_parts.append(f"Memory: {state['current_memory']}")
+    msgs: list[dict[str, Any]] = [
+        {"role": "user", "content": "\n".join(head_parts)}
+    ]
+    if state.get("current_subtask"):
+        msgs.append(
+            {"role": "user", "content": f"Current subtask: {state['current_subtask']}"}
+        )
+    return msgs
+
+
+def _msgs_for_memory(state: dict[str, Any]) -> list[dict[str, Any]]:
+    """``memory_update`` recipe layout."""
+    msgs: list[dict[str, Any]] = [
+        {"role": "user", "content": state.get("task") or ""}
+    ]
+    if state.get("current_memory"):
+        msgs.append(
+            {
+                "role": "assistant",
+                "content": f"Previous memory: {state['current_memory']}",
+            }
+        )
+    if state.get("current_subtask"):
+        msgs.append(
+            {
+                "role": "user",
+                "content": f"Completed subtask: {state['current_subtask']}",
+            }
+        )
+    return msgs
+
+
+def _msgs_for_interjection(state: dict[str, Any]) -> list[dict[str, Any]]:
+    """``user_interjection_response`` recipe layout."""
+    msgs: list[dict[str, Any]] = [
+        {"role": "user", "content": state.get("task") or ""}
+    ]
+    if state.get("current_plan"):
+        msgs.append(
+            {"role": "assistant", "content": f"Previous plan:\n{state['current_plan']}"}
+        )
+    interjection = state.get("recent_interjection")
+    if interjection:
+        msgs.append({"role": "user", "content": interjection})
+    return msgs
+
+
+def _msgs_for_vqa(question: str) -> list[dict[str, Any]]:
+    """``ask_vqa_*`` recipe layout (text-only at inference)."""
+    return [{"role": "user", "content": question}]
+
+
+def _maybe_observation(provider: Any) -> dict | None:
+    """Pull one observation from ``provider`` if it's set, else ``None``.
+
+    Errors from the provider are logged at debug level and swallowed —
+    text generation still runs (in text-only mode) so a flaky frame
+    source doesn't kill the REPL.
+    """
+    if provider is None:
+        return None
+    try:
+        return provider()
+    except Exception as exc:  # noqa: BLE001
+        logger.debug("observation_provider raised %s — falling back to text-only", exc)
+        return None
+
+
+def _generate_with_policy(
+    policy: Any,
+    messages: list[dict[str, Any]],
+    *,
+    observation: dict | None = None,
+    state: dict[str, Any] | None = None,
+    label: str = "select_message",
+) -> str:
+    """Drive ``policy.select_message`` with a chat batch (and optional obs).
+
+    When ``observation`` carries ``observation.images.*`` and
+    ``observation.state``, those are merged into the batch so
+    ``select_message`` runs the same VLM prefix the policy was trained
+    on. Without an observation the runtime falls back to a text-only
+    prompt — the text head still runs, but generations may drift from
+    the training distribution.
+
+    Failures are surfaced both to the module logger (``warning``) and,
+    when ``state`` is given, to the runtime's user-visible log via
+    :func:`push_log`, so the REPL no longer "looks dead" when
+    something goes wrong inside generation.
     """
     if not hasattr(policy, "select_message"):
+        if state is not None:
+            push_log(state, f"  [warn] policy has no select_message — skipping {label}")
         return ""
     text_batch = _build_text_batch(policy, messages)
-    # ``select_message`` expects a real batch with OBS_LANGUAGE_TOKENS.
-    # The minimal text-only batch we build doesn't have images / state,
-    # so we either run a text-only forward (handled by SmolVLA2 when
-    # supported) or skip and return empty. v1 returns empty when the
-    # policy can't handle it; the runtime logs and continues.
     try:
-        # Convert to the OBS_LANGUAGE_TOKENS / OBS_LANGUAGE_ATTENTION_MASK
-        # keys ``select_message`` uses internally.
         from lerobot.utils.constants import (  # noqa: PLC0415
             OBS_LANGUAGE_ATTENTION_MASK,
             OBS_LANGUAGE_TOKENS,
         )
 
-        batch = {
+        batch: dict[str, Any] = {
             OBS_LANGUAGE_TOKENS: text_batch["lang_tokens"],
             OBS_LANGUAGE_ATTENTION_MASK: text_batch["lang_masks"],
         }
+        if observation:
+            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"])
     except Exception as exc:  # noqa: BLE001
-        logger.debug("select_message fell back: %s", exc)
+        logger.warning("%s failed: %s", label, exc, exc_info=logger.isEnabledFor(logging.DEBUG))
+        if state is not None:
+            push_log(state, f"  [warn] {label} failed: {type(exc).__name__}: {exc}")
         return ""
 
 

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

@@ -0,0 +1,119 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Rich-based REPL layout for the SmolVLA2 runtime.
+
+Two-zone terminal layout:
+
+    [chat scrollback — user messages / robot responses, scrolls naturally]
+
+    ┌── State ──────────────────────────────────────────┐
+    │ task     please clean up the kitchen              │
+    │ subtask  grasp the handle of the sponge           │
+    │ plan     1. grasp sponge  2. wipe  3. tidy        │
+    │ memory   sponge picked up; counter still dirty    │
+    └───────────────────────────────────────────────────┘
+    > _
+
+The state panel re-renders on every state change. Chat lines are
+``console.print``'d above the live region so they accumulate naturally
+in scrollback. Implemented with :class:`rich.live.Live` plus
+:func:`rich.console.Console.input` for the prompt — when an input is
+pending, ``rich.Live`` auto-suspends so the input doesn't fight the
+panel for cursor position.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+try:  # rich is optional; only required for the interactive REPL.
+    from rich.console import Console
+    from rich.panel import Panel
+    from rich.table import Table
+    from rich.text import Text
+
+    _HAS_RICH = True
+except ImportError:  # pragma: no cover
+    _HAS_RICH = False
+    Console = Any  # type: ignore[assignment]
+    Panel = Any  # type: ignore[assignment]
+    Table = Any  # type: ignore[assignment]
+    Text = Any  # type: ignore[assignment]
+
+
+_STATE_KEYS = (
+    ("task", "task"),
+    ("current_subtask", "subtask"),
+    ("current_plan", "plan"),
+    ("current_memory", "memory"),
+)
+
+
+def make_state_panel(state: dict[str, Any]) -> Any:
+    """Render the persistent state panel for the live region.
+
+    Returns a :class:`rich.panel.Panel`. Caller passes it to
+    ``Live.update(panel)`` whenever the state changes.
+    """
+    if not _HAS_RICH:
+        raise RuntimeError(
+            "rich is required for the interactive REPL. "
+            "`pip install rich` (it's a transitive dep of lerobot)."
+        )
+    table = Table.grid(padding=(0, 2), expand=True)
+    table.add_column(justify="right", style="dim", no_wrap=True, width=10)
+    table.add_column(justify="left")
+    for key, label in _STATE_KEYS:
+        value = state.get(key)
+        if value is None:
+            rendered = Text("(not set)", style="dim italic")
+        else:
+            rendered = Text(str(value), style="bold")
+        table.add_row(label, rendered)
+    queue = state.get("action_queue")
+    queue_len = len(queue) if hasattr(queue, "__len__") else 0
+    pending = state.get("tool_calls_pending") or []
+    footer = Text.assemble(
+        ("queued actions: ", "dim"),
+        (str(queue_len), "bold cyan"),
+        ("    pending tool calls: ", "dim"),
+        (str(len(pending)), "bold magenta"),
+    )
+    table.add_row("", footer)
+    return Panel(table, title="[bold]SmolVLA2 state[/]", border_style="cyan")
+
+
+def print_user_line(console: Any, line: str) -> None:
+    """Append a user-typed line to the chat scrollback."""
+    if not _HAS_RICH:
+        print(f"you: {line}", flush=True)
+        return
+    console.print(f"[bold cyan]you:[/] {line}")
+
+
+def print_robot_lines(console: Any, lines: list[str]) -> None:
+    """Append robot/runtime log lines to the chat scrollback."""
+    if not _HAS_RICH:
+        for line in lines:
+            print(f"robot: {line.lstrip()}", flush=True)
+        return
+    for line in lines:
+        # The runtime uses leading whitespace + "label: text"; render
+        # the label in green and the value in default for readability.
+        stripped = line.lstrip()
+        if ":" in stripped:
+            label, _, value = stripped.partition(":")
+            console.print(f"[bold green]robot[/] [dim]({label.strip()})[/] {value.strip()}")
+        else:
+            console.print(f"[bold green]robot:[/] {stripped}")

src/lerobot/policies/smolvla2/modeling_smolvla2.py

@@ -60,7 +60,7 @@ class SmolVLA2Policy(SmolVLAPolicy):
     config_class = SmolVLA2Config
     name = "smolvla2"
 
-    def __init__(self, config: SmolVLA2Config, dataset_stats: dict[str, dict[str, Tensor]] | None = None):
+    def __init__(self, config: SmolVLA2Config, **kwargs):
         if not isinstance(config, SmolVLA2Config):
             config = SmolVLA2Config(
                 **{
@@ -69,7 +69,7 @@ class SmolVLA2Policy(SmolVLAPolicy):
                     if hasattr(config, f.name)
                 }
             )
-        super().__init__(config, dataset_stats=dataset_stats)
+        super().__init__(config, **kwargs)
         if config.unfreeze_lm_head and config.text_loss_weight > 0:
             self._unfreeze_lm_head()
 
@@ -200,7 +200,7 @@ class SmolVLA2Policy(SmolVLAPolicy):
             past_key_values=None,
             inputs_embeds=[prefix_embs, None],
             use_cache=False,
-            fill_kv_cache=False,
+            fill_kv_cache=True,
         )
         prefix_out = out_pair[0] if isinstance(out_pair, (tuple, list)) else out_pair
         if prefix_out is None:
@@ -228,8 +228,8 @@ class SmolVLA2Policy(SmolVLAPolicy):
                 f"num_state={num_state})."
             )
 
-        lang_hidden = prefix_out[:, lang_start:lang_end]
         vlm = self.model.vlm_with_expert.vlm
+        lang_hidden = prefix_out[:, lang_start:lang_end].to(vlm.lm_head.weight.dtype)
         logits = vlm.lm_head(lang_hidden)  # (B, num_lang, vocab)
 
         if text_labels.shape[1] != num_lang:
@@ -244,12 +244,14 @@ class SmolVLA2Policy(SmolVLAPolicy):
         # for the same convention.
         shift_logits = logits[:, :-1, :].contiguous()
         shift_labels = text_labels[:, 1:].contiguous().long()
+        valid_labels = shift_labels != -100
         loss = F.cross_entropy(
             shift_logits.reshape(-1, shift_logits.shape[-1]),
             shift_labels.reshape(-1),
             ignore_index=-100,
+            reduction="sum",
         )
-        return loss
+        return loss / valid_labels.sum().clamp(min=1)
 
     # ------------------------------------------------------------------
     # Inference: text generation
@@ -301,40 +303,67 @@ class SmolVLA2Policy(SmolVLAPolicy):
         if eos_token_id is None:
             eos_token_id = tokenizer.eos_token_id
 
+        # Match training's text-loss forward path (see
+        # ``_compute_text_loss`` above): build the full prefix via
+        # ``embed_prefix`` so images + state conditioning is intact,
+        # then loop AR with ``fill_kv_cache=True, use_cache=False``.
+        # That flag combo routes every layer through
+        # ``forward_attn_layer`` (which gracefully skips ``None``
+        # expert inputs via ``if hidden_states is None or layer is
+        # None: continue``) and short-circuits the cache-update logic
+        # so we don't have to manage past_kv. Each step just
+        # re-forwards the cumulative ``[prefix + generated]``
+        # sequence.
+        #
+        # This is O(n²) in generated text length but cheap in
+        # absolute terms: image encoding happens once via the initial
+        # ``embed_prefix`` call, and the per-step cost is just one
+        # SmolVLM transformer pass over a sequence that grows by one
+        # token each time. Standard SmolVLM ``generate`` was the
+        # other tempting path, but it can't accept SmolVLA's custom
+        # ``state_proj`` output and its tile-grid expectations
+        # disagree with our preprocessor — both lead to garbage
+        # generation, which is what the prior approach produced.
         images, img_masks = self.prepare_images(batch)
         state = self.prepare_state(batch)
         lang_tokens = batch[OBS_LANGUAGE_TOKENS]
         lang_masks = batch[OBS_LANGUAGE_ATTENTION_MASK]
 
-        # 1) Embed prefix (images + lang + state) and run with KV cache.
         prefix_embs, prefix_pad_masks, prefix_att_masks = self.model.embed_prefix(
             images, img_masks, lang_tokens, lang_masks, state=state
         )
-        prefix_2d = make_att_2d_masks(prefix_pad_masks, prefix_att_masks)
-        prefix_pos = torch.cumsum(prefix_pad_masks, dim=1) - 1
-        out_pair, past_kv = self.model.vlm_with_expert.forward(
-            attention_mask=prefix_2d,
-            position_ids=prefix_pos,
-            past_key_values=None,
-            inputs_embeds=[prefix_embs, None],
-            use_cache=True,
-            fill_kv_cache=True,
-        )
-        prefix_out = out_pair[0] if isinstance(out_pair, (tuple, list)) else out_pair
-        if prefix_out is None:
-            raise RuntimeError("select_message: prefix forward returned no hidden states.")
 
-        vlm = self.model.vlm_with_expert.vlm
-
-        # 2) Initial logits — sample first new token from the last
-        #    prefix position.
-        last_hidden = prefix_out[:, -1:]
-        device = last_hidden.device
+        device = prefix_embs.device
         bsize = prefix_embs.shape[0]
-        cur_pos = int(prefix_embs.shape[1])
+        vlm = self.model.vlm_with_expert.vlm
+        emb_dim = prefix_embs.shape[-1]
+        text_emb_scale = math.sqrt(emb_dim)
+
+        current_embs = prefix_embs
+        current_pad = prefix_pad_masks
+        current_att = prefix_att_masks
+        ones_step = torch.ones((bsize, 1), dtype=torch.bool, device=device)
 
         generated: list[int] = []
         for _ in range(max_new_tokens):
+            full_2d = make_att_2d_masks(current_pad, current_att)
+            full_pos = torch.cumsum(current_pad, dim=1) - 1
+
+            out_pair, _ = self.model.vlm_with_expert.forward(
+                attention_mask=full_2d,
+                position_ids=full_pos,
+                past_key_values=None,
+                inputs_embeds=[current_embs, None],
+                use_cache=False,
+                fill_kv_cache=True,
+            )
+            prefix_out = out_pair[0] if isinstance(out_pair, (tuple, list)) else out_pair
+            if prefix_out is None:
+                raise RuntimeError(
+                    "select_message: vlm_with_expert.forward returned no hidden states."
+                )
+
+            last_hidden = prefix_out[:, -1:].to(vlm.lm_head.weight.dtype)
             logits_step = vlm.lm_head(last_hidden)[:, -1]  # (B, V)
             next_ids = self._sample_next_token(logits_step, temperature, top_p)
             tok_id = int(next_ids[0].item())
@@ -342,26 +371,13 @@ class SmolVLA2Policy(SmolVLAPolicy):
             if eos_token_id is not None and tok_id == eos_token_id:
                 break
 
-            # 3) Embed the new token and forward with KV cache.
             new_emb = self.model.vlm_with_expert.embed_language_tokens(
                 next_ids.unsqueeze(0)
             )
-            new_emb = new_emb * math.sqrt(new_emb.shape[-1])
-
-            new_pos = torch.full((bsize, 1), cur_pos, device=device, dtype=torch.long)
-            new_attn = torch.ones((bsize, cur_pos + 1), device=device, dtype=torch.bool)
-
-            out_pair, past_kv = self.model.vlm_with_expert.forward(
-                attention_mask=new_attn,
-                position_ids=new_pos,
-                past_key_values=past_kv,
-                inputs_embeds=[new_emb, None],
-                use_cache=True,
-                fill_kv_cache=True,
-            )
-            new_prefix_out = out_pair[0] if isinstance(out_pair, (tuple, list)) else out_pair
-            last_hidden = new_prefix_out[:, -1:]
-            cur_pos += 1
+            new_emb = new_emb * text_emb_scale
+            current_embs = torch.cat([current_embs, new_emb], dim=1)
+            current_pad = torch.cat([current_pad, ones_step], dim=1)
+            current_att = torch.cat([current_att, ones_step], dim=1)
 
         return tokenizer.decode(generated, skip_special_tokens=True).strip()
 

src/lerobot/processor/batch_processor.py

@@ -175,9 +175,6 @@ class AddBatchDimensionComplementaryDataStep(ComplementaryDataProcessorStep):
             if isinstance(task_index_value, Tensor) and task_index_value.dim() == 0:
                 complementary_data["task_index"] = task_index_value.unsqueeze(0)
 
-        complementary_data.pop("language_persistent", None)
-        complementary_data.pop("language_events", None)
-
         if "messages" in complementary_data:
             messages = complementary_data["messages"]
             if isinstance(messages, list) and (not messages or isinstance(messages[0], dict)):

src/lerobot/processor/render_messages_processor.py

@@ -51,6 +51,9 @@ class RenderMessagesStep(ProcessorStep):
         if not persistent and not events:
             return transition
 
+        if _is_batched_language(persistent) or _is_batched_language(events):
+            return self._call_batch(transition, complementary_data, persistent, events)
+
         timestamp = complementary_data.get("timestamp")
         if timestamp is None:
             raise KeyError("RenderMessagesStep requires sample timestamp in complementary data.")
@@ -69,13 +72,64 @@ class RenderMessagesStep(ProcessorStep):
             return None
 
         new_transition = transition.copy()
-        new_complementary_data = dict(complementary_data)
+        new_complementary_data = dict(new_transition.get(TransitionKey.COMPLEMENTARY_DATA) or {})
         new_complementary_data.pop(LANGUAGE_PERSISTENT, None)
         new_complementary_data.pop(LANGUAGE_EVENTS, None)
         new_complementary_data.update(rendered)
         new_transition[TransitionKey.COMPLEMENTARY_DATA] = new_complementary_data
         return new_transition
 
+    def _call_batch(
+        self,
+        transition: EnvTransition,
+        complementary_data: dict[str, Any],
+        persistent_batch: list,
+        events_batch: list,
+    ) -> EnvTransition | None:
+        timestamp = complementary_data.get("timestamp")
+        if timestamp is None:
+            raise KeyError("RenderMessagesStep requires sample timestamp in complementary data.")
+
+        batch_size = max(len(persistent_batch), len(events_batch))
+        messages: list[list[dict[str, Any]]] = []
+        message_streams: list[list[str | None]] = []
+        target_message_indices: list[list[int]] = []
+        keep_indices: list[int] = []
+
+        for i in range(batch_size):
+            rendered = render_sample(
+                recipe=self.recipe,
+                persistent=persistent_batch[i] if i < len(persistent_batch) else [],
+                events=events_batch[i] if i < len(events_batch) else [],
+                t=_batch_value(timestamp, i),
+                sample_idx=int(_batch_value(complementary_data.get("index", 0), i)),
+                task=_batch_value(complementary_data.get("task"), i),
+                dataset_ctx=self.dataset_ctx,
+            )
+            if rendered is None:
+                continue
+            keep_indices.append(i)
+            messages.append(rendered["messages"])
+            message_streams.append(rendered["message_streams"])
+            target_message_indices.append(rendered["target_message_indices"])
+
+        if not messages:
+            return None
+
+        new_transition = (
+            _select_batch_indices(transition, keep_indices)
+            if len(keep_indices) != batch_size
+            else transition.copy()
+        )
+        new_complementary_data = dict(new_transition.get(TransitionKey.COMPLEMENTARY_DATA) or {})
+        new_complementary_data.pop(LANGUAGE_PERSISTENT, None)
+        new_complementary_data.pop(LANGUAGE_EVENTS, None)
+        new_complementary_data["messages"] = messages
+        new_complementary_data["message_streams"] = message_streams
+        new_complementary_data["target_message_indices"] = target_message_indices
+        new_transition[TransitionKey.COMPLEMENTARY_DATA] = new_complementary_data
+        return new_transition
+
     def transform_features(
         self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
     ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
@@ -90,3 +144,37 @@ def _scalar(value: Any) -> float | int:
     if isinstance(value, list) and len(value) == 1:
         return _scalar(value[0])
     return value
+
+
+def _is_batched_language(value: Any) -> bool:
+    return isinstance(value, list) and bool(value) and isinstance(value[0], list)
+
+
+def _batch_value(value: Any, index: int) -> Any:
+    if value is None:
+        return None
+    if isinstance(value, list):
+        return value[index]
+    if hasattr(value, "ndim") and getattr(value, "ndim") > 0:
+        return _scalar(value[index])
+    return _scalar(value)
+
+
+def _select_batch_indices(transition: EnvTransition, indices: list[int]) -> EnvTransition:
+    selected = transition.copy()
+    for key in (TransitionKey.OBSERVATION, TransitionKey.COMPLEMENTARY_DATA):
+        data = selected.get(key)
+        if isinstance(data, dict):
+            selected[key] = {k: _select_value(v, indices) for k, v in data.items()}
+    action = selected.get(TransitionKey.ACTION)
+    if action is not None:
+        selected[TransitionKey.ACTION] = _select_value(action, indices)
+    return selected
+
+
+def _select_value(value: Any, indices: list[int]) -> Any:
+    if isinstance(value, list) and len(value) >= len(indices):
+        return [value[i] for i in indices]
+    if hasattr(value, "index_select") and hasattr(value, "new_tensor") and getattr(value, "ndim", 0) > 0:
+        return value.index_select(0, value.new_tensor(indices).long())
+    return value

src/lerobot/scripts/lerobot_smolvla2_runtime.py

@@ -23,11 +23,21 @@ speech) as they happen.
 Examples
 --------
 
-Dry run on a checkpoint, no robot connected — useful for sanity-
+Dry run on a Hub checkpoint, no robot connected — useful for sanity-
 checking text generation::
 
     uv run lerobot-smolvla2-runtime \\
-        --policy.path=outputs/train/smolvla2_super_poulain/000020000/pretrained_model \\
+        --policy.path=pepijn223/smolvla2_hirobot_super_poulain_tool2 \\
+        --no_robot \\
+        --task="please clean the kitchen"
+
+Same, but feed real frames from an annotated dataset so plan / subtask
+/ memory / VQA generation runs against actual video + state::
+
+    uv run lerobot-smolvla2-runtime \\
+        --policy.path=pepijn223/smolvla2_hirobot_super_poulain_tool2 \\
+        --dataset.repo_id=pepijn223/super_poulain_annotated \\
+        --dataset.episode=0 \\
         --no_robot \\
         --task="please clean the kitchen"
 
@@ -38,6 +48,9 @@ With a real robot::
         --robot.type=so101 --robot.port=/dev/tty.usbmodem... \\
         --tts.voice=alba
 
+``--policy.path`` accepts either a local directory or a Hugging Face
+Hub repo id. ``--dataset.repo_id`` likewise.
+
 Tool dispatch (TTS via ``SayTool``) is enabled by default when
 ``pocket-tts`` is installed; pass ``--no_tts`` to disable.
 """
@@ -47,8 +60,7 @@ from __future__ import annotations
 import argparse
 import logging
 import sys
-from pathlib import Path
-from typing import Any
+from typing import Any, Callable
 
 logger = logging.getLogger("lerobot.smolvla2.runtime")
 
@@ -61,9 +73,51 @@ def _parse_args(argv: list[str] | None = None) -> argparse.Namespace:
     p.add_argument(
         "--policy.path",
         dest="policy_path",
-        type=Path,
+        type=str,
         required=True,
-        help="Path to a trained SmolVLA2 ``pretrained_model`` directory.",
+        help=(
+            "Local directory or Hugging Face Hub repo id pointing at a "
+            "trained SmolVLA2 ``pretrained_model``."
+        ),
+    )
+    p.add_argument(
+        "--dataset.repo_id",
+        dest="dataset_repo_id",
+        type=str,
+        default=None,
+        help=(
+            "Optional dataset (local path or Hub repo id) used to drive "
+            "observations during dry-run inference. When set, the runtime "
+            "reads camera frames + state from the chosen episode and feeds "
+            "them into all forward passes — so plan / subtask / memory / "
+            "VQA generation see the same visual context the policy was "
+            "trained on."
+        ),
+    )
+    p.add_argument(
+        "--dataset.episode",
+        dest="dataset_episode",
+        type=int,
+        default=0,
+        help="Episode index to walk through (default: 0).",
+    )
+    p.add_argument(
+        "--dataset.start_frame",
+        dest="dataset_start_frame",
+        type=int,
+        default=0,
+        help="Frame index within the episode to start from (default: 0).",
+    )
+    p.add_argument(
+        "--dataset.advance_per_tick",
+        dest="dataset_advance_per_tick",
+        type=int,
+        default=1,
+        help=(
+            "How many dataset frames to advance per runtime tick. The "
+            "default of 1 means the runtime walks the episode forward "
+            "frame by frame; set to 0 to freeze on ``start_frame``."
+        ),
     )
     p.add_argument(
         "--task",
@@ -111,16 +165,149 @@ def _parse_args(argv: list[str] | None = None) -> argparse.Namespace:
     return p.parse_args(argv)
 
 
-def _load_policy(path: Path):  # noqa: ANN202
-    """Load a SmolVLA2 checkpoint from ``path``."""
-    from lerobot.policies.factory import make_policy_from_path  # noqa: PLC0415
+def _load_policy_and_preprocessor(
+    policy_path: str,
+    dataset_repo_id: str | None,
+) -> tuple[Any, Any, Any]:
+    """Load a SmolVLA2 checkpoint (local path or Hub repo id).
+
+    When ``dataset_repo_id`` is provided, the dataset's metadata is used
+    to derive policy features (matching the standard
+    ``make_policy(cfg, ds_meta=...)`` flow used by ``lerobot-train`` and
+    ``lerobot-record``). When it isn't, we fall back to instantiating
+    the policy directly via ``from_pretrained`` — this skips the
+    feature-derivation path that ``make_policy`` insists on, but also
+    means we can't load the saved preprocessor pipeline (which depends
+    on ``input_features`` / ``output_features``). For inference-only
+    dry-runs this is fine; the policy still loads.
+
+    Returns ``(policy, preprocessor, ds_meta)`` where ``preprocessor``
+    and ``ds_meta`` may be ``None`` if no dataset was provided.
+    """
+    from lerobot.configs import PreTrainedConfig  # noqa: PLC0415
+    from lerobot.policies.factory import make_policy, make_pre_post_processors  # noqa: PLC0415
+
+    cfg = PreTrainedConfig.from_pretrained(policy_path)
+    cfg.pretrained_path = policy_path
+
+    ds_meta = None
+    preprocessor = None
+    if dataset_repo_id is not None:
+        from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata  # noqa: PLC0415
+
+        ds_meta = LeRobotDatasetMetadata(dataset_repo_id)
+        policy = make_policy(cfg, ds_meta=ds_meta)
+        # NOTE: we deliberately pass ``pretrained_path=None`` here even
+        # though the checkpoint ships a ``policy_preprocessor.json``.
+        # ``RenderMessagesStep`` carries a ``TrainingRecipe`` field that
+        # isn't faithfully serialized into that JSON, so the saved
+        # pipeline can't currently be round-tripped via
+        # ``PolicyProcessorPipeline.from_pretrained`` — it crashes with
+        # ``RenderMessagesStep.__init__() missing 1 required argument:
+        # 'recipe'``. Building fresh from ``cfg`` re-runs
+        # ``make_smolvla2_pre_post_processors``, which loads the recipe
+        # YAML referenced by ``cfg.recipe_path`` and wires it back into
+        # ``RenderMessagesStep`` correctly. Normalization stats come
+        # from ``ds_meta.stats`` (the same dataset the user is feeding
+        # into the runtime), so no quality loss in practice.
+        preprocessor, _ = make_pre_post_processors(
+            cfg,
+            pretrained_path=None,
+            dataset_stats=ds_meta.stats,
+        )
+    else:
+        # No dataset: instantiate the policy class directly so we don't
+        # need ds_meta. This bypasses ``make_policy``'s feature-shape
+        # derivation, which is fine for a pretrained checkpoint where
+        # the saved config already carries those shapes.
+        from lerobot.policies.factory import get_policy_class  # noqa: PLC0415
+
+        policy_cls = get_policy_class(cfg.type)
+        policy = policy_cls.from_pretrained(policy_path, config=cfg)
+        policy.to(cfg.device)
 
-    policy = make_policy_from_path(str(path))
     policy.eval()
-    return policy
+    return policy, preprocessor, ds_meta
 
 
-def _build_tools(policy_path: Path, no_tts: bool, tts_voice: str) -> dict[str, Any]:
+def _build_observation_provider(
+    *,
+    dataset_repo_id: str,
+    episode: int,
+    start_frame: int,
+    advance_per_tick: int,
+    preprocessor: Any,
+    device: str,
+) -> Callable[[], dict | None]:
+    """Build a closure that feeds dataset frames into the runtime.
+
+    Each call returns a preprocessed observation batch (images +
+    state, batched, on the policy's device, normalized) suitable for
+    ``policy.select_action`` and ``policy.select_message``. The
+    closure walks the chosen episode forward by ``advance_per_tick``
+    frames per call, looping back to the episode start when it falls
+    off the end.
+
+    The dataset's ``language_persistent`` / ``language_events``
+    columns are stripped before the sample reaches the preprocessor,
+    so ``RenderMessagesStep`` and ``SmolVLA2ChatTokenizerStep`` are
+    no-ops; the runtime supplies its own messages from current state.
+    """
+    import torch  # noqa: PLC0415
+
+    from lerobot.datasets.lerobot_dataset import LeRobotDataset  # noqa: PLC0415
+
+    ds = LeRobotDataset(dataset_repo_id, episodes=[episode])
+    if len(ds) == 0:
+        raise ValueError(
+            f"Dataset {dataset_repo_id!r} episode {episode} is empty."
+        )
+
+    state = {"cursor": max(0, min(start_frame, len(ds) - 1))}
+
+    def _provider() -> dict | None:
+        idx = state["cursor"]
+        if advance_per_tick > 0:
+            state["cursor"] = (idx + advance_per_tick) % len(ds)
+
+        sample = ds[idx]
+        # Strip the language columns so the preprocessor's render step
+        # is a no-op — the runtime drives messages itself.
+        for k in ("language_persistent", "language_events"):
+            sample.pop(k, None)
+
+        if preprocessor is not None:
+            sample = preprocessor(sample)
+
+        # Keep only observation keys; the runtime's text path will
+        # merge these with its own lang_tokens / lang_masks.
+        observation = {
+            k: v
+            for k, v in sample.items()
+            if isinstance(k, str) and k.startswith("observation.")
+        }
+        # Defensive: if something further upstream forgot the batch
+        # dim, add it now so downstream Tensor ops don't crash.
+        for k, v in list(observation.items()):
+            if isinstance(v, torch.Tensor) and v.ndim > 0 and v.shape[0] != 1:
+                # ``add_batch_dim`` already ran inside the preprocessor;
+                # an unbatched tensor at this point means a step
+                # somewhere produced an unbatched output. Best-effort
+                # fix.
+                if v.shape[0] != 1 and v.ndim < 4 and "image" not in k:
+                    observation[k] = v.unsqueeze(0)
+        # Move to device (the preprocessor's DeviceProcessorStep should
+        # already have done this when ``preprocessor is not None``;
+        # this is a belt-and-braces no-op in the common case).
+        for k, v in list(observation.items()):
+            if isinstance(v, torch.Tensor):
+                observation[k] = v.to(device)
+        return observation
+
+    return _provider
+
+
+def _build_tools(no_tts: bool, tts_voice: str) -> dict[str, Any]:
     """Instantiate the tools declared on this dataset/policy."""
     if no_tts:
         return {}
@@ -133,27 +320,70 @@ def _build_tools(policy_path: Path, no_tts: bool, tts_voice: str) -> dict[str, A
         return {}
 
 
+def _silence_noisy_loggers() -> None:
+    """Drop chatty third-party loggers down to WARNING.
+
+    HuggingFace / httpx / urllib3 emit one log line per HTTP request,
+    which the REPL has to print between the state block and the
+    prompt — completely unreadable. We never need that detail in the
+    REPL and the user can opt back into it via ``-v`` (verbose mode
+    keeps DEBUG on the lerobot loggers but still gates the noisy ones
+    here unless they explicitly want them).
+    """
+    for name in (
+        "httpcore",
+        "httpcore.connection",
+        "httpcore.http11",
+        "httpcore.proxy",
+        "httpx",
+        "urllib3",
+        "urllib3.connectionpool",
+        "huggingface_hub",
+        "huggingface_hub.repocard",
+        "huggingface_hub.file_download",
+        "transformers",
+        "transformers.modeling_utils",
+        "transformers.tokenization_utils_base",
+        "datasets",
+        "filelock",
+    ):
+        logging.getLogger(name).setLevel(logging.WARNING)
+
+
 def main(argv: list[str] | None = None) -> int:
     args = _parse_args(argv)
     logging.basicConfig(
         level=logging.DEBUG if args.verbose else logging.INFO,
         format="%(asctime)s %(levelname)s %(message)s",
     )
-
-    if not args.policy_path.exists():
-        print(f"[smolvla2] policy path not found: {args.policy_path}", file=sys.stderr)
-        return 1
+    _silence_noisy_loggers()
 
     print(f"[smolvla2] loading policy from {args.policy_path}", flush=True)
-    policy = _load_policy(args.policy_path)
+    policy, preprocessor, _ds_meta = _load_policy_and_preprocessor(
+        args.policy_path, args.dataset_repo_id
+    )
 
-    tools = _build_tools(args.policy_path, args.no_tts, args.tts_voice)
+    observation_provider: Callable[[], dict | None] | None = None
+    if args.dataset_repo_id is not None:
+        print(
+            f"[smolvla2] streaming observations from {args.dataset_repo_id} "
+            f"episode={args.dataset_episode} "
+            f"start_frame={args.dataset_start_frame}",
+            flush=True,
+        )
+        observation_provider = _build_observation_provider(
+            dataset_repo_id=args.dataset_repo_id,
+            episode=args.dataset_episode,
+            start_frame=args.dataset_start_frame,
+            advance_per_tick=args.dataset_advance_per_tick,
+            preprocessor=preprocessor,
+            device=str(getattr(policy.config, "device", "cpu")),
+        )
+
+    tools = _build_tools(args.no_tts, args.tts_voice)
     if tools:
         print(f"[smolvla2] tools loaded: {list(tools)}", flush=True)
 
-    # Robot wiring is left as a follow-up — for v1 we run language-only
-    # / dry-run so REPL development doesn't require a connected robot.
-    observation_provider = None
     robot_executor = None
     if not args.no_robot:
         print(
@@ -162,33 +392,135 @@ def main(argv: list[str] | None = None) -> int:
             flush=True,
         )
 
-    from lerobot.policies.smolvla2.inference import (  # noqa: PLC0415
-        SmolVLA2Runtime,
-        StdinReader,
-    )
+    from lerobot.policies.smolvla2.inference import SmolVLA2Runtime  # noqa: PLC0415
 
     runtime = SmolVLA2Runtime(
         policy=policy,
         tools=tools,
         observation_provider=observation_provider,
         robot_executor=robot_executor,
-        event_collector=StdinReader().poll,
+        # No background event collector — the REPL drives ticks
+        # synchronously after each user input. The runtime's own
+        # ``run()`` loop is bypassed here in favour of ``step_once()``
+        # so the input prompt and the live state panel co-exist
+        # cleanly.
+        event_collector=None,
         chunk_hz=args.chunk_hz,
         ctrl_hz=args.ctrl_hz,
         high_level_hz=args.high_level_hz,
     )
     if args.task:
         runtime.set_task(args.task)
-    print(
-        "[smolvla2] runtime ready. Type a task to begin, then any line for "
-        "interjections, questions ending in '?' for VQA, or 'stop' to exit.",
-        flush=True,
-    )
+
+    return _run_repl(runtime, initial_task=args.task, max_ticks=args.max_ticks)
+
+
+def _run_repl(runtime: Any, *, initial_task: str | None, max_ticks: int | None) -> int:
+    """Claude-Code-style block REPL.
+
+    Each turn redraws a status block (task / subtask / plan / memory)
+    at the top, prints any robot log lines that came in since the last
+    turn, then asks for input on a clean ``> `` prompt at the bottom.
+    No live region, no panel re-renders, no rendering races with HTTP
+    log lines — just clear-screen + reprint each turn, the way a
+    chat-style REPL is meant to look.
+    """
     try:
-        runtime.run(max_ticks=args.max_ticks)
+        from rich.console import Console  # noqa: PLC0415
+    except ImportError:
+        print(
+            "[smolvla2] rich is required for the interactive REPL. "
+            "`pip install rich` and re-run.",
+            file=sys.stderr,
+        )
+        return 2
+
+    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:
+        # Already shown the help line in _redraw when task is None.
+        pass
+    ticks_done = 0
+    try:
+        while True:
+            try:
+                line = console.input("[bold cyan]> [/]").strip()
+            except EOFError:
+                break
+            if not line:
+                _redraw(last_logs)
+                continue
+            lower = line.lower()
+            if lower in {"stop", "quit", "exit"}:
+                break
+
+            # Inject the user input as the right kind of event,
+            # then run a single pipeline tick to consume it.
+            if not runtime.state.get("task"):
+                task = line[5:].strip() if lower.startswith("task:") else line
+                runtime.set_task(task)
+            elif lower.endswith("?"):
+                runtime.state["recent_vqa_query"] = line
+                runtime.state.setdefault("events_this_tick", []).append(
+                    "user_vqa_query"
+                )
+            else:
+                runtime.state["recent_interjection"] = line
+                runtime.state.setdefault("events_this_tick", []).append(
+                    "user_interjection"
+                )
+
+            last_logs = runtime.step_once() or []
+            _redraw(last_logs)
+
+            ticks_done += 1
+            if max_ticks is not None and ticks_done >= max_ticks:
+                break
     except KeyboardInterrupt:
-        runtime.stop()
-        print("\n[smolvla2] interrupted by user", flush=True)
+        console.print("\n[dim]interrupted[/]")
+    console.print("[dim]runtime stopped[/]")
     return 0
 
 

src/lerobot/utils/collate.py

@@ -22,7 +22,7 @@ from torch.utils.data._utils.collate import default_collate
 
 from lerobot.datasets.language import LANGUAGE_COLUMNS
 
-_PYTHON_LIST_KEYS = {"messages", "message_streams", "target_message_indices"}
+_PYTHON_LIST_KEYS = {"messages", "message_streams", "target_message_indices", *LANGUAGE_COLUMNS}
 
 
 def lerobot_collate_fn(batch: list[dict[str, Any] | None]) -> dict[str, Any] | None: