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,