fix molmoact2 hf image key resolution

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)

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}.")]
+        keys = self._observation_image_keys(observation)
-        if not keys:
-            keys = [key for key in observation if str(key).startswith("observation.image")]
         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,

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(