chore(pi052): trim logging and recipes

docs/source/language_and_recipes.mdx

@@ -144,7 +144,7 @@ The renderer does not apply a tokenizer chat template. Policy processors decide
 ## Blends
 
 Blend recipes select one weighted sub-recipe deterministically from the sample index.
-`recipes/subtasks_vqa.yaml` trains the core blend — high-level subtask prediction, low-level execution, and VQA. `recipes/subtask_mem_vqa_speech.yaml` is the fuller variant that also adds memory updates and spoken interjection responses.
+`recipes/subtask_mem.yaml` trains the compact core blend — high-level subtask prediction, low-level execution, and memory. `recipes/subtask_mem_vqa_speech.yaml` is the fuller variant that also adds VQA and spoken interjection responses.
 
 ## Graceful absence
 

src/lerobot/common/wandb_utils.py

@@ -215,27 +215,9 @@ class WandBLogger:
         *,
         camera_keys: list[str],
         n_samples: int = 4,
-        policy=None,
-        predict_actions: bool = False,
         mode: str = "train",
     ) -> None:
-        """Push a ``wandb.Table`` of training-example rows for the current batch.
-
-        Each row is one batch element with:
-          * one ``wandb.Image`` column per camera in ``camera_keys`` (CHW or
-            HWC, uint8 or float in [0,1] — auto-detected),
-          * any text fields present in the batch (``task`` / ``subtask`` /
-            ``memory`` / ``instruction``),
-          * ground-truth action first/last frame (the action chunk's
-            endpoints — gives a quick sense of trajectory direction),
-          * if ``predict_actions=True`` and ``policy`` is supplied, the model's
-            ``predict_action_chunk`` first/last frame alongside.
-
-        This is opt-in via ``--wandb.log_examples_freq=N`` on the CLI; the
-        training loop calls it once every N steps. Cheap to keep on: with
-        N=4 samples and 3 cameras you upload 12 small image files per dump and (if
-        enabled) run one extra inference forward pass.
-        """
+        """Log a small W&B table with sampled images/text and action endpoints."""
         import logging  # noqa: PLC0415
 
         import numpy as np  # noqa: PLC0415
@@ -253,53 +235,19 @@ class WandBLogger:
             return
         n = min(int(n_samples), bsz)
 
-        # Optional predicted-action forward pass on the first n samples.
-        pred_actions: np.ndarray | None = None
-        if predict_actions and policy is not None:
-            was_training = policy.training
-            try:
-                policy.eval()
-                sub_batch = {}
-                for k, v in batch.items():
-                    if isinstance(v, torch.Tensor):
-                        sub_batch[k] = v[:n]
-                    elif isinstance(v, (list, tuple)):
-                        sub_batch[k] = list(v[:n])
-                    else:
-                        sub_batch[k] = v
-                with torch.no_grad():
-                    pred = policy.predict_action_chunk(sub_batch)
-                pred_actions = pred.detach().cpu().float().numpy()
-            except Exception as exc:  # noqa: BLE001
-                logging.warning(
-                    "log_training_examples: predict_action_chunk failed (%s) — "
-                    "skipping predicted-action columns",
-                    exc,
-                )
-                pred_actions = None
-            finally:
-                if was_training:
-                    policy.train()
-
         present_cameras = [c for c in camera_keys if c in batch]
         text_keys = [k for k in ("task", "subtask", "memory", "instruction") if k in batch]
 
         columns = ["sample"]
         columns.extend(c.removeprefix("observation.images.") or c for c in present_cameras)
         columns.extend(text_keys)
-        columns.append("gt_action_first")
-        columns.append("gt_action_last")
-        if pred_actions is not None:
-            columns.append("pred_action_first")
-            columns.append("pred_action_last")
+        columns += ["gt_action_first", "gt_action_last"]
 
         table = self._wandb.Table(columns=columns)
 
         def _to_uint8_hwc(t: torch.Tensor) -> np.ndarray:
-            # Strip an outer time dim if present: (T, C, H, W) -> first frame.
             if t.ndim == 4:
                 t = t[0]
-            # CHW -> HWC.
             if t.ndim == 3 and t.shape[0] in (1, 3, 4) and t.shape[-1] not in (1, 3, 4):
                 t = t.permute(1, 2, 0)
             arr = t.detach().cpu().float().numpy()
