From 7a32f8a72a7f7a56205035ea77d9c0b3f31894ac Mon Sep 17 00:00:00 2001 From: Pepijn Date: Wed, 13 May 2026 14:13:07 +0200 Subject: [PATCH] =?UTF-8?q?refactor(recipes):=20=CF=800.5-style=20split=20?= =?UTF-8?q?=E2=80=94=20action=20expert=20conditions=20on=20subtask=20only?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previously ``action_execution`` rendered ``task + plan + memory + subtask`` into one prefix and ran the flow loss on it. That meant the action expert was conditioned on the full hierarchical context (closer to π0.7 §V.A), not just the subtask. The π0.5 paper's hierarchical inference has the action expert see only the *subtask* (plus images and state). Split the recipe to match: high_level_subtask (0.50) user(task + plan + memory) → assistant(subtask) [+ assistant(new_memory) at boundary frames] All ``stream: high_level`` → text-CE only, no flow loss. low_level_execution (0.30) user(subtask) → assistant(subtask) Both ``stream: low_level`` → flow loss fires; text CE on the subtask is a small redundant extra signal. Prefix the action expert sees: [images, subtask, state]. plan_generation (0.10) — unchanged. ask_vqa_{top,wrist} (0.05 each) — unchanged. Runtime: the low-level loop in ``smolvla2/inference/steps.py`` now sends ``[user(subtask), assistant(subtask)]`` to ``predict_action_chunk`` instead of the full task+plan+memory context. Falls back to ``state['task']`` when no subtask has been generated yet so the first frame still has something to condition on. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../configs/recipes/pi052_hirobot.yaml | 19 +++++--- .../configs/recipes/smolvla2_hirobot.yaml | 48 ++++++++++++------- .../policies/smolvla2/inference/steps.py | 15 ++++-- 3 files changed, 54 insertions(+), 28 deletions(-) diff --git a/src/lerobot/configs/recipes/pi052_hirobot.yaml b/src/lerobot/configs/recipes/pi052_hirobot.yaml index b56d04334..15ae92b0f 100644 --- a/src/lerobot/configs/recipes/pi052_hirobot.yaml +++ b/src/lerobot/configs/recipes/pi052_hirobot.yaml @@ -8,19 +8,24 @@ blend: - action_execution: - weight: 0.75 + high_level_subtask: + weight: 0.50 bindings: new_memory: "emitted_at(t, style=memory)" messages: - role: user stream: high_level content: "${task}\nPlan: ${plan}\nMemory: ${memory}" - - {role: assistant, content: "${subtask}", stream: low_level, target: true, if_present: subtask} - # Boundary-frame tail: at a subtask transition, predict the - # new memory as a second assistant turn (same forward pass). + - {role: assistant, content: "${subtask}", stream: high_level, target: true, if_present: subtask} - {role: assistant, content: "${new_memory}", stream: high_level, target: true, if_present: new_memory} + low_level_execution: + weight: 0.30 + messages: + # Action expert prefix = [images, subtask, state] only — π0.5 style. + - {role: user, content: "${subtask}", stream: low_level, if_present: subtask} + - {role: assistant, content: "${subtask}", stream: low_level, target: true, if_present: subtask} + plan_generation: weight: 0.10 bindings: @@ -30,7 +35,7 @@ blend: - {role: assistant, content: "${current_plan}", stream: high_level, target: true, if_present: current_plan} ask_vqa_top: - weight: 0.075 + weight: 0.05 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)" @@ -44,7 +49,7 @@ blend: - {role: assistant, content: "${vqa}", stream: high_level, target: true, if_present: vqa} ask_vqa_wrist: - weight: 0.075 + weight: 0.05 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)" diff --git a/src/lerobot/configs/recipes/smolvla2_hirobot.yaml b/src/lerobot/configs/recipes/smolvla2_hirobot.yaml index 6da9e362b..8ff6a1e93 100644 --- a/src/lerobot/configs/recipes/smolvla2_hirobot.yaml +++ b/src/lerobot/configs/recipes/smolvla2_hirobot.yaml @@ -1,29 +1,45 @@ -# SmolVLA2 Hi-Robot blend — three flavors: +# SmolVLA2 Hi-Robot blend — π0.5-style split: # -# 1. action_execution — fused (task + plan + memory) prompt; -# supervises the current subtask (low_level: flow + text CE) -# and, at memory-boundary frames, the new memory too. -# 2. plan_generation — task → plan (text only). Trains the -# model to produce a plan from a bare task description so -# the runtime can call it at episode start / replanning. -# 3. ask_vqa_{top,wrist} — text-only VQA on a camera image, -# gated by ``if_present`` so they only fire on annotated frames. +# The action expert is conditioned on (images, state, subtask) +# only — NOT on task / plan / memory. We achieve this by splitting +# the work across two main sub-recipes: +# +# 1. high_level_subtask — text-only. Trains the LM head to predict +# the current subtask from (task + plan + memory). At a memory +# boundary, also predicts the new memory in the same forward. +# 2. low_level_execution — action. Renders just the subtask as the +# language conditioning so the action expert's prefix is +# [images, subtask, state]. Flow loss + (redundant) text CE on +# the subtask itself. +# 3. plan_generation — text only. task → plan. +# 4. ask_vqa_{top,wrist} — text only. camera-grounded VQA. blend: - action_execution: - weight: 0.75 + high_level_subtask: + weight: 0.50 bindings: new_memory: "emitted_at(t, style=memory)" messages: - role: user stream: high_level content: "${task}\nPlan: ${plan}\nMemory: ${memory}" - - {role: assistant, content: "${subtask}", stream: low_level, target: true, if_present: subtask} - # Boundary-frame tail: at a subtask transition, predict the - # new memory as a second assistant turn (same forward pass). + - {role: assistant, content: "${subtask}", stream: high_level, target: true, if_present: subtask} + # Boundary-frame tail: at a subtask transition, also predict + # the new memory in the same forward pass. - {role: assistant, content: "${new_memory}", stream: high_level, target: true, if_present: new_memory} + low_level_execution: + weight: 0.30 + messages: + # Just the subtask in the prompt — π0.5 style. The action + # expert sees only [images, this subtask, state]. Marking the + # assistant target as ``stream: low_level`` triggers + # ``predict_actions=True`` so the flow loss fires; text CE on + # the subtask is a (small) redundant extra signal. + - {role: user, content: "${subtask}", stream: low_level, if_present: subtask} + - {role: assistant, content: "${subtask}", stream: low_level, target: true, if_present: subtask} + plan_generation: weight: 0.10 bindings: @@ -33,7 +49,7 @@ blend: - {role: assistant, content: "${current_plan}", stream: high_level, target: true, if_present: current_plan} ask_vqa_top: - weight: 0.075 + weight: 0.05 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)" @@ -47,7 +63,7 @@ blend: - {role: assistant, content: "${vqa}", stream: high_level, target: true, if_present: vqa} ask_vqa_wrist: - weight: 0.075 + weight: 0.05 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)" diff --git a/src/lerobot/policies/smolvla2/inference/steps.py b/src/lerobot/policies/smolvla2/inference/steps.py index 878338e17..7bc2d1e16 100644 --- a/src/lerobot/policies/smolvla2/inference/steps.py +++ b/src/lerobot/policies/smolvla2/inference/steps.py @@ -111,11 +111,16 @@ class LowLevelForward(InferenceStep): if observation is None: return None - # Same prompt construction as before — task + plan + memory, - # optional current subtask — then merge into the obs batch. - ctx = _control_context_messages(state) - if state.get("current_subtask"): - ctx = ctx + [{"role": "assistant", "content": state["current_subtask"]}] + # π0.5-style: the action expert is conditioned on just the + # subtask (+ images + state). No task / plan / memory in the + # low-level prompt — those are only used by the high-level + # loop to *generate* the subtask. Matches the training-time + # ``low_level_execution`` recipe shape. + subtask = state.get("current_subtask") or state.get("task") or "" + ctx = [ + {"role": "user", "content": subtask}, + {"role": "assistant", "content": subtask}, + ] text_batch = _build_text_batch(self.policy, ctx) from lerobot.utils.constants import ( # noqa: PLC0415 OBS_LANGUAGE_ATTENTION_MASK,