refactor(recipes): π0.5-style split — action expert conditions on subtask only

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) <noreply@anthropic.com>
This commit is contained in:
Pepijn
2026-05-13 14:13:07 +02:00
parent 129aa207e3
commit 7a32f8a72a
3 changed files with 54 additions and 28 deletions
+12 -7
View File
@@ -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)"
@@ -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)"
@@ -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,