@@ -309,15 +257,13 @@ class WandBLogger:
 
         def _action_endpoints(a: torch.Tensor) -> tuple[str, str]:
             arr = a.detach().cpu().float().numpy()
-            if arr.ndim == 2:  # (T, D)
-                return (
-                    str(np.round(arr[0], 3).tolist()),
-                    str(np.round(arr[-1], 3).tolist()),
-                )
+            if arr.ndim == 2:
+                return str(np.round(arr[0], 3).tolist()), str(np.round(arr[-1], 3).tolist())
             if arr.ndim == 1:
                 rounded = np.round(arr, 3).tolist()
-                return (str(rounded), str(rounded))
-            return (str(arr.tolist()), str(arr.tolist()))
+                return str(rounded), str(rounded)
+            text = str(arr.tolist())
+            return text, text
 
         for i in range(n):
             row: list = [i]
@@ -334,7 +280,7 @@ class WandBLogger:
                     row.append(None)
             for tk in text_keys:
                 v = batch[tk]
-                if isinstance(v, (list, tuple)):
+                if isinstance(v, list | tuple):
                     row.append(str(v[i]) if i < len(v) else "")
                 else:
                     row.append(str(v))
@@ -346,11 +292,6 @@ class WandBLogger:
             else:
                 row.append("")
                 row.append("")
-            if pred_actions is not None:
-                p = torch.from_numpy(pred_actions[i])
-                pfirst, plast = _action_endpoints(p)
-                row.append(pfirst)
-                row.append(plast)
             table.add_data(*row)
 
         self._wandb.log({f"{mode}/examples": table}, step=step)

src/lerobot/configs/default.py

@@ -62,72 +62,22 @@ class WandBConfig:
     run_id: str | None = None
     mode: str | None = None  # Allowed values: 'online', 'offline' 'disabled'. Defaults to 'online'
     add_tags: bool = True  # If True, save configuration as tags in the WandB run.
-    # Periodic training-example dump (independent of ``log_freq``). When > 0,
-    # every ``log_examples_freq`` steps the trainer pushes a ``wandb.Table``
-    # with one row per sampled batch element containing each camera view
-    # (rendered as ``wandb.Image``), any text fields present in the batch
-    # (``task`` / ``subtask`` / ``memory`` / ``instruction``), and the
-    # ground-truth action chunk's first + last frames. Defaults to 5000 — set
-    # to 0 to disable. Only fires when ``enable=True``, so runs without wandb
-    # are unaffected.
+    # Periodic W&B table with sampled images/text and action endpoints. Set to 0 to disable.
     log_examples_freq: int = 5000
-    # Number of batch elements to include in each example dump.
     log_examples_n: int = 4
-    # If True (default), also run ``policy.predict_action_chunk`` on the logged
-    # samples (in eval mode, no_grad) and add predicted vs ground-truth action
-    # columns to the table. Costs one extra forward pass per dump — negligible
-    # at the 5k-step default cadence. Set to ``False`` if your policy doesn't
-    # implement ``predict_action_chunk`` or you want to skip the extra forward.
-    log_examples_predict_actions: bool = True
 
 
 @dataclass
 class EMAConfig:
-    """Exponential Moving Average of trainable policy parameters.
-
-    Diffusion / flow-matching policies (Diffusion Policy, π0/π0.5,
-    pi052) benefit substantially from averaging late-training
-    parameter oscillations — see Chi et al. 2023 §V.D. The official
-    JAX openpi trainer ships EMA with ``ema_decay=0.99`` (default) and
-    ``0.999`` for its pi05_libero config; the openpi PyTorch port
-    explicitly lists EMA as unsupported, and LeRobot main inherited
-    that gap. Enabling this flag plugs ema-pytorch
-    (https://github.com/lucidrains/ema-pytorch) into the LeRobot
-    training loop with a shadow ``nn.Module`` clone of the policy.
-
-    Cost: 1× model params in fp32 shadow (~13 GB for pi052's 3.3B
-    params) + one elementwise update per training step (~1% step time).
-
-    Off by default (opt-in): EMA is only beneficial for flow-matching /
-    diffusion policies (pi0/pi05/pi052), and the fp32 shadow copy is pure
-    overhead for other policies (e.g. VLA-JEPA). Set ``--ema.enable=true``
-    to turn it on (the pi05/pi052 training recipes do this). openpi (JAX)
-    ships EMA on for every config; enable it explicitly to match that.
-    """
+    """EMA shadow for flow/diffusion policies. Off by default because it doubles model memory."""
 
     enable: bool = False
