[None][feat] Cosmos3 action generation

examples/visual_gen/models/cosmos3/README.md

@@ -8,6 +8,7 @@ Cosmos3 supports the following generation modes from a single checkpoint:
 - **V2V** — video-conditioned video (`prompts/v2v.json`). Condition on a reference video via `--video_path` (a local MP4/AVI file). Only the first (or last, per `condition_video_keep`) `max(condition_video_latent_indexes) * 4 + 1` input frames condition the output (5 by default); the encoded bytes pass through and each worker decodes just that window on NVDEC (see [Media I/O dependencies](#media-io-dependencies)). Validated for Nano / Super only.
 - **Transfer** — control-video conditioning (`edge`/`blur`/`depth`/`seg`/`wsm` hints via `--extra_params`). The control constrains structure frame by frame; the prompt supplies appearance. `edge` and `blur` are auto-computed from `--video_path`; any hint also accepts a precomputed control clip (`{"edge": "control.mp4"}` — the example reads it and sends encoded bytes, the same contract as the `video` reference). Multiple hints compose (each adds a full control-token copy of the video sequence); long videos run chunked (93 frames/chunk, stitched on overlap frames) — but only past the first chunk, so raise `num_frames` above the pipeline default to generate one: it bounds how many frames are decoded from the inputs, and so how long the output can be. A single-hint request picks up that hint's tuned sampling preset — guidance scale, control guidance and flow shift — for any of those the request leaves unset; requests with several hints fall back to the generic video defaults. The active hint names are also appended to the prompt as a one-sentence control-adherence directive; pass `"emphasize_control_in_prompt": false` to suppress it for clean baselines or ablations.
 - **T2AV** — text-to-video with synchronized audio (`prompts/t2av.json` with `enable_audio: true`, or pass `--enable_audio`). Combine with a `vision_path` for image-conditioned audio-video (TI2AV).
+- **Action** — policy / forward dynamics / inverse dynamics generation (pass `--action_mode`); `inverse_dynamics` reads its observation clip from `--video_path` (MP4/AVI, decoded on worker NVDEC like V2V). Action and audio generation are mutually exclusive. A predicted trajectory has no representation in a video container, so action runs are saved as `safetensors` or `pt`, keeping the rollout and the action tensor in one payload — over `trtllm-serve` the default `format=auto` selects that payload automatically, and an explicit `mp4`/`avi` is rejected.
 
 ## Checkpoints
 
@@ -169,4 +170,36 @@ python cosmos3.py --model nvidia/Cosmos3-Edge \
 python cosmos3.py --model nvidia/Cosmos3-Nano \
     --prompt "A cute puppy playing with a ball in a park" \
     --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml
+
+# Action — policy (first frame + instruction -> predicted action + rollout video)
+python cosmos3.py --model nvidia/Cosmos3-Nano \
+    --prompt_file prompts/action_policy.json \
+    --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml \
+    --action_mode policy \
+    --domain_name bridge_orig_lerobot \
+    --raw_action_dim 10 \
+    --output_path policy_rollout.safetensors \
+    --action_output_path policy_action.json
+
+# Action — forward dynamics (first frame + action trajectory -> rollout video)
+# action_trajectory.json is a [T, D] list of lists; D is the embodiment's action
+# width (9 for av) and a mismatch is rejected.
+python cosmos3.py --model nvidia/Cosmos3-Nano \
+    --prompt_file prompts/action_forward_dynamics.json \
+    --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml \
+    --action_mode forward_dynamics \
+    --domain_name av \
+    --action_json action_trajectory.json \
+    --output_path forward_dynamics.safetensors
+
+# Action — inverse dynamics (video -> predicted action)
+python cosmos3.py --model nvidia/Cosmos3-Nano \
+    --prompt_file prompts/action_inverse_dynamics.json \
+    --video_path /path/to/observation_clip.mp4 \
+    --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml \
+    --action_mode inverse_dynamics \
+    --domain_name bridge_orig_lerobot \
+    --raw_action_dim 10 \
+    --output_path inverse_video.safetensors \
+    --action_output_path inverse_action.json
 ```

examples/visual_gen/models/cosmos3/cosmos3.py

@@ -31,6 +31,8 @@
 from tensorrt_llm._torch.visual_gen.models.cosmos3.transfer import TRANSFER_HINT_KEYS
 
 _SCRIPT_DIR = Path(__file__).resolve().parent
+_ACTION_MODES = ("policy", "forward_dynamics", "inverse_dynamics")
+_TENSOR_OUTPUT_SUFFIXES = {".pt", ".safetensors"}
 
 DEFAULT_PROMPT_FILE = "prompts/t2v.json"
 DEFAULT_NEGATIVE_PROMPT_FILE = "cosmos3_negative_prompt.json"
@@ -199,6 +201,95 @@ def resolve_prompt_and_options(
     return resolved_prompt, resolved_image, resolved_enable_audio, resolved_output_type
 
 
+def _validate_action_args(
+    args: argparse.Namespace, resolved_image_path: Optional[str] = None
+) -> None:
+    if args.action_mode is None:
+        return
+
+    # The first frame may come from --image_path or a prompt file's vision_path.
+    has_first_frame = resolved_image_path is not None or args.video_path is not None
+
+    mode = args.action_mode.strip().lower()
+    if mode not in _ACTION_MODES:
+        raise SystemExit(
+            f"Invalid --action_mode {args.action_mode!r}; expected one of {list(_ACTION_MODES)}."
+        )
+    args.action_mode = mode
+    if args.enable_audio:
+        raise SystemExit("Cosmos3 does not support joint action and audio generation.")
+    if args.output_type != "video":
+        raise SystemExit("Action generation requires --output_type video.")
+
+    if mode == "forward_dynamics":
+        if args.action_json is None:
+            raise SystemExit(f"{mode} requires --action_json.")
+        if not has_first_frame:
+            raise SystemExit(
+                f"{mode} requires --image_path, a prompt-file vision_path, or --video_path "
+                "for the first frame."
+            )
+    elif mode == "policy":
+        if not has_first_frame:
+            raise SystemExit(
+                f"{mode} requires --image_path, a prompt-file vision_path, or --video_path "
+                "for the first frame."
+            )
+        if args.raw_action_dim is None and args.domain_name is None and args.domain_id is None:
+            raise SystemExit(f"{mode} requires --raw_action_dim, --domain_name, or --domain_id.")
+    elif mode == "inverse_dynamics":
+        if args.video_path is None:
+            raise SystemExit(f"{mode} requires --video_path (an .mp4 or .avi file).")
+        if args.raw_action_dim is None and args.domain_name is None and args.domain_id is None:
+            raise SystemExit(f"{mode} requires --raw_action_dim, --domain_name, or --domain_id.")
+
+
+def _resolved_output_path(path: str, action_mode: Optional[str]) -> str:
+    if action_mode is None:
+        return path
+    output_path = Path(path)
+    if output_path.suffix.lower() in _TENSOR_OUTPUT_SUFFIXES:
+        return str(output_path)
+    return str(output_path.with_suffix(".safetensors"))
+
+
+def _default_action_output_path(output_path: str) -> str:
+    stem = Path(output_path)
+    return str(stem.with_suffix(".action.json"))
+
+
+def _save_action_output(output, path: str, args: argparse.Namespace) -> None:
+    """Write the trajectory plus the request that produced it.
+
+    The mode and embodiment are this script's own inputs, so they are read
+    from *args* rather than echoed back through the output schema.
+    """
+    if output.action is None:
+        return
+
+    action = output.action
+    if action.ndim == 3 and action.shape[0] == 1:
+        action_data = action[0].tolist()
+        shape = list(action.shape[1:])
+    else:
+        action_data = action.tolist()
+        shape = list(action.shape)
+
+    payload = {
+        "action_mode": args.action_mode,
+        "domain_name": args.domain_name,
+        "domain_id": args.domain_id,
+        "raw_action_dim": action.shape[-1],
+        "shape": shape,
+        "dtype": str(action.dtype).replace("torch.", ""),
+        "data": action_data,
+    }
+    out_path = Path(path)
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    with out_path.open("w", encoding="utf-8") as f:
+        json.dump(payload, f, indent=2)
+
+
 def main():
     parser = argparse.ArgumentParser(description="Cosmos3 Text(+Image)-to-Video(+Audio) example")
     parser.add_argument(
@@ -274,11 +365,77 @@ def main():
         ),
     )
     parser.add_argument("--enable_audio", action="store_true", help="Enable audio generation")
+    parser.add_argument(
+        "--action_mode",
+        type=str,
+        default=None,
+        choices=list(_ACTION_MODES),
+        help="Action mode: policy, forward_dynamics, or inverse_dynamics",
+    )
+    parser.add_argument(
+        "--domain_name",
+        type=str,
+        default=None,
+        help="Embodiment domain name (e.g. bridge_orig_lerobot, av, droid_lerobot)",
+    )
+    parser.add_argument(
+        "--domain_id",
+        type=int,
+        default=None,
+        help="Embodiment domain id (alternative to --domain_name)",
+    )
+    parser.add_argument(
+        "--raw_action_dim",
+        type=int,
+        default=None,
+        help="Raw action DOF for policy/inverse_dynamics",
+    )
+    parser.add_argument(
+        "--action_chunk_size",
+        type=int,
+        default=None,
+        help="Action tokens to generate. Defaults to the domain preset or model default.",
+    )
+    parser.add_argument(
+        "--action_json",
+        type=str,
+        default=None,
+        help="JSON file with action trajectory [T, D] for forward_dynamics",
+    )
     parser.add_argument(
         "--video_path",
         type=str,
         default=None,
-        help="Reference video for V2V: a local MP4/AVI file (decoded on worker NVDEC)",
+        help=(
+            "Reference video (MP4/AVI, decoded on worker NVDEC): V2V conditioning, "
+            "or the observation clip for action inverse_dynamics"
+        ),
+    )
+    parser.add_argument(
+        "--action_resolution",
+        type=int,
+        default=None,
+        choices=[256, 480, 704, 720],
+        help=("Resolution bucket for action image sizing. Defaults to the domain preset or 480."),
+    )
+    parser.add_argument(
+        "--action_fps",
+        type=float,
+        default=None,
+        help="Action-token temporal rate for mRoPE (Hz). Defaults to frame_rate.",
+    )
+    parser.add_argument(
+        "--view_point",
+        type=str,
+        default=None,
+        choices=["ego_view", "third_person_view", "wrist_view", "concat_view"],
+        help="Camera perspective for the action caption (default: ego_view).",
+    )
+    parser.add_argument(
+        "--action_output_path",
+        type=str,
+        default=None,
+        help="Path to save predicted action JSON (default: <output_stem>.action.json)",
     )
     parser.add_argument(
         "--output_type", type=str, default="video", help="Output type (video, image)"
@@ -310,6 +467,7 @@ def main():
         enable_audio=args.enable_audio,
         output_type=args.output_type,
     )
+    _validate_action_args(args, resolved_image_path=image_path)
 
     # Engine config from shared YAML (optional); model-specific defaults apply otherwise.
     extra_args = VisualGenArgs.from_yaml(args.visual_gen_args) if args.visual_gen_args else None
@@ -336,6 +494,25 @@ def main():
     params.extra_params["use_guardrails"] = not args.disable_guardrails
     params.extra_params["output_type"] = output_type
 
+    if args.action_mode is not None:
+        params.extra_params["action_mode"] = args.action_mode
+    if args.domain_name is not None:
+        params.extra_params["domain_name"] = args.domain_name
+    if args.domain_id is not None:
+        params.extra_params["domain_id"] = args.domain_id
+    if args.raw_action_dim is not None:
+        params.extra_params["raw_action_dim"] = args.raw_action_dim
+    if args.action_chunk_size is not None:
+        params.extra_params["action_chunk_size"] = args.action_chunk_size
+    if args.action_resolution is not None:
+        params.extra_params["action_resolution"] = args.action_resolution
+    if args.action_fps is not None:
+        params.extra_params["action_fps"] = args.action_fps
+    if args.view_point is not None:
+        params.extra_params["view_point"] = args.view_point
+    if args.action_json is not None:
+        with open(args.action_json, encoding="utf-8") as f:
+            params.extra_params["action"] = json.load(f)
     if args.video_path is not None:
         params.extra_params["video"] = Path(args.video_path).read_bytes()
     if args.extra_params:
@@ -352,8 +529,18 @@ def main():
         params=params,
     )
 
-    output.save(args.output_path)
-    print(f"Saved: {args.output_path}")
+    output_path = _resolved_output_path(args.output_path, args.action_mode)
+    output.save(output_path)
+    print(f"Saved: {output_path}")
+
+    if args.action_mode is not None:
+        action_path = args.action_output_path or _default_action_output_path(output_path)
+        _save_action_output(output, action_path, args)
+        if output.action is not None:
+            print(f"Saved action: {action_path}")
+            print(f"Action shape: {tuple(output.action.shape)}")
+        else:
+            print("Warning: action_mode was set but the output carried no action tensor.")
 
     print(output.metrics)
 

examples/visual_gen/models/cosmos3/prompts/action_forward_dynamics.json

@@ -0,0 +1,5 @@
+{
+    "model_mode": "image2video",
+    "prompt": "Robot manipulation rollout: the right arm extends to the fruit display, picks up a pear, and places it into the bag in the shopping cart.",
+    "vision_path": "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/vision/robot_153.jpg"
+}

examples/visual_gen/models/cosmos3/prompts/action_inverse_dynamics.json

@@ -0,0 +1,4 @@
+{
+    "model_mode": "image2video",
+    "prompt": "Recover the robot action trajectory from this clip of the arm picking up a pear and placing it into the bag."
+}

examples/visual_gen/models/cosmos3/prompts/action_policy.json

@@ -0,0 +1,5 @@
+{
+    "model_mode": "image2video",
+    "prompt": "Pick up the pear from the fruit display and place it into the plastic bag in the shopping cart.",
+    "vision_path": "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/vision/robot_153.jpg"
+}