From c092194cf27bc1ecea0e0ab5b1b733ef9016407c Mon Sep 17 00:00:00 2001
From: hq-fang <71356829+hq-fang@users.noreply.github.com>
Date: Tue, 19 May 2026 21:58:06 +0000
Subject: [PATCH] align molmoact2 feature validation with eo pattern

---
 .../molmoact2/configuration_molmoact2.py      | 21 +++++++++++++++++--
 1 file changed, 19 insertions(+), 2 deletions(-)

diff --git a/src/lerobot/policies/molmoact2/configuration_molmoact2.py b/src/lerobot/policies/molmoact2/configuration_molmoact2.py
index 362b5ffc9..7d07eb0db 100644
--- a/src/lerobot/policies/molmoact2/configuration_molmoact2.py
+++ b/src/lerobot/policies/molmoact2/configuration_molmoact2.py
@@ -334,7 +334,24 @@ class MolmoAct2Config(PreTrainedConfig):
                 self.dataset_feature_names[key] = feature["names"]
 
     def validate_features(self) -> None:
+        """Validate and set up MolmoAct2 input and output features."""
+        image_features = [key for key, feat in self.input_features.items() if feat.type == FeatureType.VISUAL]
+        if not image_features:
+            raise ValueError(
+                "MolmoAct2 policy requires at least one visual input feature. "
+                "No features of type FeatureType.VISUAL found in input_features."
+            )
+
         if OBS_STATE not in self.input_features:
-            self.input_features[OBS_STATE] = PolicyFeature(type=FeatureType.STATE, shape=(0,))
+            state_feature = PolicyFeature(
+                type=FeatureType.STATE,
+                shape=(0,),
+            )
+            self.input_features[OBS_STATE] = state_feature
+
         if ACTION not in self.output_features:
-            self.output_features[ACTION] = PolicyFeature(type=FeatureType.ACTION, shape=(0,))
+            action_feature = PolicyFeature(
+                type=FeatureType.ACTION,
+                shape=(self.expected_max_action_dim,),
+            )
+            self.output_features[ACTION] = action_feature