-    # Target EMA decay β in θ_ema ← β·θ_ema + (1-β)·θ_live (passed to
-    # ema-pytorch as ``beta``).
-    #   0.999  — last ~1000 steps; pi05_libero default in openpi
-    #   0.99   — last ~100 steps; openpi top-level default
-    #   0.75   — very fast EMA (Diffusion Policy original setting)
-    #   0.9999 — very slow EMA (long classification runs)
+    # Target EMA decay beta in theta_ema <- beta * theta_ema + (1 - beta) * theta_live.
     decay: float = 0.99
-    # Skip the first N calls to ``ema.update()``; during this window
-    # the shadow is just a hard copy of the live weights (no averaging).
-    # Lets early-training rapid changes settle before averaging begins.
-    # Maps to ema-pytorch's ``update_after_step`` (NOT a smooth decay
-    # ramp like older lerobot EMA implementations).
+    # Initial update calls that keep the shadow as a hard copy before averaging starts.
     warmup_steps: int = 0
-    # When True, the periodic eval block uses the EMA shadow model
-    # directly (``ema.ema_model``) instead of the live policy. Standard
-    # practice for diffusion-style policies — eval scores are usually
-    # 1–3% higher than the live policy at the same step.
+    # Use the EMA model for periodic eval.
     use_for_eval: bool = True
-    # When True, the periodic wandb training-example dump uses the EMA
-    # shadow for the optional predicted-action columns (so what you see
-    # in W&B matches eval behavior).
-    use_for_wandb_examples: bool = True
 
 
 @dataclass

src/lerobot/configs/recipes/subtask_mem.yaml

@@ -1,29 +1,17 @@
-# subtask_mem_vqa_speech — Hi-Robot blend + memory + spoken responses.
+# subtask_mem — compact Hi-Robot blend with memory.
 #
-# Superset of subtasks_vqa.yaml. Keeps the core subtask + action + VQA
-# training, and adds two text-supervised tasks:
+# Trains the core subtask + action objectives and memory updates:
 #
 #   high_level_subtask         — predict the subtask from the task.
 #   low_level_execution        — flow loss with [images, subtask, state].
 #   memory_update              — compress progress into a memory note.
-#   user_interjection_response — reply to a user interjection with a
-#                                spoken `say` tool call (no plan, no
-#                                subtask text — just the spoken reply).
-#   ask_vqa_{top,wrist}        — camera-grounded VQA.
 #
 # Plan is intentionally left out — memory is the only persistent
 # high-level state here, keeping the prompt short.
 #
-# Requires the dataset to carry `memory`, `interjection` and `say`-tool
-# annotations (the annotation pipeline's memory + interjection modules)
-# in addition to `subtask` and `vqa`. Sub-recipes whose `if_present`
-# bindings are missing simply don't render for that sample, so a
-# dataset without interjections still trains the rest of the blend.
-#
-# Tool-call note: the `say` tool call on the interjection-response turn
-# is flattened to a `<say>...</say>` text marker by the tokenizer step
-# (`_flatten_say_tool_calls`) so the LM head learns to emit exactly the
-# marker the runtime parses back (`_split_plan_and_say`).
+# Requires the dataset to carry `subtask` and `memory` annotations.
+# Sub-recipes whose `if_present` bindings are missing simply don't
+# render for that sample.
 
 blend:
 
@@ -65,4 +53,4 @@ blend:
       - {role: user, content: "${task}", stream: high_level}
       - {role: assistant, content: "Previous memory: ${prior_memory}", stream: high_level, if_present: prior_memory}
       - {role: user, content: "Completed subtask: ${completed_subtask}", stream: high_level, if_present: completed_subtask}
-      - {role: assistant, content: "${current_memory}", stream: high_level, target: true, if_present: current_memory}
+      - {role: assistant, content: "${current_memory}", stream: high_level, target: true, if_present: current_memory}

