From 1cd1ec468e11011ecec12f4fbeac1905ef0c7bb2 Mon Sep 17 00:00:00 2001 From: Khalil Meftah Date: Fri, 3 Jul 2026 16:46:56 +0200 Subject: [PATCH] feat(molmoact2): add RECAP advantage conditioning via recipe system for MolmoAct2 - Add recipe_path, advantage_prefix, cfg_beta to MolmoAct2Config - Place advantage clause in the assistant section of _build_robot_text - Add MolmoAct2NormalizeTaskStep for consistent task normalization - Parse recipe-rendered advantage in MolmoAct2PackInputsProcessorStep - Insert RenderMessagesStep pipeline when recipe_path is configured - Add recap_advantage_molmoact2.yaml recipe --- .../recipes/recap_advantage_molmoact2.yaml | 28 ++++ .../molmoact2/configuration_molmoact2.py | 13 ++ .../policies/molmoact2/processor_molmoact2.py | 139 +++++++++++++++--- 3 files changed, 157 insertions(+), 23 deletions(-) create mode 100644 src/lerobot/configs/recipes/recap_advantage_molmoact2.yaml diff --git a/src/lerobot/configs/recipes/recap_advantage_molmoact2.yaml b/src/lerobot/configs/recipes/recap_advantage_molmoact2.yaml new file mode 100644 index 000000000..2763bbcc1 --- /dev/null +++ b/src/lerobot/configs/recipes/recap_advantage_molmoact2.yaml @@ -0,0 +1,28 @@ +# RECAP advantage recipe for MolmoAct2. +# +# Renders task + advantage into the task field as " Advantage: ". +# MolmoAct2PackInputsProcessorStep parses this, extracts the advantage value, +# and places it AFTER the full user prompt but BEFORE action tokens — matching +# the RECAP paper (Section V-B): "The advantage indicator appears in the training +# sequence after ˆℓ but before the actions, such that only the action +# log-likelihoods are affected." +# +# Final prompt layout: +# <|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\nAdvantage: positive. ... +# +# When advantage is absent (CFG dropout), if_present guard skips this message +# and RenderedMessagesToTaskStep leaves the task unchanged — no advantage clause. + +bindings: + advantage: "active_at(t, style=advantage)" + +messages: + - role: user + content: "${task} Advantage: ${advantage}" + stream: high_level + if_present: advantage + + - role: assistant + content: "" + stream: low_level + target: true diff --git a/src/lerobot/policies/molmoact2/configuration_molmoact2.py b/src/lerobot/policies/molmoact2/configuration_molmoact2.py index bf9437ba9..5327325a3 100644 --- a/src/lerobot/policies/molmoact2/configuration_molmoact2.py +++ b/src/lerobot/policies/molmoact2/configuration_molmoact2.py @@ -73,6 +73,19 @@ class MolmoAct2Config(PreTrainedConfig): num_inference_steps: int | None = None mask_action_dim_padding: bool = True enable_inference_cuda_graph: bool = True + + # Language conditioning (e.g. RECAP advantage). When set, RenderMessagesStep + # resolves language_persistent rows via the recipe YAML. Same mechanism as PI05. + recipe_path: str | None = None + + # Inference-time advantage indicator (e.g. "Advantage: positive. "). + # Used during rollout when no language_persistent data is available. + # Placed after the user prompt, before action tokens. + advantage_prefix: str = "" + + # Classifier-Free Guidance (CFG) scale for inference (RECAP Eq. 13). + # 1.0 = no guidance. >1.0 = dual-path: v = v_uncond + beta * (v_cond - v_uncond) + cfg_beta: float = 1.0 # MolmoAct2-local eval option. When enabled, stochastic continuous action # generation uses a rollout-local generator derived from eval_seed. per_episode_seed: bool = False diff --git a/src/lerobot/policies/molmoact2/processor_molmoact2.py b/src/lerobot/policies/molmoact2/processor_molmoact2.py index d2db817ef..f80b022fa 100644 --- a/src/lerobot/policies/molmoact2/processor_molmoact2.py +++ b/src/lerobot/policies/molmoact2/processor_molmoact2.py @@ -359,6 +359,7 @@ def _build_robot_text( add_setup_tokens: bool, add_control_tokens: bool, num_images: int, + advantage: str = "", ) -> str: setup_text = _wrap_setup_text(setup_type, add_setup_tokens=add_setup_tokens) control_text = _wrap_control_text(control_mode, add_control_tokens=add_control_tokens) @@ -375,7 +376,10 @@ def _build_robot_text( image_prefix = "<|image|>" else: image_prefix = "".join(f"Image {idx + 1}<|image|>" for idx in range(num_images)) - return f"{image_prefix}<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n{ACTION_OUTPUT_TOKEN}" + # Per RECAP paper (Section V-B): advantage indicator goes after context, + # before actions, so only action log-likelihoods are affected. + advantage_clause = f"Advantage: {advantage}. " if advantage else "" + return f"{image_prefix}<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n{advantage_clause}{ACTION_OUTPUT_TOKEN}" def _as_text_list(value: Any, batch_size: int) -> list[str]: @@ -695,6 +699,39 @@ class MolmoAct2ClampNormalizedProcessorStep(ProcessorStep): return features +@ProcessorStepRegistry.register(name="molmoact2_normalize_task") +@dataclass +class MolmoAct2NormalizeTaskStep(ProcessorStep): + """Normalize the task text in complementary_data before recipe rendering. + + Ensures ${task} in recipe templates gets the same normalized form that + MolmoAct2PackInputsProcessorStep would produce, so training prompts + match inference prompts. + """ + + def __call__(self, transition: EnvTransition) -> EnvTransition: + complementary = transition.get(TransitionKey.COMPLEMENTARY_DATA) + if not isinstance(complementary, dict): + return transition + task = complementary.get("task") + if task is None: + return transition + + transition = transition.copy() + complementary = dict(complementary) + if isinstance(task, str): + complementary["task"] = _normalize_question_text(task) + elif isinstance(task, list): + complementary["task"] = [_normalize_question_text(t) for t in task] + transition[TransitionKey.COMPLEMENTARY_DATA] = complementary + return transition + + def transform_features( + self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]] + ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]: + return features + + @ProcessorStepRegistry.register(name="molmoact2_pack_inputs") @dataclass class MolmoAct2PackInputsProcessorStep(ProcessorStep): @@ -715,6 +752,8 @@ class MolmoAct2PackInputsProcessorStep(ProcessorStep): chunk_size: int = 30 max_action_dim: int = 32 env_action_dim: int | None = None + # RECAP: advantage indicator for inference (e.g. "Advantage: positive. ") + advantage_prefix: str = "" def __post_init__(self) -> None: require_package("transformers", extra="molmoact2") @@ -757,6 +796,7 @@ class MolmoAct2PackInputsProcessorStep(ProcessorStep): "chunk_size": self.chunk_size, "max_action_dim": self.max_action_dim, "env_action_dim": self.env_action_dim, + "advantage_prefix": self.advantage_prefix, } def _resolve_max_sequence_length( @@ -919,8 +959,40 @@ class MolmoAct2PackInputsProcessorStep(ProcessorStep): if task_source is None: task_source = complementary.get("language_instruction") tasks = _as_text_list(task_source, batch_size) - if self.normalize_language: - tasks = [_normalize_question_text(task) for task in tasks] + + # Resolve the advantage indicator. Per RECAP paper (Section V-B), it goes + # after all context but before actions — handled by _build_robot_text. + # Source priority: recipe-rendered "advantage" key > config advantage_prefix. + advantages: list[str] = [] + recipe_rendered = "base_task" in complementary + if recipe_rendered: + # Recipe rendered the task as " Advantage: ". + # Extract the advantage value and restore the clean task. + clean_tasks: list[str] = [] + for t in tasks: + if " Advantage: " in t: + split_idx = t.rindex(" Advantage: ") + clean_task = t[:split_idx] + adv = t[split_idx + len(" Advantage: ") :] + advantages.append(adv) + clean_tasks.append(clean_task) + else: + advantages.append("") + clean_tasks.append(t) + tasks = clean_tasks + else: + if self.normalize_language: + tasks = [_normalize_question_text(task) for task in tasks] + if self.advantage_prefix: + # Extract just the value from prefix like "Advantage: positive. " + prefix = self.advantage_prefix.strip() + if prefix.startswith("Advantage:"): + adv_val = prefix[len("Advantage:") :].strip().rstrip(".") + else: + adv_val = prefix + advantages = [adv_val] * batch_size + else: + advantages = [""] * batch_size complementary["task"] = tasks action_padded = None @@ -953,6 +1025,7 @@ class MolmoAct2PackInputsProcessorStep(ProcessorStep): add_setup_tokens=self.add_setup_tokens, add_control_tokens=self.add_control_tokens, num_images=len(images), + advantage=advantages[batch_idx], ) prompt_texts.append(prompt) if build_action_labels: @@ -1164,28 +1237,48 @@ def make_molmoact2_pre_post_processors( stats=masked_dataset_stats, ), MolmoAct2ClampNormalizedProcessorStep(normalization_masks=normalization_masks), - MolmoAct2PackInputsProcessorStep( - checkpoint_path=config.checkpoint_path, - checkpoint_revision=config.checkpoint_revision, - checkpoint_force_download=config.checkpoint_force_download, - action_mode=config.action_mode, - discrete_action_tokenizer=config.discrete_action_tokenizer, - image_keys=image_keys, - allow_image_key_fallback=not bool(config.image_keys), - setup_type=setup_type, - control_mode=control_mode, - normalize_language=config.normalize_language, - add_setup_tokens=config.add_setup_tokens, - add_control_tokens=config.add_control_tokens, - num_state_tokens=config.num_state_tokens, - max_sequence_length=config.max_sequence_length, - chunk_size=chunk_size, - max_action_dim=config.expected_max_action_dim, - env_action_dim=env_action_dim, - ), - DeviceProcessorStep(device=config.device), ] + # Insert language rendering steps when a recipe is configured (e.g. RECAP advantage) + if config.recipe_path is not None: + from lerobot.configs.recipe import load_recipe + from lerobot.processor.render_messages_processor import RenderMessagesStep + from lerobot.processor.rendered_messages_to_task import RenderedMessagesToTaskStep + + recipe = load_recipe(config.recipe_path) + # Normalize task text before recipe uses ${task}, ensuring consistency + # between training (recipe-rendered) and inference (advantage_prefix). + if config.normalize_language: + input_steps.append(MolmoAct2NormalizeTaskStep()) + input_steps.append(RenderMessagesStep(recipe=recipe)) + input_steps.append(RenderedMessagesToTaskStep()) + + input_steps.extend( + [ + MolmoAct2PackInputsProcessorStep( + checkpoint_path=config.checkpoint_path, + checkpoint_revision=config.checkpoint_revision, + checkpoint_force_download=config.checkpoint_force_download, + action_mode=config.action_mode, + discrete_action_tokenizer=config.discrete_action_tokenizer, + image_keys=image_keys, + allow_image_key_fallback=not bool(config.image_keys), + setup_type=setup_type, + control_mode=control_mode, + normalize_language=config.normalize_language, + add_setup_tokens=config.add_setup_tokens, + add_control_tokens=config.add_control_tokens, + num_state_tokens=config.num_state_tokens, + max_sequence_length=config.max_sequence_length, + chunk_size=chunk_size, + max_action_dim=config.expected_max_action_dim, + env_action_dim=env_action_dim, + advantage_prefix=config.advantage_prefix, + ), + DeviceProcessorStep(device=config.device), + ] + ) + output_steps: list[ProcessorStep] = [ MolmoAct2ClampActionProcessorStep(), MolmoAct2MaskedUnnormalizerProcessorStep(