mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-24 18:26:11 +00:00
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:
@@ -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,
|
||||||
|
|||||||
Reference in New Issue
Block a user