src/lerobot/configs/recipes/subtask_mem_vqa_robocasa.yaml

@@ -1,99 +0,0 @@
-# subtask_mem_vqa_robocasa — Hi-Robot blend tuned for RoboCasa cameras.
-#
-# Same supervision as ``subtask_mem.yaml`` (subtask + memory) plus
-# camera-grounded VQA across the three RoboCasa camera keys produced
-# by ``slurm_build_robocasa_composite_seen.py``:
-#
-#   observation.images.robot0_agentview_left   (left scene view)
-#   observation.images.robot0_agentview_right  (right scene view)
-#   observation.images.robot0_eye_in_hand      (wrist)
-#
-# The annotation pipeline (``examples/annotations/run_hf_job.py``) emits
-# VQA per camera, so each anchor frame produces three (user, assistant)
-# rows tagged with their source camera. Each VQA sub-recipe consumes
-# the rows for one camera via ``camera=...`` resolver bindings.
-#
-# Spatial VQA targets (bbox / point) are rewritten from JSON to
-# PaliGemma ``<locDDDD>`` tokens by ``_messages_vqa_to_loc`` —
-# ``register_paligemma_loc_tokens`` already collapses them to single
-# detection-vocab ids so the LM head learns the pretrained pointing /
-# detection prior, not a 7-piece BPE salad.
-#
-# Interjections / spoken responses are intentionally absent — the
-# annotation job runs with ``--interjections.enabled=false``.
-
-blend:
-
-  high_level_subtask:
-    weight: 0.25
-    messages:
-      - {role: user, content: "${task}", stream: high_level}
-      - {role: assistant, content: "${subtask}", stream: high_level, target: true, if_present: subtask}
-
-  low_level_execution:
-    weight: 0.45
-    messages:
-      # Action expert is conditioned on the SUBTASK; at inference the
-      # high-level loop generates it via the LM head and feeds it here.
-      # ``stream: low_level`` flips ``predict_actions=True`` so the flow
-      # loss fires; subtask CE is owned by ``high_level_subtask``.
-      - {role: user, content: "${subtask}", stream: low_level, if_present: subtask}
-
-  memory_update:
-    # Trained densely with ``active_at`` — every frame inside a subtask
-    # interval — so the (prior_memory, completed_subtask) → current_memory
-    # mapping is supervised against varied observations. The *when* to
-    # emit lives in the inference trigger (subtask_change), not the
-    # model. See ``subtask_mem.yaml`` for the long version of this note.
-    weight: 0.15
-    bindings:
-      prior_memory: "nth_prev(style=memory, offset=1)"
-      current_memory: "active_at(t, style=memory)"
-      completed_subtask: "nth_prev(style=subtask, offset=1)"
-    messages:
-      - {role: user, content: "${task}", stream: high_level}
-      - {role: assistant, content: "Previous memory: ${prior_memory}", stream: high_level, if_present: prior_memory}
-      - {role: user, content: "Completed subtask: ${completed_subtask}", stream: high_level, if_present: completed_subtask}
-      - {role: assistant, content: "${current_memory}", stream: high_level, target: true, if_present: current_memory}
-
-  ask_vqa_agentview_left:
-    weight: 0.05
-    bindings:
-      vqa_query: "emitted_at(t, style=vqa, role=user, camera=observation.images.robot0_agentview_left)"
-      vqa: "emitted_at(t, style=vqa, role=assistant, camera=observation.images.robot0_agentview_left)"
-    messages:
-      - role: user
-        stream: high_level
-        if_present: vqa_query
-        content:
-          - {type: image, feature: observation.images.robot0_agentview_left}
-          - {type: text, text: "${vqa_query}"}
-      - {role: assistant, content: "${vqa}", stream: high_level, target: true, if_present: vqa}
-
-  ask_vqa_agentview_right:
-    weight: 0.05
-    bindings:
-      vqa_query: "emitted_at(t, style=vqa, role=user, camera=observation.images.robot0_agentview_right)"
-      vqa: "emitted_at(t, style=vqa, role=assistant, camera=observation.images.robot0_agentview_right)"
-    messages:
-      - role: user
-        stream: high_level
-        if_present: vqa_query
-        content:
-          - {type: image, feature: observation.images.robot0_agentview_right}
-          - {type: text, text: "${vqa_query}"}
-      - {role: assistant, content: "${vqa}", stream: high_level, target: true, if_present: vqa}
-
-  ask_vqa_wrist:
-    weight: 0.05
-    bindings:
-      vqa_query: "emitted_at(t, style=vqa, role=user, camera=observation.images.robot0_eye_in_hand)"
-      vqa: "emitted_at(t, style=vqa, role=assistant, camera=observation.images.robot0_eye_in_hand)"
-    messages:
-      - role: user
-        stream: high_level
-        if_present: vqa_query
-        content:
-          - {type: image, feature: observation.images.robot0_eye_in_hand}
-          - {type: text, text: "${vqa_query}"}
-      - {role: assistant, content: "${vqa}", stream: high_level, target: true, if_present: vqa}

