From 7540a0cfafccdbdd21074bf006b575160f14a566 Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Tue, 4 Aug 2026 01:22:44 -0700 Subject: [PATCH] inference/action: let forward dynamics run for per-dataset action widths `get_action_sample_data` validated `domain_name` against `EMBODIMENT_TO_RAW_ACTION_DIM`. That table is a width lookup, not the list of valid domains -- `EMBODIMENT_TO_DOMAIN_ID` is, as the `get_domain_id(domain_name)` call a few lines down shows. And it deliberately omits `hand_pose` and `libero`, whose raw width is chosen per dataset (`keypoint_option` / `rotation_format`, `rotation_space`). So every hand-pose and LIBERO run died with: ValueError: invalid domain_name 'hand_pose'; expected one of ['abc_yam', 'agibotworld', ...] including forward dynamics, which never needed the lookup: the caller supplies an action file and the raw width is its last dimension. The table's own comment says only inverse_dynamics and WAM are unsupported for these domains, and the next line agreed: assert action_path is not None or raw_action_dim is not None That assert was already dead -- the membership check above it guaranteed a non-None width. Validate against `EMBODIMENT_TO_DOMAIN_ID`, and when no width is registered, accept it only for the two domains that are known to size their action per dataset, and only in forward dynamics. `_load_actions` now returns the width it resolved, so forward dynamics reads it off the action file while keeping the table as a cross-check wherever it has an entry. Behavioural delta is exactly: newly accepted (forward dynamics only): hand_pose, libero still rejected, as before: no_action the 18 domains with a canonical width: unchanged Verified on GB200 with Cosmos3-Nano: the action forward-dynamics cookbook's hand-pose section raised before any sampling started and now runs to completion. Co-Authored-By: Claude Opus 5 (1M context) --- cosmos_framework/inference/action.py | 55 +++++++++++++++++++++------- 1 file changed, 41 insertions(+), 14 deletions(-) diff --git a/cosmos_framework/inference/action.py b/cosmos_framework/inference/action.py index 69257f46..efb02641 100644 --- a/cosmos_framework/inference/action.py +++ b/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,6 +29,13 @@ 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, @@ -32,21 +43,27 @@ def _load_actions( 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,