diff --git a/docs/source/language_and_recipes.mdx b/docs/source/language_and_recipes.mdx
index 92f97d6bb..7df272973 100644
--- a/docs/source/language_and_recipes.mdx
+++ b/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
diff --git a/src/lerobot/common/wandb_utils.py b/src/lerobot/common/wandb_utils.py
index f6154549e..e681ffad6 100644
--- a/src/lerobot/common/wandb_utils.py
+++ b/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)
diff --git a/src/lerobot/configs/default.py b/src/lerobot/configs/default.py
index 205b68940..e813586d4 100644
--- a/src/lerobot/configs/default.py
+++ b/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
diff --git a/src/lerobot/configs/recipes/subtask_mem.yaml b/src/lerobot/configs/recipes/subtask_mem.yaml
index 6903b3585..823d49d75 100644
--- a/src/lerobot/configs/recipes/subtask_mem.yaml
+++ b/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 `...` 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}
\ No newline at end of file
+ - {role: assistant, content: "${current_memory}", stream: high_level, target: true, if_present: current_memory}
diff --git a/src/lerobot/configs/recipes/subtask_mem_vqa_robocasa.yaml b/src/lerobot/configs/recipes/subtask_mem_vqa_robocasa.yaml
deleted file mode 100644
index 607e20e5d..000000000
--- a/src/lerobot/configs/recipes/subtask_mem_vqa_robocasa.yaml
+++ /dev/null
@@ -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 ```` 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}
diff --git a/src/lerobot/configs/recipes/subtask_mem_vqa_speech.yaml b/src/lerobot/configs/recipes/subtask_mem_vqa_speech.yaml
index 2cd1e7ae5..b6424bdf5 100644
--- a/src/lerobot/configs/recipes/subtask_mem_vqa_speech.yaml
+++ b/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:
diff --git a/src/lerobot/configs/recipes/subtasks_vqa.yaml b/src/lerobot/configs/recipes/subtasks_vqa.yaml
deleted file mode 100644
index 48a6ced54..000000000
--- a/src/lerobot/configs/recipes/subtasks_vqa.yaml
+++ /dev/null
@@ -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}
diff --git a/src/lerobot/policies/groot/groot_n1.py b/src/lerobot/policies/groot/groot_n1.py
index 6987f7f37..79de45185 100644
--- a/src/lerobot/policies/groot/groot_n1.py
+++ b/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
diff --git a/src/lerobot/policies/pi052/__init__.py b/src/lerobot/policies/pi052/__init__.py
index c726b7790..dc3928bdf 100644
--- a/src/lerobot/policies/pi052/__init__.py
+++ b/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.
"""
diff --git a/src/lerobot/policies/pi052/configuration_pi052.py b/src/lerobot/policies/pi052/configuration_pi052.py
index 2b1576929..e645e0cb0 100644
--- a/src/lerobot/policies/pi052/configuration_pi052.py
+++ b/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)."""
diff --git a/src/lerobot/policies/pi052/inference/steps.py b/src/lerobot/policies/pi052/inference/steps.py
index 819f73cd4..a9af3b7a1 100644
--- a/src/lerobot/policies/pi052/inference/steps.py
+++ b/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 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})
diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py
index dbf7cede8..9987556f1 100644
--- a/src/lerobot/scripts/lerobot_train.py
+++ b/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)