src/lerobot/configs/recipes/subtask_mem_vqa_speech.yaml

@@ -1,7 +1,6 @@
 # subtask_mem_vqa_speech — Hi-Robot blend + memory + spoken responses.
 #
-# Superset of subtasks_vqa.yaml. Keeps the core subtask + action + VQA
-# training, and adds two text-supervised tasks:
+# Extends the compact subtask_mem recipe with VQA and spoken interjection responses:
 #
 #   high_level_subtask         — predict the subtask from the task.
 #   low_level_execution        — flow loss with [images, subtask, state].
@@ -83,8 +82,7 @@ blend:
       - {role: assistant, stream: high_level, target: true, if_present: speech, tool_calls_from: speech}
 
   # VQA is view-dependent — each camera gets its own sub-recipe so the
-  # resolver disambiguates via `camera=...`. Camera keys match
-  # subtasks_vqa.yaml (`front` + `wrist`); adjust to your dataset.
+  # resolver disambiguates via `camera=...`. Adjust camera keys to your dataset.
   ask_vqa_top:
     weight: 0.075
     bindings:

src/lerobot/configs/recipes/subtasks_vqa.yaml

@@ -1,61 +0,0 @@
-# subtasks_vqa — Hi-Robot blend for PI052 (PaliGemma backbone).
-#
-#   Trains two things only: subtasks and VQA. Plan and memory are
-#   intentionally left out — keeps the prompt short and the training
-#   surface small. The fuller blend with memory + spoken replies is
-#   ``subtask_mem_vqa_speech.yaml``.
-#
-#     high_level_subtask  — predict the subtask from the task.
-#     low_level_execution — flow loss with [images, subtask, state].
-#     ask_vqa_{top,wrist} — camera-grounded VQA.
-#
-# PI052's text tokenizer renders these messages as plain
-# ``Role: content`` text (PaliGemma is not chat-pretrained).
-
-blend:
-
-  high_level_subtask:
-    weight: 0.40
-    messages:
-      - {role: user, content: "${task}", stream: high_level}
-      - {role: assistant, content: "${subtask}", stream: high_level, target: true, if_present: subtask}
-
-  low_level_execution:
-    weight: 0.40
-    messages:
-      # The action expert is conditioned on the SUBTASK — at inference
-      # the high-level loop (``HighLevelSubtaskFwd``) generates the
-      # subtask via the LM head and feeds it here. The action expert's
-      # prefix is [images, subtask, state]. ``stream: low_level`` flips
-      # ``predict_actions=True`` so the flow loss fires; no text-CE
-      # target here (subtask prediction is owned by
-      # ``high_level_subtask``).
-      - {role: user, content: "${subtask}", stream: low_level, if_present: subtask}
-
-  ask_vqa_top:
-    weight: 0.10
-    bindings:
-      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.front}
-          - {type: text, text: "${vqa_query}"}
-      - {role: assistant, content: "${vqa}", stream: high_level, target: true, if_present: vqa}
-
-  ask_vqa_wrist:
-    weight: 0.10
-    bindings:
-      vqa_query: "emitted_at(t, style=vqa, role=user, camera=observation.images.wrist)"
-      vqa: "emitted_at(t, style=vqa, role=assistant, camera=observation.images.wrist)"
-    messages:
-      - role: user
-        stream: high_level
-        if_present: vqa_query
-        content:
-          - {type: image, feature: observation.images.wrist}
-          - {type: text, text: "${vqa_query}"}
-      - {role: assistant, content: "${vqa}", stream: high_level, target: true, if_present: vqa}

