[feat] cosmos3 PR3: action (domain-aware) DiT pathway + packing + CFG glue

Add the action (multi-embodiment world-model) modality. The Nano checkpoint
ships action weights (action_gen=true, 32 embodiment domains); this activates
them, bit-exact vs the framework.

- DiT forward: action encode/decode mirroring sound but domain-aware — pack each
  [T, action_dim] latent with a per-token embodiment domain id, `action_proj_in`
  (= framework `action2llm`, a `DomainAwareLinear`) + `action_modality_embed` +
  timestep scatter; decode via `action_proj_out` (= `llm2action`) on noisy hidden
  states -> `_unpack_action`. Mirrors `_encode_action`/`_decode_action`.
- Sound/action share the vision "full" split; sequence/flat order is
  [vision | action | sound] (matches the framework per-sample concat).
- Packing: `Cosmos3ActionItem` + action fields. Action tokens are `(T,)`-shaped,
  domain-tagged, with 3D-MRoPE at the vision temporal offset and
  `start_frame_offset=1` (framework `_pack_action_tokens`).
- CFG velocity / `Cosmos3DenoiseEngine` take optional `action_specs`
  (`Cosmos3ActionSpec`); the joint denoise machinery now handles action.

Parity (`test_cosmos3_action_parity.py`, framework model+pack = oracle, 3
embodiment domains):
- action packing fields + `position_ids` exact;
- `preds_vision` + `preds_action` (forward): max abs diff = 0.0, mean = 0.0;
- combined [vision|action] CFG velocity: max abs diff = 0.0, mean = 0.0.

Action is processed bit-exactly at the model + packing + denoise-glue level. A
full action2world real-weights run is gated on real robot-action input data (not
available here); the bit-exact parity is the correctness proof. Full suite: 139
passed, 0 skipped.

fastvideo/models/dits/cosmos3.py

@@ -808,6 +808,13 @@ class Cosmos3VFMTransformer(BaseDiT):
         sound_mse_loss_indexes: torch.Tensor | None = None,
         sound_noisy_frame_indexes: list[torch.Tensor] | None = None,
         fps_sound: torch.Tensor | None = None,
