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: blend:
action_execution: high_level_subtask:
weight: 0.75 weight: 0.50
bindings: bindings:
new_memory: "emitted_at(t, style=memory)" new_memory: "emitted_at(t, style=memory)"
messages: messages:
- role: user - role: user
stream: high_level stream: high_level
content: "${task}\nPlan: ${plan}\nMemory: ${memory}" content: "${task}\nPlan: ${plan}\nMemory: ${memory}"
- {role: assistant, content: "${subtask}", stream: low_level, target: true, if_present: subtask} - {role: assistant, content: "${subtask}", stream: high_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: "${new_memory}", stream: high_level, target: true, if_present: new_memory} - {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: plan_generation:
weight: 0.10 weight: 0.10
bindings: bindings:
@@ -30,7 +35,7 @@ blend:
- {role: assistant, content: "${current_plan}", stream: high_level, target: true, if_present: current_plan} - {role: assistant, content: "${current_plan}", stream: high_level, target: true, if_present: current_plan}
ask_vqa_top: ask_vqa_top:
weight: 0.075 weight: 0.05
bindings: bindings:
vqa_query: "emitted_at(t, style=vqa, role=user, camera=observation.images.front)" 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)" 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} - {role: assistant, content: "${vqa}", stream: high_level, target: true, if_present: vqa}
ask_vqa_wrist: ask_vqa_wrist:
weight: 0.075 weight: 0.05
bindings: bindings:
vqa_query: "emitted_at(t, style=vqa, role=user, camera=observation.images.wrist)" 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)" 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; # The action expert is conditioned on (images, state, subtask)
# supervises the current subtask (low_level: flow + text CE) # only — NOT on task / plan / memory. We achieve this by splitting
# and, at memory-boundary frames, the new memory too. # the work across two main sub-recipes:
# 2. plan_generation — task → plan (text only). Trains the #
# model to produce a plan from a bare task description so # 1. high_level_subtask — text-only. Trains the LM head to predict
# the runtime can call it at episode start / replanning. # the current subtask from (task + plan + memory). At a memory
# 3. ask_vqa_{top,wrist} — text-only VQA on a camera image, # boundary, also predicts the new memory in the same forward.
# gated by ``if_present`` so they only fire on annotated frames. # 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: blend:
action_execution: high_level_subtask:
weight: 0.75 weight: 0.50
bindings: bindings:
new_memory: "emitted_at(t, style=memory)" new_memory: "emitted_at(t, style=memory)"
messages: messages:
- role: user - role: user
stream: high_level stream: high_level
content: "${task}\nPlan: ${plan}\nMemory: ${memory}" content: "${task}\nPlan: ${plan}\nMemory: ${memory}"
- {role: assistant, content: "${subtask}", stream: low_level, target: true, if_present: subtask} - {role: assistant, content: "${subtask}", stream: high_level, target: true, if_present: subtask}
# Boundary-frame tail: at a subtask transition, predict the # Boundary-frame tail: at a subtask transition, also predict
# new memory as a second assistant turn (same forward pass). # the new memory in the same forward pass.
- {role: assistant, content: "${new_memory}", stream: high_level, target: true, if_present: new_memory} - {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: plan_generation:
weight: 0.10 weight: 0.10
bindings: bindings:
@@ -33,7 +49,7 @@ blend:
- {role: assistant, content: "${current_plan}", stream: high_level, target: true, if_present: current_plan} - {role: assistant, content: "${current_plan}", stream: high_level, target: true, if_present: current_plan}
ask_vqa_top: ask_vqa_top:
weight: 0.075 weight: 0.05
bindings: bindings:
vqa_query: "emitted_at(t, style=vqa, role=user, camera=observation.images.front)" 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)" 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} - {role: assistant, content: "${vqa}", stream: high_level, target: true, if_present: vqa}
ask_vqa_wrist: ask_vqa_wrist:
weight: 0.075 weight: 0.05
bindings: bindings:
vqa_query: "emitted_at(t, style=vqa, role=user, camera=observation.images.wrist)" 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)" vqa: "emitted_at(t, style=vqa, role=assistant, camera=observation.images.wrist)"
@@ -111,11 +111,16 @@ class LowLevelForward(InferenceStep):
if observation is None: if observation is None:
return None return None
# Same prompt construction as before — task + plan + memory, # π0.5-style: the action expert is conditioned on just the
# optional current subtask — then merge into the obs batch. # subtask (+ images + state). No task / plan / memory in the
ctx = _control_context_messages(state) # low-level prompt — those are only used by the high-level
if state.get("current_subtask"): # loop to *generate* the subtask. Matches the training-time
ctx = ctx + [{"role": "assistant", "content": state["current_subtask"]}] # ``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) text_batch = _build_text_batch(self.policy, ctx)
from lerobot.utils.constants import ( # noqa: PLC0415 from lerobot.utils.constants import ( # noqa: PLC0415
OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_ATTENTION_MASK,