From e1afb96474368ad268e4a824b7d191fd791044cc Mon Sep 17 00:00:00 2001 From: hq-fang <71356829+hq-fang@users.noreply.github.com> Date: Tue, 19 May 2026 22:04:34 +0000 Subject: [PATCH] fix molmoact2 hf image key resolution --- .../molmoact2/configuration_molmoact2.py | 2 - .../policies/molmoact2/processor_molmoact2.py | 28 +++++++-- tests/policies/molmoact2/test_molmoact2.py | 61 ++++++++++++++++++- 3 files changed, 83 insertions(+), 8 deletions(-) diff --git a/src/lerobot/policies/molmoact2/configuration_molmoact2.py b/src/lerobot/policies/molmoact2/configuration_molmoact2.py index 409385ad2..a97170e31 100644 --- a/src/lerobot/policies/molmoact2/configuration_molmoact2.py +++ b/src/lerobot/policies/molmoact2/configuration_molmoact2.py @@ -444,8 +444,6 @@ class MolmoAct2Config(PreTrainedConfig): self.setup_type = str(metadata["setup_type"]) if not self.control_mode and metadata.get("control_mode") is not None: self.control_mode = str(metadata["control_mode"]) - if not self.image_keys and isinstance(metadata.get("camera_keys"), list): - self.image_keys = [str(key) for key in metadata["camera_keys"]] def saved_policy_action_mode(self) -> str | None: pretrained_path = getattr(self, "pretrained_path", None) diff --git a/src/lerobot/policies/molmoact2/processor_molmoact2.py b/src/lerobot/policies/molmoact2/processor_molmoact2.py index 94fd2c5cb..f5e037fa1 100644 --- a/src/lerobot/policies/molmoact2/processor_molmoact2.py +++ b/src/lerobot/policies/molmoact2/processor_molmoact2.py @@ -30,7 +30,7 @@ import torch from huggingface_hub import snapshot_download from torch import Tensor -from lerobot.configs import PipelineFeatureType, PolicyFeature +from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature from lerobot.processor import ( AddBatchDimensionProcessorStep, DeviceProcessorStep, @@ -501,6 +501,7 @@ class MolmoAct2PackInputsProcessorStep(ProcessorStep): action_mode: str = "both" discrete_action_tokenizer: str = "allenai/MolmoAct2-FAST-Tokenizer" image_keys: list[str] = field(default_factory=list) + allow_image_key_fallback: bool = False setup_type: str = "" control_mode: str = "" normalize_language: bool = True @@ -547,6 +548,7 @@ class MolmoAct2PackInputsProcessorStep(ProcessorStep): "action_mode": self.action_mode, "discrete_action_tokenizer": self.discrete_action_tokenizer, "image_keys": list(self.image_keys), + "allow_image_key_fallback": self.allow_image_key_fallback, "setup_type": self.setup_type, "control_mode": self.control_mode, "normalize_language": self.normalize_language, @@ -590,15 +592,23 @@ class MolmoAct2PackInputsProcessorStep(ProcessorStep): return int(value.shape[0]) if getattr(value, "ndim", 0) == 4 else 1 return 1 + @staticmethod + def _observation_image_keys(observation: dict[str, Any]) -> list[str]: + keys = [key for key in observation if str(key).startswith(f"{OBS_IMAGES}.")] + if not keys: + keys = [key for key in observation if str(key).startswith("observation.image")] + return sorted(keys) + def _resolve_image_keys(self, observation: dict[str, Any]) -> list[str]: if self.image_keys: missing = [key for key in self.image_keys if key not in observation] if missing: + fallback_keys = self._observation_image_keys(observation) + if self.allow_image_key_fallback and fallback_keys: + return fallback_keys raise ValueError(f"MolmoAct2 image_keys missing from observation: {missing}.") return list(self.image_keys) - keys = [key for key in observation if str(key).startswith(f"{OBS_IMAGES}.")] - if not keys: - keys = [key for key in observation if str(key).startswith("observation.image")] + keys = self._observation_image_keys(observation) if not keys: raise ValueError("MolmoAct2 requires at least one image observation.") return sorted(keys) @@ -835,8 +845,15 @@ def make_molmoact2_pre_post_processors( ) image_keys = list(config.image_keys) + visual_feature_keys = [ + key for key, feature in config.input_features.items() if feature.type == FeatureType.VISUAL + ] if not image_keys and isinstance(hf_metadata.get("camera_keys"), list): - image_keys = [str(key) for key in hf_metadata["camera_keys"]] + metadata_image_keys = [str(key) for key in hf_metadata["camera_keys"]] + if not visual_feature_keys or all(key in config.input_features for key in metadata_image_keys): + image_keys = metadata_image_keys + if not image_keys: + image_keys = visual_feature_keys setup_type = config.setup_type or str(hf_metadata.get("setup_type") or "") control_mode = config.control_mode or str(hf_metadata.get("control_mode") or "") chunk_size = int(hf_metadata.get("action_horizon") or config.chunk_size) @@ -865,6 +882,7 @@ def make_molmoact2_pre_post_processors( 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, diff --git a/tests/policies/molmoact2/test_molmoact2.py b/tests/policies/molmoact2/test_molmoact2.py index 23925b074..800ca2594 100644 --- a/tests/policies/molmoact2/test_molmoact2.py +++ b/tests/policies/molmoact2/test_molmoact2.py @@ -27,7 +27,10 @@ import torch.nn.functional as F # noqa: N812 from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature from lerobot.policies import get_policy_class, make_policy_config -from lerobot.policies.molmoact2 import modeling_molmoact2 as molmoact2_modeling +from lerobot.policies.molmoact2 import ( + modeling_molmoact2 as molmoact2_modeling, + processor_molmoact2 as molmoact2_processor, +) from lerobot.policies.molmoact2.configuration_molmoact2 import ( MolmoAct2Config, MolmoAct2CosineDecayWithWarmupSchedulerConfig, @@ -41,6 +44,7 @@ from lerobot.policies.molmoact2.processor_molmoact2 import ( _add_gripper_masks_to_stats, _build_discrete_state_string, _normalize_question_text, + make_molmoact2_pre_post_processors, ) from lerobot.policies.rtc.configuration_rtc import RTCConfig from lerobot.types import TransitionKey @@ -234,6 +238,61 @@ def test_molmoact2_uses_config_feature_names_without_dataset_meta(): assert masked_stats[ACTION]["mask"] == [True, False] +def test_molmoact2_processor_uses_available_visual_features_over_missing_metadata_keys(monkeypatch): + monkeypatch.setattr( + molmoact2_processor, + "_load_hf_norm_stats_for_tag", + lambda *args, **kwargs: ( + {}, + {"camera_keys": ["observation.images.image", "observation.images.wrist_image"]}, + ), + ) + monkeypatch.setattr(MolmoAct2PackInputsProcessorStep, "__post_init__", lambda self: None) + cfg = MolmoAct2Config( + checkpoint_path="/tmp/not-a-real-checkpoint", + norm_tag="libero", + input_features={ + "observation.images.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 224, 224)), + "observation.images.image2": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 224, 224)), + OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(7,)), + }, + output_features={ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(7,))}, + ) + + preprocessor, _ = make_molmoact2_pre_post_processors(cfg) + pack_step = next( + step for step in preprocessor.steps if isinstance(step, MolmoAct2PackInputsProcessorStep) + ) + + assert pack_step.image_keys == ["observation.images.image", "observation.images.image2"] + assert pack_step.allow_image_key_fallback is True + + +def test_molmoact2_metadata_image_keys_can_fall_back_to_observation_keys(): + step = object.__new__(MolmoAct2PackInputsProcessorStep) + step.image_keys = ["observation.images.image", "observation.images.wrist_image"] + step.allow_image_key_fallback = True + observation = { + "observation.images.image": torch.zeros(3, 4, 4), + "observation.images.image2": torch.zeros(3, 4, 4), + } + + assert step._resolve_image_keys(observation) == ["observation.images.image", "observation.images.image2"] + + +def test_molmoact2_explicit_image_keys_stay_strict(): + step = object.__new__(MolmoAct2PackInputsProcessorStep) + step.image_keys = ["observation.images.image", "observation.images.wrist_image"] + step.allow_image_key_fallback = False + observation = { + "observation.images.image": torch.zeros(3, 4, 4), + "observation.images.image2": torch.zeros(3, 4, 4), + } + + with pytest.raises(ValueError, match="wrist_image"): + step._resolve_image_keys(observation) + + def test_enable_lora_vlm_builds_policy_local_peft_config(): pytest.importorskip("peft") policy_cfg = MolmoAct2Config(