inference/action: let forward dynamics run for per-dataset action widths

cosmos_framework/inference/action.py

@@ -14,7 +14,11 @@
     make_batched_action_processing_fields,
     pad_action_to_max_dim,
 )
-from cosmos_framework.data.generator.action.domain_utils import EMBODIMENT_TO_RAW_ACTION_DIM, get_domain_id
+from cosmos_framework.data.generator.action.domain_utils import (
+    EMBODIMENT_TO_DOMAIN_ID,
+    EMBODIMENT_TO_RAW_ACTION_DIM,
+    get_domain_id,
+)
 from cosmos_framework.data.generator.action.json_formatter import ActionPromptJsonFormatter
 from cosmos_framework.data.generator.action.transforms import (
     build_sequence_plan_from_mode,
@@ -25,28 +29,41 @@
 from cosmos_framework.inference.vision import read_media_frames
 from cosmos_framework.utils.generator.data_utils import get_vision_data_resolution
 
+# Domains whose raw action width is chosen per dataset at construction time rather than
+# being a property of the embodiment -- ``hand_pose`` varies with ``keypoint_option`` and
+# ``rotation_format``, ``libero`` with ``rotation_space``. They are absent from
+# ``EMBODIMENT_TO_RAW_ACTION_DIM`` for that reason, so forward dynamics has to take the
+# width from the action file instead of looking it up.
+_PER_DATASET_ACTION_WIDTH = frozenset({"hand_pose", "libero"})
+
 
 def _load_actions(
     action_path: Path | str | None,
     model_mode: ModelMode,
     action_chunk_size: int,
     max_action_dim: int,
     raw_action_dim: int | None,
-) -> torch.Tensor:
-    """Load actions from JSON (or zeros for policy mode and inverse dynamics mode). Returns padded action tensor."""
+) -> tuple[torch.Tensor, int]:
+    """Load actions from JSON (or zeros for policy mode and inverse dynamics mode).
+
+    Returns the padded action tensor and the resolved raw (unpadded) action width.
+    In forward-dynamics mode the width comes from the action file itself, so
+    ``raw_action_dim`` is only a cross-check and may be ``None`` for domains that
+    have no single canonical width (e.g. ``hand_pose``, ``libero``).
+    """
     match model_mode:
         case ModelMode.FORWARD_DYNAMICS:
             assert action_path is not None, "action_path is required for forward_dynamics mode"
             p = Path(str(action_path))
             raw = torch.tensor(json.loads(p.read_text()), dtype=torch.float32)
-            raw_dim = raw.shape[-1]
-            assert raw_dim == raw_action_dim, (
+            raw_dim = int(raw.shape[-1])
+            assert raw_action_dim is None or raw_dim == raw_action_dim, (
                 f"Raw action dimension from file ({raw_dim}) does not match expected dimension ({raw_action_dim})"
             )
-            return pad_action_to_max_dim(raw, max_action_dim)
+            return pad_action_to_max_dim(raw, max_action_dim), raw_dim
         case ModelMode.WAM | ModelMode.INVERSE_DYNAMICS:
             assert raw_action_dim is not None, "raw_action_dim is required for policy and inverse_dynamics modes"
-            return torch.zeros(action_chunk_size, max_action_dim, dtype=torch.float32)
+            return torch.zeros(action_chunk_size, max_action_dim, dtype=torch.float32), raw_action_dim
         case _:
             raise ValueError(f"Unsupported action model_mode: {model_mode}")
 
@@ -163,17 +180,27 @@ def get_action_sample_data(
 ) -> dict:
     """Load observation image/video + optional actions and build an Action inference batch."""
     domain_name = domain_name.lower().strip()
-    if domain_name not in EMBODIMENT_TO_RAW_ACTION_DIM:
+    if domain_name not in EMBODIMENT_TO_DOMAIN_ID:
         raise ValueError(
-            f"invalid domain_name {domain_name!r}; expected one of {sorted(EMBODIMENT_TO_RAW_ACTION_DIM.keys())}"
+            f"invalid domain_name {domain_name!r}; expected one of {sorted(EMBODIMENT_TO_DOMAIN_ID.keys())}"
         )
 
-    raw_action_dim = EMBODIMENT_TO_RAW_ACTION_DIM[domain_name]
+    raw_action_dim = EMBODIMENT_TO_RAW_ACTION_DIM.get(domain_name)
+    if raw_action_dim is None:
+        if domain_name not in _PER_DATASET_ACTION_WIDTH:
+            raise ValueError(
+                f"no raw action width registered for domain_name {domain_name!r}; domains with a "
+                f"canonical width are {sorted(EMBODIMENT_TO_RAW_ACTION_DIM.keys())}"
+            )
+        if model_mode is not ModelMode.FORWARD_DYNAMICS:
+            raise ValueError(
+                f"domain_name {domain_name!r} sizes its raw action per dataset, so {model_mode.value} "
+                f"inference is unsupported for it; only forward_dynamics can resolve the width, from "
+                f"the action file it is given"
+            )
+
     frames, _ = read_media_frames(Path(vision_path), max_frames=action_chunk_size + 1)
-    assert action_path is not None or raw_action_dim is not None, (
-        "Either action_path or raw_action_dim must be provided"
-    )
-    action = _load_actions(action_path, model_mode, action_chunk_size, max_action_dim, raw_action_dim)
+    action, raw_action_dim = _load_actions(action_path, model_mode, action_chunk_size, max_action_dim, raw_action_dim)
 
     return build_action_batch(
         video=frames,