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)