mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-06 09:37:06 +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:
|
||||
|
||||
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,
|
||||
|
||||
Reference in New Issue
Block a user