src/lerobot/policies/groot/groot_n1.py

@@ -26,14 +26,9 @@ from lerobot.utils.import_utils import _transformers_available
 
 # Conditional import for type checking and lazy loading
 if TYPE_CHECKING or _transformers_available:
-    from huggingface_hub.dataclasses import strict
     from transformers import AutoConfig, AutoModel, PretrainedConfig, PreTrainedModel
     from transformers.feature_extraction_utils import BatchFeature
 else:
-
-    def strict(cls):
-        return cls
-
     AutoConfig = None
     AutoModel = None
     PretrainedConfig = object

src/lerobot/policies/pi052/__init__.py

@@ -24,8 +24,8 @@ Extends :class:`lerobot.policies.pi05.PI05Policy` with:
 * per-component prompt dropout (Pi 0.7 §V.E) for regularising the
   text head against missing context at inference.
 
-See ``src/lerobot/configs/recipes/subtasks_vqa.yaml`` for the
-canonical training recipe and
+See ``src/lerobot/configs/recipes/subtask_mem.yaml`` for the compact
+training recipe and
 ``examples/training/pi052_hirobot.slurm`` for the launcher.
 """
 

src/lerobot/policies/pi052/configuration_pi052.py

@@ -55,10 +55,10 @@ class PI052Config(PI05Config):
     """
 
     # Recipe / language stack ---------------------------------------------
-    recipe_path: str | None = "recipes/subtasks_vqa.yaml"
+    recipe_path: str | None = "recipes/subtask_mem.yaml"
     """Path (absolute or relative to ``src/lerobot/configs/``) to a
-    ``TrainingRecipe`` YAML. Defaults to the canonical Hi-Robot blend
-    shipped alongside this policy. Set to ``None`` to disable recipe
+    ``TrainingRecipe`` YAML. Defaults to the compact Hi-Robot blend
+    shipped with this policy. Set to ``None`` to disable recipe
     rendering and fall back to π0.5's single-task ``Task: ... Action:``
     prompt path (unannotated datasets keep working that way)."""
 

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

@@ -17,7 +17,7 @@ Each step is a tiny class with a ``trigger`` and an ``__call__(state)``;
 the runtime applies them in order each tick. When a step's trigger
 doesn't fire, the step is a no-op and the runtime moves on.
 
-Stream-to-step mapping mirrors the ``subtasks_vqa.yaml`` recipe:
+Stream-to-step mapping mirrors the PI052 training recipe:
 
 * ``LowLevelForward``        — calls ``policy.select_action`` for the
                                 action chunk; trained by
@@ -153,8 +153,7 @@ class LowLevelForward(InferenceStep):
             )
             push_log(
                 state,
-                f"  [warn] predict_action_chunk failed: "
-                f"{type(exc).__name__}: {exc}",
+                f"  [warn] predict_action_chunk failed: {type(exc).__name__}: {exc}",
             )
             return None
 
@@ -288,9 +287,7 @@ def _build_text_batch(
         register_paligemma_loc_tokens,
     )
 
-    tok_name = (
-        getattr(policy.config, "tokenizer_name", None) or "google/paligemma-3b-pt-224"
-    )
+    tok_name = getattr(policy.config, "tokenizer_name", None) or "google/paligemma-3b-pt-224"
     # Register PaliGemma's <locDDDD> tokens so inference encoding /
     # decoding sees them as single vocab ids — must match training.
     # The tokenizer is read-only after registration, so cache it: rebuilding it
@@ -482,9 +479,7 @@ class HighLevelSubtaskFwd(InferenceStep):
                 # despite the chunk having drained (visual scene may
                 # have changed but the LM is replaying training
                 # tokens).
-                state["subtask_repeat_count"] = (
-                    state.get("subtask_repeat_count", 0) + 1
-                )
+                state["subtask_repeat_count"] = state.get("subtask_repeat_count", 0) + 1
         # Silently skip empty completions — common when the model
         # warms up or generates only EOS; logging it every tick at
         # ctrl_hz is just noise.