+        action_tokens: list[torch.Tensor] | None = None,
+        action_token_shapes: list[tuple[int, ...]] | None = None,
+        action_sequence_indexes: torch.Tensor | None = None,
+        action_timesteps: torch.Tensor | None = None,
+        action_mse_loss_indexes: torch.Tensor | None = None,
+        action_noisy_frame_indexes: list[torch.Tensor] | None = None,
+        action_domain_id: list[torch.Tensor] | None = None,
         **kwargs: Any,
     ) -> dict[str, Any]:
         """Video-path forward returning ``{"last_hidden_state", "preds_vision"}``
@@ -863,6 +870,26 @@ class Cosmos3VFMTransformer(BaseDiT):
                                                              sound_token_shapes)
             packed_sequence[sound_sequence_indexes] = packed_sound
 
+        # 2c. Action (DomainAwareLinear): pack [T, D] per sample with a per-token
+        #     domain id -> action_proj_in(domain) + action_modality_embed ->
+        #     timestep scatter -> scatter into the gen split (framework
+        #     ``_encode_action``). Action is domain-conditioned (per embodiment).
+        if action_tokens:
+            packed_action = torch.cat(
+                [a[:shp[0]] for a, shp in zip(action_tokens, action_token_shapes)], dim=0,
+            )  # [total_action_tokens, action_dim]
+            per_token_domain = torch.cat(
+                [d.reshape(1).expand(shp[0]) for d, shp in zip(action_domain_id, action_token_shapes)], dim=0,
+            )  # [total_action_tokens]
+            packed_action = self.action_proj_in(packed_action.to(target_dtype), per_token_domain)
+            packed_action = packed_action + self.action_modality_embed.view(1, -1)
+            if action_mse_loss_indexes is not None and action_mse_loss_indexes.numel() > 0:
+                a_ts = action_timesteps * self.timestep_scale
+                a_embeds = self.time_embedder(a_ts).to(target_dtype)
+                packed_action = self._scatter_timestep_embeds(packed_action, a_embeds, action_noisy_frame_indexes,
+                                                              action_token_shapes)
+            packed_sequence[action_sequence_indexes] = packed_action
+
         # 3. RoPE for the full packed sequence (cos/sin per token); split und/gen
         cos, sin = self.rotary_emb(position_ids, device=device, dtype=target_dtype)  # [N, head_dim]
         und_idx, gen_idx = self._mode_indices(split_lens, attn_modes, device)
@@ -895,9 +922,38 @@ class Cosmos3VFMTransformer(BaseDiT):
             preds_sound, _ = self.audio_proj_out(last_hidden_state[sound_mse_loss_indexes])
             output["preds_sound"] = self._unpack_sound(preds_sound, sound_token_shapes, sound_noisy_frame_indexes)
 
+        # 6c. Action prediction: action_proj_out(per-token domain) on noisy hidden
+        #     states -> unpack to per-sample [T, D] (framework ``_decode_action``).
+        if action_tokens and action_mse_loss_indexes is not None and action_mse_loss_indexes.numel() > 0:
+            noisy_domain = torch.cat(
+                [d.reshape(1).expand(len(nfi)) for d, nfi in zip(action_domain_id, action_noisy_frame_indexes)], dim=0,
+            )
+            preds_action = self.action_proj_out(last_hidden_state[action_mse_loss_indexes], noisy_domain)
+            output["preds_action"] = self._unpack_action(preds_action, action_token_shapes, action_noisy_frame_indexes)
+
         output["last_hidden_state"] = last_hidden_state
         return output
 
+    def _unpack_action(
+        self,
+        packed_preds: torch.Tensor,
+        token_shapes: list[tuple[int, ...]],
+        noisy_frame_indexes: list[torch.Tensor],
+    ) -> list[torch.Tensor]:
+        """Scatter packed noisy action preds back to per-sample ``[T, D]`` (clean
+        frames left zero). Mirrors framework ``unpack_action``."""
+        outputs: list[torch.Tensor] = []
+        start_idx = 0
+        for shape, nfi in zip(token_shapes, noisy_frame_indexes):
+            t = shape[0]
+            out = torch.zeros((t, self.action_dim), device=packed_preds.device, dtype=packed_preds.dtype)
+            t_n = len(nfi)
+            if t_n > 0:
+                out[nfi] = packed_preds[start_idx:start_idx + t_n]
+                start_idx += t_n
+            outputs.append(out)
+        return outputs
+
     def _unpack_sound(
         self,
         packed_preds: torch.Tensor,

fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py

@@ -48,6 +48,7 @@ from fastvideo.logger import init_logger
 from fastvideo.models.schedulers.scheduling_unipc_multistep import (
     UniPCMultistepScheduler, )
 from fastvideo.pipelines.basic.cosmos3.sequence_packing import (
+    Cosmos3ActionItem,
     Cosmos3SampleInputs,
     Cosmos3SoundItem,
     Cosmos3VisionItem,
@@ -215,6 +216,37 @@ class Cosmos3SoundSpec:
         return int(math.prod(self.shape))
 
 
+@dataclass
+class Cosmos3ActionSpec:
+    """Geometry + conditioning for one action item in a denoise run.
+
+    Args:
+        shape: ``(T, action_dim)`` of the action latent.
+        condition_frame_indexes: Frame indices kept clean (conditioning actions).
+        domain_id: Embodiment domain id for the domain-aware action projection.
+        fps: Action FPS; used iff fps modulation is on.
+    """
+
+    shape: tuple[int, int]
+    condition_frame_indexes: list[int] = field(default_factory=list)
+    domain_id: int = 0
+    fps: float | None = None
+
+    @property
+    def numel(self) -> int:
+        return int(math.prod(self.shape))
+
+
+def _split_action_latent(flat: torch.Tensor, specs: list[Cosmos3ActionSpec]) -> list[torch.Tensor]:
+    """Split a flat action vector into per-item ``[T, action_dim]`` tensors."""
+    out: list[torch.Tensor] = []
+    offset = 0
+    for spec in specs:
+        out.append(flat[offset:offset + spec.numel].reshape(spec.shape))
+        offset += spec.numel
+    return out
+
+
 def cosmos3_get_cfg_velocity(
     *,
     transformer: Any,
@@ -236,6 +268,8 @@ def cosmos3_get_cfg_velocity(
     normalize_cfg: bool = False,
     sound_specs: list[Cosmos3SoundSpec] | None = None,
     sound_fps_per_item: list[float] | None = None,
+    action_specs: list[Cosmos3ActionSpec] | None = None,
+    action_fps_per_item: list[float] | None = None,
 ) -> torch.Tensor:
     """Sequential-CFG velocity for one denoise step (framework math).
 
@@ -257,11 +291,16 @@ def cosmos3_get_cfg_velocity(
     assert timestep.numel() == 1, "timestep must be a scalar"
     timestep_value = float(timestep.reshape(()).item())
 
-    # Combined flat layout (t2vs): [all vision | all sound], matching the
-    # framework per-sample concat order ([vision_i | sound_i]); single sample here.
+    # Combined flat layout: [all vision | all action | all sound], matching the
+    # framework per-sample concat order ([vision_i | action_i | sound_i]); single
+    # sample here.
     vision_total = sum(spec.numel for spec in specs)
+    action_total = sum(spec.numel for spec in action_specs) if action_specs else 0
     noise_x_vision = _split_flat_latent(flat_latent[:vision_total], specs)
-    noise_x_sound = (_split_flat_latent(flat_latent[vision_total:], sound_specs) if sound_specs else None)
+    noise_x_action = (_split_action_latent(flat_latent[vision_total:vision_total +
+                                                       action_total], action_specs) if action_specs else None)
+    noise_x_sound = (_split_flat_latent(flat_latent[vision_total +
+                                                    action_total:], sound_specs) if sound_specs else None)
     device = next(transformer.parameters()).device
 
     def _run(token_ids: list[int]) -> torch.Tensor:
@@ -274,6 +313,16 @@ def cosmos3_get_cfg_velocity(
                     fps=(sound_fps_per_item[i] if sound_fps_per_item is not None else None),
                 ) for i, ss in enumerate(sound_specs)
             ]
