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,