@@ -729,9 +724,7 @@ def _looks_like_gibberish(text: str) -> bool:
     # Length-independent: many tokens but a tiny unique ratio. The
     # earlier ``< 80`` check missed these because the looped string
     # blows well past 80 chars.
-    if len(tokens) >= 8 and len(unique_alpha) <= max(3, len(tokens) // 10):
-        return True
-    return False
+    return len(tokens) >= 8 and len(unique_alpha) <= max(3, len(tokens) // 10)
 
 
 def _control_context_messages(
@@ -742,7 +735,7 @@ def _control_context_messages(
 ) -> list[dict[str, Any]]:
     """Build a chat-template-ready prompt from current runtime state.
 
-    Mirrors what ``subtasks_vqa.yaml`` renders into ``${task}\nPlan:
+    Mirrors what the recipe renders into ``${task}\nPlan:
     ${plan}\nMemory: ${memory}`` for the high-level branches.
     """
     # Always emit ``Plan: `` / ``Memory: `` labels — even with empty
@@ -762,7 +755,7 @@ def _control_context_messages(
 
 # ---------------------------------------------------------------------------
 # Per-recipe prompt builders. Each one mirrors a single sub-recipe's
-# message layout in ``subtasks_vqa.yaml`` so the chat-templated
+# message layout in the recipe 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.
@@ -817,26 +810,18 @@ def _msgs_for_memory(state: dict[str, Any]) -> list[dict[str, Any]]:
     ]
     prior_memory = state.get("current_memory")
     if prior_memory:
-        msgs.append(
-            {"role": "assistant", "content": f"Previous memory: {prior_memory}"}
-        )
+        msgs.append({"role": "assistant", "content": f"Previous memory: {prior_memory}"})
     completed_subtask = state.get("prior_subtask")
     if completed_subtask:
-        msgs.append(
-            {"role": "user", "content": f"Completed subtask: {completed_subtask}"}
-        )
+        msgs.append({"role": "user", "content": f"Completed subtask: {completed_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 ""}
-    ]
+    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']}"}
-        )
+        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})

src/lerobot/scripts/lerobot_train.py

@@ -725,12 +725,10 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         )
         ema.to(accelerator.device)
         logging.info(
-            "EMA enabled (ema-pytorch): beta=%g, update_after_step=%d, "
-            "use_for_eval=%s, use_for_wandb_examples=%s",
+            "EMA enabled (ema-pytorch): beta=%g, update_after_step=%d, use_for_eval=%s",
             cfg.ema.decay,
             cfg.ema.warmup_steps,
             cfg.ema.use_for_eval,
-            cfg.ema.use_for_wandb_examples,
         )
 
         # Resume the EMA shadow if a previous run wrote one.
@@ -874,10 +872,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
                     wandb_logger.log_dict(wandb_log_dict, step)
             train_tracker.reset_averages()
 
-        # Periodic training-example dump to wandb (camera images + text
-        # fields + action endpoints). Opt-in via ``--wandb.log_examples_freq``;
-        # independent of ``--log_freq`` so you can keep scalar logs frequent
-        # and the heavier visual dump rare (e.g. every 5000 steps).
+        # Periodic W&B example table (camera images + text fields + action endpoints).
         if (
             wandb_logger is not None
             and cfg.wandb.log_examples_freq > 0
@@ -885,23 +880,11 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
             and is_main_process
         ):
             try:
-                # Optionally use the EMA shadow model directly for the
-                # predicted-action columns (matches what eval / deployment
-                # would see). ``ema-pytorch`` exposes the shadow as a
-                # full ``nn.Module`` at ``ema.ema_model``, so we just
-                # pass that instead of swap-and-restore.
-                target_policy = (
-                    ema.ema_model
-                    if (ema is not None and cfg.ema.use_for_wandb_examples)
-                    else accelerator.unwrap_model(policy)
-                )
                 wandb_logger.log_training_examples(
                     batch=batch,
                     step=step,
                     camera_keys=list(dataset.meta.camera_keys),
                     n_samples=cfg.wandb.log_examples_n,
-                    policy=target_policy,
-                    predict_actions=cfg.wandb.log_examples_predict_actions,
                 )
             except Exception as exc:  # noqa: BLE001
                 logging.warning("wandb log_training_examples failed: %s", exc)