+        action_items: list[Cosmos3ActionItem] = []
+        if action_specs is not None and noise_x_action is not None:
+            action_items = [
+                Cosmos3ActionItem(
+                    latent=noise_x_action[i],
+                    condition_frame_indexes=list(asp.condition_frame_indexes),
+                    domain_id=asp.domain_id,
+                    fps=(action_fps_per_item[i] if action_fps_per_item is not None else None),
+                ) for i, asp in enumerate(action_specs)
+            ]
         samples = [
             Cosmos3SampleInputs(
                 text_ids=list(token_ids),
@@ -283,6 +332,7 @@ def cosmos3_get_cfg_velocity(
                     fps=(fps_per_item[i] if fps_per_item is not None else None),
                 ),
                 sound=(sound_items[i] if i < len(sound_items) else None),
+                action=(action_items[i] if i < len(action_items) else None),
                 timestep=timestep_value,
             ) for i, (latent, spec) in enumerate(zip(noise_x_vision, specs, strict=False))
         ]
@@ -310,23 +360,38 @@ def cosmos3_get_cfg_velocity(
                 items.append(pred * keep if keep.sum() > 0 else torch.zeros_like(pred))
             vision_vel = torch.cat([v.reshape(-1) for v in items]).to(flat_latent.dtype)
 
-        if not sound_specs:
-            return vision_vel
+        parts = [vision_vel]
 
-        # Sound velocity: preds_sound are per-item [C, T], already zero on clean
-        # frames (unpack fills only noisy frames); zero on cond frames defensively.
-        sound_total = sum(spec.numel for spec in sound_specs)
-        sound_vel = torch.zeros(sound_total, device=flat_latent.device, dtype=flat_latent.dtype)
-        preds_s = out.get("preds_sound")
-        if preds_s is not None:
-            s_items: list[torch.Tensor] = []
-            for pred, cond_mask in zip(preds_s, packed.sound_condition_mask, strict=False):
-                pred = pred.squeeze(0) if pred.dim() == 3 else pred  # [C, T]
-                keep = (1.0 - cond_mask).reshape(1, -1).to(dtype=pred.dtype, device=pred.device)  # [1, T]
-                s_items.append(pred * keep)
-            sound_vel = torch.cat([v.reshape(-1) for v in s_items]).to(flat_latent.dtype)
+        if action_specs:
+            # Action velocity: preds_action are per-item [T, D], already zero on
+            # clean frames; zero on cond frames defensively.
+            action_vel = torch.zeros(action_total, device=flat_latent.device, dtype=flat_latent.dtype)
+            preds_a = out.get("preds_action")
+            if preds_a is not None:
+                a_items: list[torch.Tensor] = []
+                for pred, cond_mask in zip(preds_a, packed.action_condition_mask, strict=False):
+                    pred = pred.squeeze(0) if pred.dim() == 3 else pred  # [T, D]
+                    keep = (1.0 - cond_mask).reshape(-1, 1).to(dtype=pred.dtype, device=pred.device)  # [T, 1]
+                    a_items.append(pred * keep)
+                action_vel = torch.cat([v.reshape(-1) for v in a_items]).to(flat_latent.dtype)
+            parts.append(action_vel)
 
-        return torch.cat([vision_vel, sound_vel])
+        if sound_specs:
+            # Sound velocity: preds_sound are per-item [C, T], already zero on clean
+            # frames (unpack fills only noisy frames); zero on cond frames defensively.
+            sound_total = sum(spec.numel for spec in sound_specs)
+            sound_vel = torch.zeros(sound_total, device=flat_latent.device, dtype=flat_latent.dtype)
+            preds_s = out.get("preds_sound")
+            if preds_s is not None:
+                s_items: list[torch.Tensor] = []
+                for pred, cond_mask in zip(preds_s, packed.sound_condition_mask, strict=False):
+                    pred = pred.squeeze(0) if pred.dim() == 3 else pred  # [C, T]
+                    keep = (1.0 - cond_mask).reshape(1, -1).to(dtype=pred.dtype, device=pred.device)  # [1, T]
+                    s_items.append(pred * keep)
+                sound_vel = torch.cat([v.reshape(-1) for v in s_items]).to(flat_latent.dtype)
+            parts.append(sound_vel)
+
+        return vision_vel if len(parts) == 1 else torch.cat(parts)
 
     cond_v = _run(cond_token_ids)
     uncond_v = _run(uncond_token_ids)
@@ -382,6 +447,8 @@ class Cosmos3DenoiseEngine:
         fps_per_item: list[float] | None = None,
         sound_specs: list[Cosmos3SoundSpec] | None = None,
         sound_fps_per_item: list[float] | None = None,
+        action_specs: list[Cosmos3ActionSpec] | None = None,
+        action_fps_per_item: list[float] | None = None,
     ) -> torch.Tensor:
         return cosmos3_get_cfg_velocity(
             transformer=self.transformer,
@@ -402,6 +469,8 @@ class Cosmos3DenoiseEngine:
             fps_per_item=fps_per_item,
             sound_specs=sound_specs,
             sound_fps_per_item=sound_fps_per_item,
+            action_specs=action_specs,
+            action_fps_per_item=action_fps_per_item,
         )
 
     def denoise(
@@ -417,6 +486,8 @@ class Cosmos3DenoiseEngine:
         progress_bar: Any | None = None,
         sound_specs: list[Cosmos3SoundSpec] | None = None,
         sound_fps_per_item: list[float] | None = None,
+        action_specs: list[Cosmos3ActionSpec] | None = None,
+        action_fps_per_item: list[float] | None = None,
     ) -> torch.Tensor:
         """Run the full UniPC denoise loop, returning the final flat latent.
 
@@ -439,6 +510,8 @@ class Cosmos3DenoiseEngine:
                 fps_per_item=fps_per_item,
                 sound_specs=sound_specs,
                 sound_fps_per_item=sound_fps_per_item,
+                action_specs=action_specs,
+                action_fps_per_item=action_fps_per_item,
             )
             stepped = self.scheduler.step(
                 model_output=v_pred,

fastvideo/pipelines/basic/cosmos3/sequence_packing.py

@@ -86,14 +86,33 @@ class Cosmos3SoundItem:
     fps: float | None = None
 
 
+@dataclass
+class Cosmos3ActionItem:
+    """One action latent for a sample (action-conditioned world model).
+
+    Args:
+        latent: Action latent ``[T, action_dim]`` (per-frame action vectors).
+        condition_frame_indexes: Frame indices kept clean (conditioning actions).
+        domain_id: Embodiment domain id (scalar / ``[1]``) for the
+            domain-aware action projection.
+        fps: Action FPS; only used when ``enable_fps_modulation`` is set.
+    """
+
+    latent: torch.Tensor
+    condition_frame_indexes: list[int] = field(default_factory=list)
+    domain_id: int = 0
+    fps: float | None = None
+
+
 @dataclass
 class Cosmos3SampleInputs:
-    """Per-sample packing inputs (one text prompt + its vision item, +sound)."""
+    """Per-sample packing inputs (text prompt + vision item, +sound, +action)."""
 
     text_ids: list[int]
     vision: Cosmos3VisionItem
     timestep: float
     sound: Cosmos3SoundItem | None = None
+    action: Cosmos3ActionItem | None = None
 
 
 # ---------------------------------------------------------------------------
@@ -139,6 +158,16 @@ class Cosmos3PackedSequence:
     sound_condition_mask: list[torch.Tensor] = field(default_factory=list)
     fps_sound: torch.Tensor | None = None
 
+    # Action modality (action-conditioned world model); empty/None when no action.
+    action_tokens: list[torch.Tensor] = field(default_factory=list)
+    action_token_shapes: list[tuple[int, ...]] = field(default_factory=list)
+    action_sequence_indexes: torch.Tensor | None = None
+    action_timesteps: torch.Tensor | None = None
+    action_mse_loss_indexes: torch.Tensor | None = None
+    action_noisy_frame_indexes: list[torch.Tensor] = field(default_factory=list)
+    action_condition_mask: list[torch.Tensor] = field(default_factory=list)
+    action_domain_id: list[torch.Tensor] = field(default_factory=list)
+
     def to_dit_kwargs(self, device: torch.device | str | None = None) -> dict[str, Any]:
         """Return the kwargs dict for ``Cosmos3VFMTransformer.forward``.
 
