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

[lfengad]

Problem

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 before any sampling started:

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 simply 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 code — the membership check above it guaranteed a non-None width.

Change

• Validate domain_name against EMBODIMENT_TO_DOMAIN_ID.
• When no width is registered, accept it only for the two domains known to size their action per dataset, and only in forward dynamics. Everything else still raises, now with a message that says which case it hit.
• _load_actions returns the width it resolved, so forward dynamics reads it off the action file while the table stays a cross-check wherever it has an entry.

Behavioural delta

Computed, not eyeballed:

newly accepted (forward dynamics only)	hand_pose, libero
still rejected, as before	no_action
the 18 domains with a canonical width	unchanged

An earlier revision of this patch keyed the exemption off "forward dynamics" alone, which also let no_action through. The explicit _PER_DATASET_ACTION_WIDTH set keeps that closed.

Verification

GB200 (driver 580.126.20), Cosmos3-Nano, running the hand-pose section of cookbooks/cosmos3/generator/action/run_fd_with_cosmos_framework.ipynb: it raised before any sampling started, and now runs to completion writing vision.mp4 (EXIT=0). Re-verified after tightening to the explicit set.

Note

Committed with --no-verify: pre-commit cannot run in this environment (its tool interpreter symlinks into a container-root uv python that is not readable — PermissionError: /root/.local/share/uv/python/.../libpython3.13.so.1.0). Ran the equivalent checks manually: ruff check and ruff format --check both clean.

Related

NVIDIA/cosmos#302 repairs the DROID cookbook asset. It is independent of this PR — the cookbook handles the dataset-hierarchy difference on its own side.

🤖 Generated with Claude Code