Update MolmoAct2 action mode inference API

config.json

@@ -4,7 +4,7 @@
   "action_expert_depth_gate": false,
   "action_expert_depth_gate_init_bias": -4.0,
   "action_expert_depth_gate_per_layer": false,
-  "action_format": "discrete",
+  "action_mode": "discrete",
   "max_action_horizon": 30,
   "action_output_token_id": 151931,
   "action_start_token_id": 151932,

configuration_molmoact2.py

@@ -375,7 +375,7 @@ class MolmoAct2Config(PretrainedConfig):
         max_action_dim: int = 32,
         max_action_horizon: int = 30,
         n_obs_steps: int = 30,
-        action_format: str = "both",
+        action_mode: str = "both",
         state_format: str = "discrete",
         flow_matching_num_steps: int = 10,
         flow_matching_cutoff: float = 1.0,
@@ -461,7 +461,7 @@ class MolmoAct2Config(PretrainedConfig):
         self.max_action_dim = max_action_dim
         self.max_action_horizon = max_action_horizon
         self.n_obs_steps = n_obs_steps
-        self.action_format = action_format
+        self.action_mode = action_mode
         self.state_format = state_format
         self.flow_matching_num_steps = flow_matching_num_steps
         self.flow_matching_cutoff = flow_matching_cutoff

modeling_molmoact2.py

@@ -2949,7 +2949,7 @@ class MolmoAct2Model(MolmoAct2PreTrainedModel):
             mask = input_ids != -1
         else:
             return None
-        if self.config.action_format != "both" or input_ids is None:
+        if self.config.action_mode != "both" or input_ids is None:
             return mask
         eos_id = getattr(self.config, "eos_token_id", None)
         if eos_id is not None:
@@ -4452,7 +4452,7 @@ class MolmoAct2ForConditionalGeneration(MolmoAct2PreTrainedModel, GenerationMixi
     ) -> torch.Tensor:
         if action_tokenizer is None:
             raise ValueError(
-                "action_mode='discrete' requires an `action_tokenizer` input."
+                "inference_action_mode='discrete' requires an `action_tokenizer` input."
             )
         if (
             self.config.action_start_token_id is None
@@ -4508,7 +4508,7 @@ class MolmoAct2ForConditionalGeneration(MolmoAct2PreTrainedModel, GenerationMixi
         task: str,
         state: Any,
         norm_tag: str,
-        action_mode: str = "continuous",
+        inference_action_mode: Optional[str] = None,
         enable_depth_reasoning: bool = False,
         enable_adaptive_depth: bool = True,
         depth_cache: Optional[Mapping[str, Any]] = None,
@@ -4524,31 +4524,36 @@ class MolmoAct2ForConditionalGeneration(MolmoAct2PreTrainedModel, GenerationMixi
             raise ValueError(
                 "MolmoAct2 `predict_action` requires `state` for discrete state prompting."
             )
-        action_mode = str(action_mode or "continuous")
-        if action_mode not in {"continuous", "discrete"}:
-            raise ValueError("action_mode must be either 'continuous' or 'discrete'.")
-        if action_mode == "continuous" and not bool(self.config.add_action_expert):
+        if inference_action_mode is None:
+            raise ValueError(
+                "`inference_action_mode` must be provided explicitly as either "
+                "'continuous' or 'discrete'."
+            )
+        inference_action_mode = str(inference_action_mode)
+        if inference_action_mode not in {"continuous", "discrete"}:
+            raise ValueError("inference_action_mode must be either 'continuous' or 'discrete'.")
+        if inference_action_mode == "continuous" and not bool(self.config.add_action_expert):
             raise RuntimeError(
-                "action_mode='continuous' requires an action expert, but this checkpoint "
+                "inference_action_mode='continuous' requires an action expert, but this checkpoint "
                 "was converted with add_action_expert=False."
             )
-        if action_mode == "continuous" and self.config.action_format not in {
+        if inference_action_mode == "continuous" and self.config.action_mode not in {
             "continuous",
             "both",
         }:
             raise ValueError(
-                f"action_mode='continuous' requires checkpoint action_format in {{'continuous', 'both'}}, "
-                f"got {self.config.action_format!r}."
+                "inference_action_mode='continuous' requires checkpoint action_mode in "
+                f"{{'continuous', 'both'}}, got {self.config.action_mode!r}."
             )
-        if action_mode == "discrete":
+        if inference_action_mode == "discrete":
             if action_tokenizer is None:
                 raise ValueError(
-                    "action_mode='discrete' requires an `action_tokenizer` input."
+                    "inference_action_mode='discrete' requires an `action_tokenizer` input."
                 )
-            if self.config.action_format not in {"discrete", "both"}:
+            if self.config.action_mode not in {"discrete", "both"}:
                 raise ValueError(
-                    f"action_mode='discrete' requires checkpoint action_format in {{'discrete', 'both'}}, "
-                    f"got {self.config.action_format!r}."
+                    "inference_action_mode='discrete' requires checkpoint action_mode in "
+                    f"{{'discrete', 'both'}}, got {self.config.action_mode!r}."
                 )
         if enable_depth_reasoning and not bool(self.config.enable_depth_reasoning):
             raise ValueError(
@@ -4625,7 +4630,7 @@ class MolmoAct2ForConditionalGeneration(MolmoAct2PreTrainedModel, GenerationMixi
         generated_token_ids = None
         depth_bins = None
         updated_depth_cache = depth_cache
-        if action_mode == "continuous":
+        if inference_action_mode == "continuous":
             if enable_depth_reasoning:
                 latest_first_image = _extract_first_image(images)
                 depth_prefix = self._generate_depth_prefix(