@@ -171,6 +200,13 @@ class Cosmos3PackedSequence:
             sound_mse_loss_indexes=_mv(self.sound_mse_loss_indexes),
             sound_noisy_frame_indexes=[_mv(t) for t in self.sound_noisy_frame_indexes],
             fps_sound=_mv(self.fps_sound),
+            action_tokens=[_mv(t) for t in self.action_tokens],
+            action_token_shapes=list(self.action_token_shapes),
+            action_sequence_indexes=_mv(self.action_sequence_indexes),
+            action_timesteps=_mv(self.action_timesteps),
+            action_mse_loss_indexes=_mv(self.action_mse_loss_indexes),
+            action_noisy_frame_indexes=[_mv(t) for t in self.action_noisy_frame_indexes],
+            action_domain_id=[_mv(t) for t in self.action_domain_id],
         )
 
 
@@ -247,6 +283,15 @@ def pack_cosmos3_video_sequence(
     sound_condition_mask: list[torch.Tensor] = []
     sound_fps_values: list[float] = []
 
+    action_tokens: list[torch.Tensor] = []
+    action_token_shapes: list[tuple[int, ...]] = []
+    action_sequence_indexes: list[int] = []
+    action_timesteps: list[float] = []
+    action_mse_loss_indexes: list[int] = []
+    action_noisy_frame_indexes: list[torch.Tensor] = []
+    action_condition_mask: list[torch.Tensor] = []
+    action_domain_id: list[torch.Tensor] = []
+
     curr = 0  # running position in the packed sequence
     is_image_batch = True
 
@@ -335,6 +380,56 @@ def pack_cosmos3_video_sequence(
         curr += num_vision_tokens
         sample_len += num_vision_tokens
 
+        # ---- 2a2. Action split: shares the vision "full" split ----
+        # Mirrors framework ``_pack_action_tokens``: action latent [T, D] -> T
+        # tokens (token shape (T,)), domain-aware, 3D-MRoPE at the vision temporal
+        # offset with ``start_frame_offset=1`` (parallel to vision; tcf=1; does
+        # not advance the offset).
+        action_split_len = 0
+        if sample.action is not None:
+            action_latent = sample.action.latent  # [T, D]
+            action_t = int(action_latent.shape[0])
+            action_split_len = action_t
+
+            action_tokens.append(action_latent)
+            action_token_shapes.append((action_t, ))
+            action_sequence_indexes.extend(range(curr, curr + action_t))
+            action_domain_id.append(torch.tensor([int(sample.action.domain_id)], dtype=torch.long))
+
+            a_cond_set = {idx for idx in sample.action.condition_frame_indexes if 0 <= idx < action_t}
+            a_cond_mask = torch.zeros((action_t, 1), device=action_latent.device, dtype=action_latent.dtype)
+            for fi in a_cond_set:
+                a_cond_mask[fi, 0] = 1.0
+            action_condition_mask.append(a_cond_mask)
+
+            a_noisy = torch.tensor([idx for idx in range(action_t) if idx not in a_cond_set],
+                                   device=action_latent.device,
+                                   dtype=torch.long)
+            action_noisy_frame_indexes.append(a_noisy)
+
+            for fi in range(action_t):
+                if fi in a_cond_set:
+                    continue
+                action_mse_loss_indexes.append(curr + fi)
+                action_timesteps.append(float(sample.timestep))
+
+            action_fps = sample.action.fps if enable_fps_modulation else None
+            action_mrope, _ = compute_mrope_position_ids_vision(
+                grid_t=action_t,
+                grid_h=1,
+                grid_w=1,
+                temporal_offset=vision_start_temporal_offset,
+                fps=action_fps,
+                base_fps=base_fps,
+                temporal_compression_factor=1,  # action is at frame rate
+                base_temporal_compression_factor=temporal_compression_factor,
+                enable_fps_modulation=enable_fps_modulation,
+                start_frame_offset=1,
+            )
+            position_id_blocks.append(action_mrope)
+            curr += action_t
+            sample_len += action_t
+
         # ---- 2b. Sound split (t2vs): shares the vision "full" split ----
         # Mirrors framework ``_pack_sound_tokens``: sound latent [C, T] -> T
         # tokens (token shape (T,1,1)), packed right after vision, with 3D-MRoPE
@@ -401,9 +496,9 @@ def pack_cosmos3_video_sequence(
             eov_len = 1
             sample_len += 1
 
-        # Vision + sound + any trailing eov marker share one "full" split.
+        # Vision + action + sound + any trailing eov marker share one "full" split.
         attn_modes.append("full")
-        split_lens.append(num_vision_tokens + sound_split_len + eov_len)
+        split_lens.append(num_vision_tokens + action_split_len + sound_split_len + eov_len)
         sample_lens.append(sample_len)
 
         if latent_t != 1:
@@ -443,4 +538,12 @@ def pack_cosmos3_video_sequence(
         sound_noisy_frame_indexes=sound_noisy_frame_indexes,
         sound_condition_mask=sound_condition_mask,
         fps_sound=(torch.tensor(sound_fps_values, dtype=torch.float32) if sound_fps_values else None),
+        action_tokens=action_tokens,
+        action_token_shapes=action_token_shapes,
+        action_sequence_indexes=(torch.tensor(action_sequence_indexes, dtype=torch.long) if action_tokens else None),
+        action_timesteps=(torch.tensor(action_timesteps, dtype=timesteps_dtype) if action_tokens else None),
+        action_mse_loss_indexes=(torch.tensor(action_mse_loss_indexes, dtype=torch.long) if action_tokens else None),
+        action_noisy_frame_indexes=action_noisy_frame_indexes,
+        action_condition_mask=action_condition_mask,
+        action_domain_id=action_domain_id,
     )

tests/local_tests/cosmos3/test_cosmos3_action_parity.py

@@ -0,0 +1,271 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Numerical-parity test: FastVideo Cosmos3 action pathway vs the framework.
+
+Covers the action (multi-embodiment world-model) modality at the DiT level:
+
+  * **action packing** — native ``pack_cosmos3_video_sequence`` with a
+    ``Cosmos3ActionItem`` vs framework ``pack_input_sequence`` with
+    ``has_action``: action tokens share the vision "full" split, with ``(T,)``
+    shapes, a ``(T,1)`` condition mask, and 3D-MRoPE temporal positions at the
+    vision offset with ``start_frame_offset=1`` (parallel to vision); and
+  * **DiT action forward** — the dormant domain-aware ``action_proj_in`` /
+    ``action_proj_out`` (``DomainAwareLinear``) + ``action_modality_embed`` heads,
+    now activated, with a per-token embodiment ``domain_id``.
+
+Framework model + pack is the parity ORACLE (CPU/float32 via SDPA monkey-patch).
+We assert the native packer matches the framework field-by-field, then that
+``preds_vision`` AND ``preds_action`` match the framework forward.
+
+Run:
+    cd <worktree> && <fv-cosmos3 python> -m pytest \
+        tests/local_tests/cosmos3/test_cosmos3_action_parity.py -q -s
+"""
+from __future__ import annotations
+
+import pytest
+import torch
+
+cosmos_framework = pytest.importorskip(
+    "cosmos_framework",
+    reason="cosmos_framework not installed; run in fv-cosmos3 env.",
+)
+
+from .test_cosmos3_dit_parity import (  # noqa: E402
+    _fastvideo_inputs_from_packed_seq,
+    _framework_to_fastvideo_state_dict,
+)
+from .test_cosmos3_dit_parity_mrope import (  # noqa: E402
+    _ACTION_DIM,
+    _LATENT_CHANNEL,
+    _LATENT_PATCH_SIZE,
+    _RESET_SPATIAL_IDS,
+    _TCF,
+    _TEMPORAL_MODALITY_MARGIN,
+    _build_tiny_cosmos3_mrope,
+    _build_tiny_fastvideo_dit_mrope,
+)
+from .test_cosmos3_reference_forward import _apply_sdpa_patches  # noqa: E402
+
+pytestmark = [pytest.mark.local]
+_apply_sdpa_patches()
+
+_SPECIAL_TOKENS = {"start_of_generation": 60, "end_of_generation": 61, "eos_token_id": 62}
+
+
+def _copy_weights_with_action(vfm, dit) -> None:
+    """Copy backbone + vision weights AND the domain-aware action heads."""
+    mapped = _framework_to_fastvideo_state_dict(vfm, num_layers=dit.num_hidden_layers)
+    src = dict(vfm.named_parameters())
+    mapped["action_proj_in.fc.weight"] = src["action2llm.fc.weight"].detach().clone()
+    mapped["action_proj_in.bias.weight"] = src["action2llm.bias.weight"].detach().clone()
+    mapped["action_proj_out.fc.weight"] = src["llm2action.fc.weight"].detach().clone()
+    mapped["action_proj_out.bias.weight"] = src["llm2action.bias.weight"].detach().clone()
+    mapped["action_modality_embed"] = src["action_modality_embed"].detach().clone()
+    dst = dict(dit.named_parameters())
+    with torch.no_grad():
+        for name, tensor in mapped.items():
+            assert name in dst, f"DiT missing param {name!r}"
+            assert dst[name].shape == tensor.shape, f"shape mismatch {name}"
+            dst[name].copy_(tensor.to(dst[name].dtype))
+
+
+def _framework_pack_action(*, text_ids, vision, action, cond_vision, cond_action, domain_id, timestep,
+                           is_image_batch):
+    from cosmos_framework.data.vfm.sequence_packing import (
+        GenerationDataClean,
+        SequencePlan,
+        pack_input_sequence,
+    )
+
+    gen = GenerationDataClean(
+        batch_size=1,
+        is_image_batch=is_image_batch,
+        x0_tokens_vision=[vision],
+        fps_vision=None,
+        num_vision_items_per_sample=[1],
+        x0_tokens_action=[action],
+        fps_action=None,
+        action_domain_id=[torch.tensor([domain_id], dtype=torch.long)],
+    )
+    plans = [SequencePlan(
+        has_text=True, has_vision=True, has_action=True,
+        condition_frame_indexes_vision=list(cond_vision),
+        condition_frame_indexes_action=list(cond_action),
+    )]
+    ps = pack_input_sequence(
+        sequence_plans=plans,
+        input_text_indexes=[list(text_ids)],
+        gen_data_clean=gen,
+        input_timesteps=torch.tensor([timestep], dtype=torch.float32),
+        special_tokens=_SPECIAL_TOKENS,
+        latent_patch_size=_LATENT_PATCH_SIZE,
+        include_end_of_generation_token=False,
+        position_embedding_type="unified_3d_mrope",
+        unified_3d_mrope_reset_spatial_ids=_RESET_SPATIAL_IDS,
+        unified_3d_mrope_temporal_modality_margin=_TEMPORAL_MODALITY_MARGIN,
+        enable_fps_modulation=False,
+        base_fps=24.0,
+        temporal_compression_factor=_TCF,
+    )
+    # The framework sets action.domain_id on the packed sequence from
+    # gen_data_clean inside the model (_get_velocity); mirror that for the oracle.
+    if ps.action is not None:
+        ps.action.domain_id = [torch.tensor([domain_id], dtype=torch.long)]
+    return ps
+
+
+def _fastvideo_pack_action(*, text_ids, vision, action, cond_vision, cond_action, domain_id, timestep):
+    from fastvideo.pipelines.basic.cosmos3.sequence_packing import (
+        Cosmos3ActionItem,
+        Cosmos3SampleInputs,
+        Cosmos3VisionItem,
+        pack_cosmos3_video_sequence,
+    )
+
+    samples = [Cosmos3SampleInputs(
+        text_ids=list(text_ids),
+        vision=Cosmos3VisionItem(latent=vision, condition_frame_indexes=list(cond_vision)),
+        action=Cosmos3ActionItem(latent=action, condition_frame_indexes=list(cond_action), domain_id=domain_id),
+        timestep=float(timestep),
+    )]
+    return pack_cosmos3_video_sequence(
+        samples, _SPECIAL_TOKENS,
+        latent_patch_size=_LATENT_PATCH_SIZE, include_end_of_generation_token=False,
+        temporal_modality_margin=_TEMPORAL_MODALITY_MARGIN, reset_spatial_ids=_RESET_SPATIAL_IDS,
+        enable_fps_modulation=False, base_fps=24.0, temporal_compression_factor=_TCF,
+    )
+
+
+def _fv_inputs_with_action(ps) -> dict:
+    kw = _fastvideo_inputs_from_packed_seq(ps)
+    a = ps.action
+    kw.update(
+        action_tokens=list(a.tokens),
+        action_token_shapes=[tuple(x) for x in a.token_shapes],
+        action_sequence_indexes=a.sequence_indexes,
+        action_timesteps=a.timesteps,
+        action_mse_loss_indexes=a.mse_loss_indexes,
+        action_noisy_frame_indexes=list(a.noisy_frame_indexes),
+        action_domain_id=list(a.domain_id),
+    )
+    return kw
+
+
+def _diffs(a, b):
+    d = (a - b).abs()
+    return d.max().item(), d.mean().item()
+
+
+# (grid_t, lh, lw, action_t, n_text, cond_vision, cond_action, domain_id)
+_CASES = [
+    pytest.param(2, 4, 4, 6, 4, [], [], 0, id="a2v_2x2x2_act6_dom0"),
+    pytest.param(3, 8, 4, 9, 5, [], [], 7, id="a2v_3x4x2_act9_dom7"),
+    pytest.param(2, 4, 4, 5, 5, [0], [0], 3, id="ai2v_cond_act5_dom3"),
+]
+
+
+class TestCosmos3ActionParity:
+
+    def _build(self, num_layers=2, seed_model=42):
+        vfm = _build_tiny_cosmos3_mrope(seed=seed_model, num_layers=num_layers, action_gen=True)
+        dit = _build_tiny_fastvideo_dit_mrope(num_layers=num_layers)
+        _copy_weights_with_action(vfm, dit)
+        return vfm, dit
+
+    def _make_inputs(self, grid_t, lh, lw, act_t, n_text, cond_v, cond_a, dom, seed=7):
+        torch.manual_seed(seed)
+        return dict(
+            text_ids=torch.randint(0, 60, (n_text,)).tolist(),
+            vision=torch.randn(1, _LATENT_CHANNEL, grid_t, lh, lw),
+            action=torch.randn(act_t, _ACTION_DIM),  # [T, D]
+            cond_vision=cond_v, cond_action=cond_a, domain_id=dom, timestep=500.0,
+        )
+
+    @pytest.mark.parametrize(("grid_t", "lh", "lw", "act_t", "n_text", "cond_v", "cond_a", "dom"), _CASES)
+    def test_action_packing_matches_framework(self, grid_t, lh, lw, act_t, n_text, cond_v, cond_a, dom):
+        ins = self._make_inputs(grid_t, lh, lw, act_t, n_text, cond_v, cond_a, dom)
+        fw = _framework_pack_action(is_image_batch=(grid_t == 1), **ins)
+        fv = _fastvideo_pack_action(**ins)
+        assert fv.split_lens == list(fw.split_lens), f"split_lens fv={fv.split_lens} fw={list(fw.split_lens)}"
+        assert fv.attn_modes == list(fw.attn_modes)
+        assert int(fv.sequence_length) == int(fw.sequence_length)
+        torch.testing.assert_close(fv.position_ids, fw.position_ids, rtol=0, atol=0)
+        a = fw.action
+        torch.testing.assert_close(fv.action_sequence_indexes, a.sequence_indexes.to(torch.long), rtol=0, atol=0)
+        assert fv.action_token_shapes == [tuple(x) for x in a.token_shapes]
+        torch.testing.assert_close(fv.action_timesteps.to(torch.float32), a.timesteps.to(torch.float32))
+        torch.testing.assert_close(fv.action_mse_loss_indexes, a.mse_loss_indexes.to(torch.long), rtol=0, atol=0)
+        for x, y in zip(fv.action_noisy_frame_indexes, a.noisy_frame_indexes):
+            torch.testing.assert_close(x.to(torch.long), y.to(torch.long), rtol=0, atol=0)
+        print(f"\n[action_packing {grid_t}x{lh}x{lw} act={act_t} dom={dom}] position_ids + action fields exact")
+
+    @pytest.mark.parametrize(("grid_t", "lh", "lw", "act_t", "n_text", "cond_v", "cond_a", "dom"), _CASES)
+    def test_action_dit_forward_matches_framework(self, grid_t, lh, lw, act_t, n_text, cond_v, cond_a, dom):
+        vfm, dit = self._build()
+        ins = self._make_inputs(grid_t, lh, lw, act_t, n_text, cond_v, cond_a, dom)
+        fw_pack = _framework_pack_action(is_image_batch=(grid_t == 1), **ins)
+        fv_pack = _fastvideo_pack_action(**ins)
+        with torch.no_grad():
+            fw_out = vfm(packed_seq=fw_pack)
+            fv_out = dit(**fv_pack.to_dit_kwargs())
+            fv_on_fw = dit(**_fv_inputs_with_action(fw_pack))
+        pv_mx, pv_mn = _diffs(fv_out["preds_vision"][0], fw_out["preds_vision"][0])
+        pa_mx, pa_mn = _diffs(fv_out["preds_action"][0], fw_out["preds_action"][0])
+        paf_mx, paf_mn = _diffs(fv_on_fw["preds_action"][0], fw_out["preds_action"][0])
+        print(f"\n[action_dit {grid_t}x{lh}x{lw} act={act_t} dom={dom}] "
+              f"preds_vision max={pv_mx:.3e} mean={pv_mn:.3e} | "
+              f"preds_action max={pa_mx:.3e} mean={pa_mn:.3e} | "
+              f"preds_action(fwpack) max={paf_mx:.3e} mean={paf_mn:.3e}")
+        assert fv_out["preds_action"][0].shape == fw_out["preds_action"][0].shape
+        torch.testing.assert_close(fv_out["preds_vision"][0], fw_out["preds_vision"][0], atol=1e-4, rtol=1e-3)
+        torch.testing.assert_close(fv_out["preds_action"][0], fw_out["preds_action"][0], atol=1e-4, rtol=1e-3)
+        torch.testing.assert_close(fv_on_fw["preds_action"][0], fw_out["preds_action"][0], atol=1e-4, rtol=1e-3)
+
+    @pytest.mark.parametrize(("grid_t", "lh", "lw", "act_t", "n_text", "cond_v", "cond_a", "dom"), _CASES)
+    def test_action_cfg_velocity_matches_framework(self, grid_t, lh, lw, act_t, n_text, cond_v, cond_a, dom):
+        """Combined [vision|action] sequential-CFG velocity (action pipeline glue)
+        matches a framework-DiT oracle."""
+        from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import (
+            Cosmos3ActionSpec,
+            Cosmos3VisionSpec,
+            cosmos3_get_cfg_velocity,
+        )
+
+        vfm, dit = self._build()
+        vlat_shape = (_LATENT_CHANNEL, grid_t, lh, lw)
+        action_shape = (act_t, _ACTION_DIM)
+        torch.manual_seed(3)
+        cond_ids = torch.randint(0, 60, (n_text,)).tolist()
+        uncond_ids = torch.randint(0, 60, (max(1, n_text - 1),)).tolist()
+        vis_numel = int(torch.tensor(vlat_shape).prod())
+        act_numel = int(torch.tensor(action_shape).prod())
+        flat = torch.randn(vis_numel + act_numel)
+        guidance, ts = 6.0, 500.0
+
+        def _fw_velocity(ids):
+            vision = flat[:vis_numel].reshape(vlat_shape).unsqueeze(0)
+            action = flat[vis_numel:].reshape(action_shape)
+            ps = _framework_pack_action(text_ids=ids, vision=vision, action=action, cond_vision=cond_v,
+                                        cond_action=cond_a, domain_id=dom, timestep=ts,
+                                        is_image_batch=(grid_t == 1))
+            with torch.no_grad():
+                out = vfm(packed_seq=ps)
+            pv = out["preds_vision"][0].squeeze(0)  # [C,T,H,W] (zero on clean)
+            pa = out["preds_action"][0]  # [T,D] (zero on clean)
+            return torch.cat([pv.reshape(-1), pa.reshape(-1)])
+
+        fw_cond, fw_uncond = _fw_velocity(cond_ids), _fw_velocity(uncond_ids)
+        fw_v = fw_uncond + guidance * (fw_cond - fw_uncond)
+        fv_v = cosmos3_get_cfg_velocity(
+            transformer=dit, flat_latent=flat, timestep=torch.tensor([ts]), guidance=guidance,
+            specs=[Cosmos3VisionSpec(shape=vlat_shape, condition_frame_indexes=list(cond_v))],
+            action_specs=[Cosmos3ActionSpec(shape=action_shape, condition_frame_indexes=list(cond_a), domain_id=dom)],
+            cond_token_ids=cond_ids, uncond_token_ids=uncond_ids,
+            special_tokens=_SPECIAL_TOKENS, latent_patch_size=_LATENT_PATCH_SIZE,
+            temporal_modality_margin=_TEMPORAL_MODALITY_MARGIN, reset_spatial_ids=_RESET_SPATIAL_IDS,
+            enable_fps_modulation=False, base_fps=24.0, temporal_compression_factor=_TCF,
+        )
+        assert fv_v.shape == fw_v.shape, f"shape fv={fv_v.shape} fw={fw_v.shape}"
+        mx, mn = _diffs(fv_v, fw_v)
+        print(f"\n[action_cfg_velocity {grid_t}x{lh}x{lw} act={act_t} dom={dom}] max={mx:.3e} mean={mn:.3e}")
+        torch.testing.assert_close(fv_v, fw_v, atol=1e-4, rtol=1e-3)

tests/local_tests/cosmos3/test_cosmos3_dit_parity_mrope.py

@@ -73,9 +73,12 @@ _TCF = 4  # temporal compression factor
 # ---------------------------------------------------------------------------
 _SOUND_DIM = 64
 _SOUND_LATENT_FPS = 25
+_ACTION_DIM = 64
+_NUM_EMBODIMENT_DOMAINS = 32
 
 
-def _build_tiny_cosmos3_mrope(seed: int = 42, num_layers: int = 2, sound_gen: bool = False):
+def _build_tiny_cosmos3_mrope(seed: int = 42, num_layers: int = 2, sound_gen: bool = False,
+                              action_gen: bool = False):
     """Tiny framework ``Cosmos3VFMNetwork`` with ``unified_3d_mrope``.
 
     ``rope_theta`` / ``rope_scaling`` (carrying ``mrope_section`` +
@@ -129,6 +132,11 @@ def _build_tiny_cosmos3_mrope(seed: int = 42, num_layers: int = 2, sound_gen: bo
         temporal_compression_factor_sound=1,
         sound_latent_fps=_SOUND_LATENT_FPS,
     ) if sound_gen else {}
+    action_kwargs = dict(
+        action_gen=True,
+        action_dim=_ACTION_DIM,
+        num_embodiment_domains=_NUM_EMBODIMENT_DOMAINS,
+    ) if action_gen else {}
     vfm_cfg = Cosmos3VFMNetworkConfig(
         vision_gen=True,
         vlm_config=tiny_vlm_cfg,
@@ -141,6 +149,7 @@ def _build_tiny_cosmos3_mrope(seed: int = 42, num_layers: int = 2, sound_gen: bo
         max_latent_t=8,
         temporal_compression_factor_vision=_TCF,
         **sound_kwargs,
+        **action_kwargs,
     )
     torch.manual_seed(seed)
     lm = Qwen3VLTextForCausalLM(config=mot_cfg)