Add Bridge native Stage2 training source; GPU validation NOT_RUN

releases/bridge-native-v3/README.md

@@ -0,0 +1,160 @@
+# Bridge/WidowX handoff — memory assets verified
+
+**v3 implementation update:** `scripts/run_bridge_stage2_native.py` now exists.
+See [native training instructions](code/docs/bridge_native_training.md). Its GPU training and
+50/500-step acceptance remain NOT_RUN. The previously published v2 code snapshot
+predates this implementation. Use v3 for the current code; the GPU readiness flag stays false.
+
+The v3 bundle contains the native trainer source and the verified v2 memory assets.
+Use `releases/bridge-native-v3`. The earlier v1 upload was incomplete because
+the Space's existing `.gitignore` silently omitted `.pt` files. It is retained
+for provenance, not as a usable release. The publisher checks every remote
+Git/LFS hash and preserves all pre-existing Space files and ignore rules.
+Do not run `run_bridge_stage2_pilot.py`, `run_ipec_full_training.sh`, or
+`run_widowx_training.sh` for this experiment: those historical paths use the old
+policy contracts and/or old RoboCasa/IPEC initialization.
+
+## What is delivered
+
+- `memory/cte_30000_inference.pt`: the exact 177 model tensors of the selected
+  Bridge reached-state-delta CTE at step 30,000. No optimizer/RNG is exported.
+- `memory/static_bank_3000/`: unchanged bank and its original manifest, selection,
+  support/query partition, and the shards needed to verify the original manifest.
+  Keys are `[3000,128]`, values `[3000,256]`, one entry per support episode.
+- `memory/stage1_split.json`, `pilot_split.json`, and the separate CTE legacy and
+  policy native10 normalizers.
+- `code/`: the actual ZeVA overlay source snapshot, including local/uncommitted
+  scripts, the native geometry contract helpers, tests, and handoff tools.
+- `handoff_manifest.json`: per-file SHA256/size, exact upstream versions, data/VAE
+  provenance, CTE export-to-original tensor binding, and explicit readiness.
+
+The local original 30k CTE checkpoint is unchanged. Its SHA256 is
+`5c94378ea87b875dd709d46cf7941cdb5ca33584df071fef254311c6bbab912e`.
+It contained a Stage1 optimizer, so the downloadable inference export necessarily
+has a **different file hash**. All model tensors were checked exactly equal;
+the export manifest records both identities. The bank's original source-checkpoint
+identity is preserved, not silently rewritten to the export hash.
+
+The original bank SHA256 remains
+`bd248a1df2a908e3f16a6a6df66c50ad691514a011edb4c4704493c2e5799ffb`.
+Absolute paths in original provenance JSON are historical records only, not
+runtime paths. Use bundle-relative paths and local CLI arguments.
+
+## Download and verify on another machine
+
+First obtain the release commit from the handoff message or publication receipt.
+Replace `RELEASE_COMMIT` below with that full commit. Never substitute a moving
+`main` when reproducing an experiment.
+
+```bash
+hf download DanTim05/DynaTTT --repo-type space \
+  --revision RELEASE_COMMIT --include 'releases/bridge-native-v3/*' \
+  --local-dir ./zeva-transfer
+
+cd zeva-transfer/releases/bridge-native-v3
+python code/scripts/bridge_handoff.py verify --bundle .
+```
+
+That first check uses only the Python standard library. A successful memory check
+does **not** imply training readiness: `training_ready` remains false.
+
+Install the pinned NVIDIA Cosmos Framework following its setup guide. ZeVA is an
+overlay, not a stand-alone pip package. In a new dedicated workspace:
+
+```bash
+git clone https://github.com/NVIDIA/cosmos-framework.git cosmos-framework
+git -C cosmos-framework checkout --detach ee58e41467f33c49ddde08b4d0ef4923876a95ac
+
+# Run from the downloaded release directory; use an absolute path to the clone.
+export BRIDGE_FRAMEWORK_DIR=/absolute/path/to/cosmos-framework
+rsync -a --exclude=README.md ./code/ "$BRIDGE_FRAMEWORK_DIR/"
+cd "$BRIDGE_FRAMEWORK_DIR"
+
+# Requires the CUDA/PyTorch prerequisites documented by this pinned framework.
+uv sync --all-extras --group=cu128-train
+source .venv/bin/activate
+```
+
+Run in the resulting environment, specifying the release directory:
+
+```bash
+python -m unittest -v scripts.test_bridge_handoff scripts.test_bridge_native_contract
+python -m scripts.bridge_handoff verify \
+  --bundle /absolute/path/to/zeva-transfer/releases/bridge-native-v3 --deep
+```
+
+The export uses safe tensor/basic-type deserialization (`weights_only=True`).
+It loads strictly into CTE and is frozen. Nothing here loads a policy checkpoint,
+starts a training job, allocates GPUs, or stops another process.
+
+## Download pinned external dependencies
+
+```bash
+python -m scripts.bridge_handoff download \
+  --bundle /absolute/path/to/zeva-transfer/releases/bridge-native-v3 \
+  --destination /large/disk/bridge-run --base --dataset --vae
+```
+
+The downloader uses:
+
+- `nvidia/Cosmos3-Nano@e59a53c25979a090fa8706c9acc0c254a6e89b92`, including
+  root configuration, transformer, vision encoder and tokenizer. A URL ending
+  in `/transformer` alone is not a complete training installation.
+- `IPEC-COMMUNITY/bridge_orig_lerobot@0e9d76d07e9df3ea3eba257b2520d4913833fad2`.
+  Its metadata must match the CTE source hashes; no task strings are invented.
+- The exact pinned Wan VAE revision recorded in the manifest. Expected SHA256:
+  `20eb789667fa5e60e7516bf509512f6cb61f01b0aa0695eadaea930c13892b36`.
+
+Use each machine's own HF/W&B credentials; no credentials are included here.
+Before a full run, verify all expected per-episode parquet/image_0 video files,
+not merely the presence of `meta/info.json`.
+
+## Fixed research design / release validation
+
+`code/configs/bridge_stage2_native_contract.json` describes the research contract.
+The native launcher implements it; use the v3 source and
+[training instructions](code/docs/bridge_native_training.md), alongside the included memory assets.
+`verify --require-training-ready` still exits nonzero for this release
+because its GPU acceptance has not passed.
+
+The native path must preserve:
+
+- Native policy `[B,32,10]` (XYZ + rot6d + gripper index 9), current absolute
+  proprio `[B,10]`, independent reached CTE7 (RPY, gripper index 6), runtime-verified
+  internal width 64 with padded channels masked. Labels come from adjacent raw
+  states through NVIDIA's transforms, not reached7 converted to 10D.
+- Original Nano domain-7 action projections and shared action embedding. Compare
+  checkpoint/constructed/pre-optimizer hashes. No DROID specialist, RoboCasa,
+  old IPEC delta, old 7D policy, or old policy optimizer.
+- Frozen shared understanding/reasoner, Wan VAE and CTE; train generation
+  `moe_gen/time_embedder/vae2llm/llm2vae`, domain7 action rows, and new ZeVA/proprio
+  modules. Never put the entire policy forward in no_grad. Use a real WAM vision
+  loss so llm2vae receives nonzero gradients. Other-domain rows must not change.
+- A **new** trained/frozen Bridge query head (not supplied yet), pure frozen
+  understanding image+instruction readout 4096→1024→128, cosine top5/T=0.07.
+  Do not reuse a RoboCasa head or treat random queries as trained retrieval.
+- Support/query separation; phase/effects from query prefixes only; one
+  transition per four actual controls, first completed effect at 16 controls,
+  last four effects only. No all-episode future context and no full-dataset cache.
+- Native10 decoding must invert the tool/flange/OpenCV transforms and account
+  for the Simpler controller's previous target. Recompute reached7 from actual
+  before/after states only after env.step.
+- Real runtime shape trace; 32 train/8 development 50-step engineering gate,
+  then 500-step fit gate, then continuous. The 8 development episodes are held
+  out from Stage2, not independent of CTE training. Log sampled actions, masks,
+  per-dimension/physical/gripper errors, BIT histogram, gradients and deltas.
+- Generation/action/ZeVA LR 2e-5/1e-4/3e-5; domain tables weight_decay=0;
+  every 50 steps verify other-domain diff=0 and frozen hashes unchanged.
+- Save resumable trainable replacements, optimizer/scheduler/RNG/data cursor
+  every 10k optimizer updates; never delete old checkpoints. Estimate ~91 GiB
+  per checkpoint with FP32-master Adam; plan disk capacity accordingly.
+
+50/500-step gates, full model runtime dimensions, online environment acceptance,
+and continuous training are **NOT RUN**. CPU geometry/asset tests are not substitutes.
+
+## Licenses and provenance
+
+The source snapshot includes the original LICENSE/NOTICE/ATTRIBUTIONS. The
+dataset, NVIDIA base and Wan VAE retain their respective upstream licenses.
+The 3000-entry bank is derived from IPEC Bridge; it is not a copy of the raw videos.
+No base model, raw dataset, policy optimizer, secret, or training log is redistributed.

releases/bridge-native-v3/code/LICENSE

@@ -0,0 +1,52 @@
+Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+
+OpenMDW License Agreement, version 1.1 (OpenMDW-1.1)
+
+By exercising rights granted to you under this agreement, you accept and agree
+to its terms.
+
+As used in this agreement, "Model Materials" means the materials provided to
+you under this agreement, consisting of: (1) one or more machine learning
+models (including architecture and parameters); and (2) all related artifacts
+(including associated data, documentation and software) that are provided to
+you hereunder.
+
+Subject to your compliance with this agreement, permission is hereby granted,
+free of charge, to deal in the Model Materials without restriction, including
+under all copyright, patent, database, and trade secret rights included or
+embodied therein.
+
+If you distribute any portion of the Model Materials, you shall retain in your
+distribution (1) a copy of this agreement, and (2) all copyright notices and
+other notices of origin included in the Model Materials that are applicable to
+your distribution.
+
+If you file, maintain, or voluntarily participate in a lawsuit against any
+person or entity asserting that the Model Materials directly or indirectly
+infringe any patent or copyright, then all rights and grants made to you
+hereunder are terminated, unless that lawsuit was in response to a
+corresponding lawsuit first brought against you.
+
+This agreement does not impose any restrictions or obligations with respect to
+any use, modification, or sharing of any outputs generated by using the Model
+Materials.
+
+THE MODEL MATERIALS ARE PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
+OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE, TITLE, NONINFRINGEMENT, ACCURACY, OR THE
+ABSENCE OF LATENT OR OTHER DEFECTS OR ERRORS, WHETHER OR NOT DISCOVERABLE, ALL
+TO THE GREATEST EXTENT PERMISSIBLE UNDER APPLICABLE LAW.
+
+YOU ARE SOLELY RESPONSIBLE FOR (1) CLEARING RIGHTS OF OTHER PERSONS THAT MAY
+APPLY TO THE MODEL MATERIALS OR ANY USE THEREOF, INCLUDING WITHOUT LIMITATION
+ANY PERSON'S COPYRIGHTS OR OTHER RIGHTS INCLUDED OR EMBODIED IN THE MODEL
+MATERIALS; (2) OBTAINING ANY NECESSARY CONSENTS, PERMISSIONS OR OTHER RIGHTS
+REQUIRED FOR ANY USE OF THE MODEL MATERIALS; OR (3) PERFORMING ANY DUE
+DILIGENCE OR UNDERTAKING ANY OTHER INVESTIGATIONS INTO THE MODEL MATERIALS OR
+ANYTHING INCORPORATED OR EMBODIED THEREIN.
+
+IN NO EVENT SHALL THE PROVIDERS OF THE MODEL MATERIALS BE LIABLE FOR ANY CLAIM,
+DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
+OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE MODEL MATERIALS, THE
+USE THEREOF OR OTHER DEALINGS THEREIN.

releases/bridge-native-v3/code/NOTICE

@@ -0,0 +1,54 @@
+Cosmos
+
+Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
+
+This file contains attribution notices for third-party source code that is
+adapted, copied, or vendored in this repository. Dependency notices are listed
+separately in ATTRIBUTIONS.md.
+
+HuggingFace Transformers / Qwen
+Source: https://github.com/huggingface/transformers
+License: Apache-2.0
+Notices:
+  Copyright 2024 The Qwen team, Alibaba Group and the HuggingFace Inc. team.
+  Copyright 2025 The Qwen team, Alibaba Group and the HuggingFace Inc. team.
+  Copyright 2025 The Qwen Team and The HuggingFace Inc. team.
+  Copyright 2025 HuggingFace Inc. team.
+Files:
+  cosmos_framework/model/generator/vlm/qwen3_vl/configuration_qwen3_vl.py
+  cosmos_framework/model/generator/vlm/qwen3_vl/qwen3_vl.py
+  cosmos_framework/model/generator/vlm/qwen3_vl/utils.py
+  cosmos_framework/model/generator/vlm/qwen3_vl/video_processing_qwen3_vl.py
+
+ByteDance-Seed Bagel / Qwen2 tokenizer
+Source: https://github.com/ByteDance-Seed/Bagel
+License: Apache-2.0
+Notice:
+  Copyright 2024 The Qwen Team and The HuggingFace Inc. team.
+File:
+  cosmos_framework/model/generator/tokenizers/tokenization_qwen2.py
+
+HuggingFace Diffusers UniPC scheduler
+Source: https://github.com/huggingface/diffusers
+License: Apache-2.0
+Notices:
+  Copyright 2024 TSAIL Team and The HuggingFace Team. All rights reserved.
+  Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+File:
+  cosmos_framework/model/generator/diffusion/samplers/fm_solvers_unipc.py
+
+Detectron2 lazy config helpers
+Source: https://github.com/facebookresearch/detectron2
+License: Apache-2.0
+Notice:
+  Copyright (c) Facebook, Inc. and its affiliates.
+File:
+  cosmos_framework/utils/lazy_config/__init__.py
+
+TaylorSeer
+Source: https://github.com/Shenyi-Z/TaylorSeer
+License: GPL-3.0
+Notice:
+  Upstream source files do not contain an explicit copyright notice.
+Files:
+  cosmos_framework/model/generator/utils/taylorseer.py

releases/bridge-native-v3/code/README.md

@@ -0,0 +1,142 @@
+# zeva
+
+**In-Context Causal Learning for Generalizable Embodied Manipulation**
+
+[Project page](https://air-embodied-brain.github.io/Zeva) ·
+[Model weights](https://huggingface.co/chen123fu/zeva) ·
+[Atomic5 reproduction](docs/reproduce.md) ·
+[Cross-attempt reproduction](docs/zeva_pim.md)
+
+---
+
+## Overview
+
+zeva enables a frozen robot policy to learn from its own physical interaction
+experience at deployment time. Instead of updating model parameters, zeva
+extracts action-induced state changes, stores them in a dual-timescale causal
+memory, and retrieves relevant interaction evidence as context for subsequent
+actions.
+
+![zeva overview](assets/zeva_overview.png)
+
+## Method
+
+zeva implements **In-Context Causal Learning (ICCL)** in three parts:
+
+1. **Causal Interaction Extraction.** A Causal Transition Encoder (CTE)
+   integrates visual state, executed action, and observed state change. It
+   produces a phase token \(p_\tau\) and a causal interaction signal
+   \(e_\tau\). CTE is trained with causal-effect prediction, task-identity
+   clustering, and phase-progression objectives.
+2. **Dual-timescale Causal Memory.** A Brief Interaction Trace (BIT) retains
+   recent evidence within an attempt, while Persistent Interaction Memory (PIM)
+   consolidates phase/effect evidence across attempts in the same episode.
+3. **In-Context Policy Injection.** Phase-conditioned retrieval selects relevant
+   PIM entries. A memory encoder fuses the task, current phase, BIT, and retrieved
+   PIM evidence into a Causal Prompt \(M_\tau\) for the diffusion policy.
+
+> [!IMPORTANT]
+> During deployment, all neural parameters remain frozen. Only BIT and PIM are
+> updated online.
+
+## Installation
+
+zeva is released as an overlay for
+[NVIDIA Cosmos Framework](https://github.com/NVIDIA/cosmos-framework) at commit
+`ee58e41467f33c49ddde08b4d0ef4923876a95ac`.
+
+```bash
+git clone https://github.com/NVIDIA/cosmos-framework.git
+cd cosmos-framework
+git checkout ee58e41467f33c49ddde08b4d0ef4923876a95ac
+cd ..
+
+git clone https://github.com/dsdfasfaac/zeva-test.git zeva
+rsync -a --exclude=.git --exclude=README.md zeva/ cosmos-framework/
+cd cosmos-framework
+```
+
+Install the remaining Cosmos Framework dependencies following the upstream
+[setup instructions](https://github.com/NVIDIA/cosmos-framework/blob/main/docs/setup.md).
+
+## Reproduction
+
+> [!NOTE]
+> The released checkpoint was trained on success-only trajectories and
+> evaluated without cross-attempt self-evolution. Its benchmark success rate
+> therefore measures the frozen policy rather than improvement from accumulated
+> PIM experience.
+
+### Atomic5 benchmark reproduction
+
+Follow [docs/reproduce.md](docs/reproduce.md) to start the frozen policy server
+and reproduce the locked five-task evaluation. The released protocol uses a
+32-action horizon, 30 UniPC steps, guidance 3.0, shift 5.0, and no retries.
+
+### Fixed-seed cross-attempt evolution
+
+Follow [docs/zeva_pim.md](docs/zeva_pim.md) to run PIM-enabled evaluation on a
+fixed episode. Each case below was reproduced twice from empty PIM; `F` denotes
+failure and `S` denotes success.
+
+> [!NOTE]
+> The current release includes these reproducible case studies.
+
+| Task | Seed | Attempt sequence |
+| --- | ---: | --- |
+| OpenStandMixerHead | 199 | `FSSSS` |
+| TurnOnElectricKettle | 197 | `FSSSSS` |
+| CloseToasterOvenDoor | 199 | `FFFSSSSSSSS` |
+| TurnOnMicrowave | 200 | `FFFFSSSS` |
+| CoffeeSetupMug | 204 | `FFFFFFFFFSSSS` |
+
+Validate an exported pair of runs with:
+
+```bash
+python experiment_manifests/verify_pim_evolve_case.py \
+  --first /path/to/first_run \
+  --repeat /path/to/repeat_run \
+  --task TurnOnElectricKettle \
+  --seed 197
+```
+
+## Repository layout
+
+| Path | Description |
+| --- | --- |
+| `cosmos_framework/model/zeva/` | CTE, BIT, PIM, retrieval, Causal Prompt, and policy injection |
+| `cosmos_framework/scripts/` | Policy servers and RoboCasa365 evaluators |
+| `experiment_manifests/` | Locked evaluation and verification utilities |
+| `examples/` | Training launchers and experiment configurations |
+| `docs/reproduce.md` | Atomic5 reproduction protocol |
+| `docs/zeva_pim.md` | PIM and cross-attempt case studies |
+
+## Citation
+
+If you find this work useful, please cite:
+
+```bibtex
+@misc{zeva2026,
+  title={Zeva: In-Context Causal Learning for Generalizable Embodied Manipulation},
+  author={Fu Chen, Xin Ding, Bingjia Huang, Xiangyu Li, Mingju Wang, Jiawei He, Kun Li, Wei Sun, Yunxin Liu, Hao Wu, Ting Cao},
+  howpublished={\url{https://air-embodied-brain.github.io/Zeva}},
+  year={2026},
+}
+```
+
+## Acknowledgements
+
+zeva builds on the following open-source projects:
+
+- [NVIDIA Cosmos Framework](https://github.com/NVIDIA/cosmos-framework)
+- [RoboCasa365](https://github.com/robocasa/robocasa)
+- [BehaviorVLA](https://github.com/iLearn-Lab/ICML26-BehaviorVLA)
+
+Please follow the licenses and citation requirements of all external models,
+datasets, simulators, and software dependencies.
+
+## License
+
+This repository is released under the included [OpenMDW-1.1 license](LICENSE).
+External models, datasets, and simulators remain subject to their respective
+licenses.

releases/bridge-native-v3/code/configs/bridge_stage2_native_contract.json

@@ -0,0 +1,47 @@
+{
+  "schema": "bridge_native_stage2_design_v1",
+  "implementation_status": "IMPLEMENTED_GPU_VALIDATION_NOT_RUN",
+  "entrypoint": "python -m scripts.run_bridge_stage2_native",
+  "base_repo": "nvidia/Cosmos3-Nano",
+  "base_revision": "e59a53c25979a090fa8706c9acc0c254a6e89b92",
+  "framework_revision": "ee58e41467f33c49ddde08b4d0ef4923876a95ac",
+  "dataset_repo": "IPEC-COMMUNITY/bridge_orig_lerobot",
+  "dataset_revision": "0e9d76d07e9df3ea3eba257b2520d4913833fad2",
+  "domain_name": "bridge_orig_lerobot",
+  "domain_id": 7,
+  "policy_action_dim": 10,
+  "policy_gripper_index": 9,
+  "cte_action_dim": 7,
+  "cte_gripper_index": 6,
+  "proprio_dim": 10,
+  "policy_horizon": 32,
+  "expected_model_action_dim": 64,
+  "model_action_dim_source": "runtime configuration AND constructed action projections",
+  "global_dim": 256,
+  "phase_dim": 128,
+  "effect_dim": 128,
+  "effect_history_length": 4,
+  "boundary_controls": 4,
+  "completed_effect_controls": 16,
+  "cte_preprocess_id": "bridge_cte_image0_rgb_singleframe_256_bf16_wan22_v1",
+  "cte_normalizer_id": "bridge_reached_delta_legacy_quantile_v0",
+  "policy_normalizer": "NVIDIA bridge_orig_lerobot native10 quantile (unclipped)",
+  "proprio_normalizer": null,
+  "cte_source_step": 30000,
+  "static_bank_entries": 3000,
+  "retrieval": {"readout": "frozen Nano understanding-only image+instruction prefill", "head_dimensions": [4096, 1024, 128], "head_status": "NOT_TRAINED", "top_k": 5, "temperature": 0.07},
+  "learning_rates": {"generation": 0.00002, "action": 0.0001, "zeva": 0.00003},
+  "frozen": ["shared_reasoner", "shared_understanding", "semantic_vision_encoder", "Wan_VAE", "Bridge_CTE", "bank", "retrieval_head_after_training"],
+  "train_generation": ["moe_gen", "time_embedder", "vae2llm", "llm2vae"],
+  "train_action": ["action2llm_domain7", "llm2action_domain7", "action_modality_embed_shared"],
+  "train_zeva": ["proprio_projector_and_gate", "behavior_pbd", "behavior_adapter", "behavior_global_projector", "phase_effect_receivers"],
+  "action_table_weight_decay": 0.0,
+  "domain7_initialization": "original Nano checkpoint tensors; never reinitialize",
+  "other_domain_max_absolute_diff": 0.0,
+  "losses": {"mode": "wam", "vision_weight": 1.0, "action_weight": 10.0, "pbd_weight": 0.01, "ignore_tail_padding": true},
+  "gates": {"train_episodes": 32, "development_episodes": 8, "smoke_optimizer_steps": 50, "fit_optimizer_steps": 500, "both_required_before_continuous": true},
+  "checkpoint": {"save_every_optimizer_steps": 10000, "automatic_deletion": false, "save_frozen_base": false, "save_optimizer_and_data_cursor": true},
+  "forbidden_initialization": ["Cosmos3-Nano-Policy-DROID", "ZeVA RoboCasa Stage2", "RoboCasa Atomic5", "old ipec_full delta", "old Bridge 7D policy", "old optimizer state"],
+  "no_full_dataset_feature_cache": true,
+  "seed": 20260926
+}

releases/bridge-native-v3/code/cosmos_framework/__init__.py

@@ -0,0 +1,3 @@
+"""Zeva overlay package; extend imports to the upstream Cosmos framework."""
+from pkgutil import extend_path
+__path__ = extend_path(__path__, __name__)

releases/bridge-native-v3/code/cosmos_framework/configs/base/config.py

@@ -0,0 +1,109 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+from typing import Any, List
+
+import attrs
+
+from cosmos_framework.trainer import ImaginaireTrainer as Trainer
+from cosmos_framework.utils import config
+
+
+@attrs.define(slots=False)
+class DataSetting:
+    """Configuration for data.
+
+    Attributes:
+        qwen_max_video_token_length: Maximum video token length.
+        qwen_target_fps: Target fps for video sampling.
+        text_chat_order: Order of text items in user messages.
+    """
+
+    qwen_max_video_token_length: int = 8192
+
+
+@attrs.define(slots=False)
+class Config(config.Config):
+    data_setting: DataSetting = attrs.field(factory=DataSetting)
+    defaults: List[Any] = attrs.field(
+        factory=lambda: [
+            "_self_",
+            {"model": "mot_fsdp"},
+            {"data_train": None},
+            {"data_val": None},
+            {"optimizer": "adamw"},
+            {"scheduler": "warmup_cosine_lr"},
+            {"checkpoint": "s3"},
+            {"callbacks": ["basic", "optimization", "job_monitor", "generation"]},
+            {"ema": "power"},
+            {"tokenizer": "wan2pt2_tokenizer"},
+            {"sound_tokenizer": None},  # Optional: for audio-video generation
+            {"vlm_config": None},
+            {"ckpt_type": "dcp"},
+            {"experiment": None},
+        ]
+    )
+
+
+def make_config() -> Config:
+    c = Config(
+        model=None,
+        optimizer=None,
+        scheduler=None,
+        dataloader_train=None,
+        dataloader_val=None,
+    )
+
+    # Specifying values through instances of attrs
+    c.job.project = "cosmos3_vfm"
+    c.job.group = "debug"
+    c.job.name = "delete_${now:%Y-%m-%d}_${now:%H-%M-%S}"
+
+    c.trainer.type = Trainer
+    c.trainer.straggler_detection.enabled = False
+    c.trainer.max_iter = 400_000
+    c.trainer.logging_iter = 20
+    c.trainer.validation_iter = 100
+    c.trainer.run_validation = False
+    c.trainer.callbacks = None
+
+    c.upload_reproducible_setup = False
+
+    from cosmos_framework.configs.base.defaults.callbacks import register_callbacks
+    from cosmos_framework.configs.base.defaults.checkpointer import register_checkpoint, register_ckpt_type
+    from cosmos_framework.configs.base.defaults.ema import register_ema
+
+    # from cosmos_framework.configs.base.defaults.data import register_data
+    from cosmos_framework.configs.base.defaults.model import register_model
+    from cosmos_framework.configs.base.defaults.optimizer import register_optimizer, register_scheduler
+    from cosmos_framework.configs.base.defaults.tokenizer import register_sound_tokenizer, register_tokenizer
+    from cosmos_framework.configs.base.defaults.reasoner import register_vlm
+
+    # Call this function to register config groups for advanced overriding. the order follows the default config groups
+    # register_data()
+    register_model()
+    register_checkpoint()
+    register_ckpt_type()
+    register_optimizer()
+    register_scheduler()
+    register_callbacks()
+    register_tokenizer()
+    register_sound_tokenizer()
+    register_ema()
+    register_vlm()
+
+    # Register shipped experiments explicitly. (vision_sft_nano also defines
+    # vision_sft_nano_mapstyle_dataloader — the CosmosDataLoader variant — in the same module.)
+    import cosmos_framework.configs.base.experiment.action.posttrain_config.action_policy_droid_nano  # noqa: F401
+    import cosmos_framework.configs.base.experiment.action.posttrain_config.action_policy_robocasa365_atomic5_zeva  # noqa: F401
+    import cosmos_framework.configs.base.experiment.action.posttrain_config.action_policy_robocasa365_atomic5_zeva_transition_memory  # noqa: F401
+    import cosmos_framework.configs.base.experiment.action.posttrain_config.action_policy_robocasa365_atomic5_zeva_pim  # noqa: F401
+    import cosmos_framework.configs.base.experiment.action.posttrain_config.action_policy_robocasa365_atomic5_zeva_pim_inference  # noqa: F401
+    import cosmos_framework.configs.base.experiment.action.posttrain_config.action_policy_libero_all_nano  # noqa: F401
+    import cosmos_framework.configs.base.experiment.action.posttrain_config.action_policy_libero_nano  # noqa: F401
+    import cosmos_framework.configs.base.experiment.action.posttrain_config.action_fd_droid_posttrain  # noqa: F401
+    import cosmos_framework.configs.base.experiment.sft.vision_sft_nano  # noqa: F401
+    import cosmos_framework.configs.base.experiment.sft.vision_sft_super  # noqa: F401
+    import cosmos_framework.configs.base.experiment.sft.vision_sft_edge  # noqa: F401
+
+    return c

releases/bridge-native-v3/code/cosmos_framework/configs/base/defaults/model_config.py

@@ -0,0 +1,380 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+from typing import Any
+
+import attrs
+
+from cosmos_framework.utils.lazy_config import LazyDict
+from cosmos_framework.configs.base.defaults.activation_checkpointing import ActivationCheckpointingConfig
+from cosmos_framework.configs.base.defaults.compile import CompileConfig
+from cosmos_framework.configs.base.defaults.ema import EMAConfig
+from cosmos_framework.configs.base.defaults.parallelism import ParallelismConfig
+from cosmos_framework.configs.base.defaults.reasoner import VLMConfig
+
+
+@attrs.define(slots=False)
+class DiffusionExpertConfig:
+    # This determines the range of timesteps before the fourier feature embedding is applied.
+    timestep_range: float = 1.0
+    # Whether to load the generation pathway weights from pretrained LLM/VLM weights.
+    load_weights_from_pretrained: bool = True
+
+    patch_spatial: int = 2
+    max_vae_latent_side_after_patchify: int = (
+        20  # Max dimension (h or w) of the VAE latent after patchification (320/(8*2))
+    )
+    # Vision/action/sound position information is always provided through
+    # Qwen3VL-style 3D mRoPE attention IDs.
+    enable_fps_modulation: bool = False
+    base_fps: int = 24
+    # Base temporal compression factor for SOUND m-RoPE. None = current behavior
+    # (sound advances at base_fps positions/sec). Set to the vision tcf (4) to put
+    # sound on the same latent-frame temporal grid as vision/action.
+    sound_base_temporal_compression_factor: int | None = None
+    # Temporal coordinates used for unified_3d_mrope vision tokens.
+    # - "latent_index": legacy behavior, positions are 0, 1, ..., T_latent-1.
+    # - "uniae_source_right_edge": use UniAE padded-patch right-edge source-frame coordinates.
+    vision_temporal_position_mode: str = "latent_index"
+    # For unified_3d_mrope: whether spatial (H, W) indices reset to 0 for each vision segment
+    unified_3d_mrope_reset_spatial_ids: bool = True
+    # Setting the temporal gap on the boundary of the different modalities, default is 0, using a value greater than 0 will add an additional offset on the accumulated temporal offset.
+    unified_3d_mrope_temporal_modality_margin: int = 0
+
+
+@attrs.define(slots=False)
+class LBLConfig:
+    # For load balancing loss computation.
+    # - "local": Use the fraction of tokens routed to each expert only for the local rank.
+    # - "global": Use the fraction of tokens routed to each expert across all ranks.
+    method: str = "local"
+
+    # Coefficients for the load balancing loss.
+    # - "und": Coefficient for the load balancing loss for the "und" pathway.
+    # - "gen": Coefficient for the load balancing loss for the "gen" pathway.
+    coeff_und: float | None = None
+    coeff_gen: float | None = None
+
+
+@attrs.define(slots=False)
+class RectifiedFlowTrainingConfig:
+    shift: Any = 5  # Training time shift. If dict, maps resolution (str) to shift value (int)
+    use_dynamic_shift: bool = False  # Whether to use dynamic shifting
+    train_time_image_distribution: str = "logitnormal"  # Training time distribution for images
+    train_time_video_distribution: str = "logitnormal"  # Training time distribution for videos
+    train_time_action_distribution: str = "logitnormal"  # Training time distribution for actions
+    train_time_sound_distribution: str = "logitnormal"  # Training time distribution for sound
+    train_time_weight: str = "uniform"  # Training time weight
+    loss_scale: float = 1.0  # Loss scale
+    image_loss_scale: float | None = None  # If set, overrides loss_scale for images
+    sound_loss_scale: float | None = None  # If set, overrides loss_scale for sound
+    use_discrete_rf: bool = False  # Whether to use discrete formulation of rectified flow
+
+    # user: please adjust this value according to loss_scale to balance the action loss with the video loss.
+    # default is 10.0 to align with previous training settings.
+    action_loss_weight: float = 10.0
+
+    # Independent noise schedule for action. When False (default), action shares the sigma
+    # sampled from the vision RF on every step — legacy behavior. When True, action samples
+    # its own sigma from `rectified_flow_action` using `shift_action`. Action always uses a
+    # shared scalar sigma per sample ([B,1]), independent of vision's DF mode.
+    independent_action_schedule: bool = False
+    shift_action: int | None = None  # must be int; None → inherit `shift` (which must also be int)
+
+    # Independent noise schedule for sound. When False (default), sound shares the vision
+    # sigma schedule, reindexed to the dense audio-bearing subset. When True, sound samples
+    # its own scalar sigma per sample ([B,1]) from `rectified_flow_sound` using `shift_sound`.
+    independent_sound_schedule: bool = False
+    shift_sound: int | None = None  # must be int; None → inherit `shift` (which must also be int)
+
+    # When True, per-instance flow-matching loss is normalized by the count of
+    # active (noisy) elements rather than all elements — preserves sum/active_count
+    # semantics so conditioning-heavy samples (e.g. I2V, forward_dynamics, diffusion
+    # forcing, AR rollout teacher-forcing) contribute gradient on par with K=0
+    # samples. With .mean() the gradient of a K-conditioned sample is scaled by
+    # (T-K)/T, which undertrains the attend-to-clean-history dynamics. Kept
+    # False by default to preserve legacy loss magnitudes; enable for AR/DF training.
+    normalize_loss_by_active: bool = False
+
+    # Sample-level (vs rank-level) loss averaging for the vision modality.
+    #
+    # By default the vision loss on each rank is a mean over that rank's samples, and
+    # FSDP/DDP averages gradients across ranks — so every *rank* contributes equally
+    # regardless of how many image/video samples it holds ("rank-level averaging").
+    # When ranks carry different sample counts this over-weights samples on sparse ranks.
+    #
+    # When True, the loss is renormalized so every *sample* contributes equally across
+    # the whole data-parallel group: each iteration all-reduces the total number of image
+    # and video samples over the DP group, and image and video losses are each normalized
+    # by their own global sample count and then summed. This counteracts the framework's
+    # rank-level gradient averaging so the effective objective is a true per-sample mean.
+    # The two independently normalized terms are weighted by the existing `image_loss_scale`
+    # (images; falls back to `loss_scale` when None) and `loss_scale` (videos), so their
+    # balance is tuned with the same knobs as the legacy rank-level path.
+    sample_level_loss_averaging: bool = False
+
+
+@attrs.define(slots=False)
+class ZevaPolicyConfig:
+    """Configuration for Zeva causal-context policy conditioning.
+
+    Data paths intentionally live in the dataset recipe. This model config only
+    describes the learned policy-injection prior and is disabled by default, preserving
+    every existing Cosmos recipe.
+    """
+
+    enabled: bool = False
+    global_dim: int = 256
+    phase_dim: int = 128
+    # Ordered causal observed-effect history from Stage-1 effect-v3.
+    effect_dim: int = 128
+    effect_history_length: int = 4
+    action_dim: int = 8
+    horizon: int = 32
+    num_anchors: int = 8
+    hidden_dim: int = 256
+    num_heads: int = 4
+    prior_loss_weight: float = 0.01
+    prior_dropout_rate: float = 0.4
+    prior_inference_guidance_scale: float = 0.5
+    global_prefix_tokens: int = 1
+    # Number of leading action tokens that encode a clean initial state rather
+    # than policy commands. DROID/RoboTwin use one; RoboCasa365 uses zero
+    # because its 16-D proprio state is not isomorphic to its 12-D action.
+    leading_condition_steps: int = 1
+    # Optional same-task-session transition context. The context is already
+    # encoded by TransitionMemoryEncoder before entering the model.
+    online_memory_enabled: bool = False
+    online_context_dim: int = 256
+    online_prefix_tokens: int = 1
+    # Paper-aligned Persistent Interaction Memory.  PIM is injected as a
+    # gated residual into the existing global-behavior prefix, so enabling the
+    # module with gate_init=0 preserves the verified Stage-2 forward exactly.
+    pim_memory_enabled: bool = False
+    pim_persistent_length: int = 4
+    pim_context_dim: int = 256
+    pim_gate_init: float = 0.0
+    # Static inference-only kill switch. The PIM modules remain instantiated
+    # and load their trained weights, but the forward takes the original
+    # Stage-2 prefix path without executing the PIM adapter. This is stronger
+    # than an empty validity mask because it also preserves the compiled graph
+    # used by the frozen base policy.
+    pim_force_bypass: bool = False
+
+
+# Serialized releases may still reference this fully qualified type name.
+BehaviorStage2Config = ZevaPolicyConfig
+
+
+@attrs.define(slots=False)
+class ProprioConditionConfig:
+    """Optional low-dimensional robot-state prefix for action policies."""
+
+    enabled: bool = False
+    input_dim: int = 16
+    prefix_tokens: int = 1
+
+
+@attrs.define(slots=False)
+class RectifiedFlowInferenceConfig:
+    scheduler_type: str = "unipc"  # Scheduler type
+    num_train_timesteps: int = 1000
+    shift: int = 1
+    use_dynamic_shifting: bool = False
+
+
+@attrs.define(slots=False)
+class FixedStepSamplerConfig:
+    """Config for the fixed-step sampler used by distilled models.
+
+    Uses a fixed sigma schedule instead of a smooth multi-step solver.
+
+    Mirrors the constructor args of ``FixedStepSampler``.
+    """
+
+    # Discrete noise-level schedule (descending, excluding the final 0.0 step).
+    # Convention: exclude the final 0.0 step — FixedStepSampler appends it automatically.
+    # Values must be descending. Using 0.999 instead of 1.0 avoids numeric edge cases at sigma=1.
+    t_list: list[float] = [0.999, 0.75, 0.5, 0.25]
+    # Distilled fixed-step sampling uses stochastic re-noising at each step.
+    sample_type: str = "sde"
+
+
+# Don't have any defaults and init only in config file.
+@attrs.define(slots=False)
+class OmniMoTModelConfig:
+    """
+    Config for Omni MoT model.
+    """
+
+    tokenizer: LazyDict = None
+    load_vision_tokenizer: bool = True
+    """Instantiate the vision tokenizer.
+
+    Reasoner-only inference disables this to avoid loading the generation VAE.
+    """
+    net: LazyDict = None
+    ema: EMAConfig = EMAConfig()
+
+    # Parallelism (CP, CFGP, FSDP, DP) and FSDP reduce-dtype configuration.
+    parallelism: ParallelismConfig = ParallelismConfig()
+
+    # torch.compile knobs (enabled, compiled_region, dynamic, ...).
+    compile: CompileConfig = CompileConfig()
+
+    # Activation-checkpointing policy (trade-off between memory and speed).
+    activation_checkpointing: ActivationCheckpointingConfig = ActivationCheckpointingConfig()
+
+    # Model parameter / activation dtype (consumed by MixedPrecisionPolicy and
+    # ``model.precision`` for LowPrecisionCallback). One of "bfloat16",
+    # "float16", "float32".
+    precision: str = "bfloat16"
+
+    # LoRA (parameter-efficient fine-tuning). When `lora_enabled=True`,
+    # `OmniMoTModel.build_net` injects custom LoRA adapters BEFORE FSDP wrap on
+    # the meta-device network, then re-initializes lora_A/lora_B after
+    # to_empty + init_weights. Pair with `optimizer.keys_to_select=["lora_"]`
+    # and `checkpoint.keys_to_skip_loading=[..., "lora_"]`.
+    lora_enabled: bool = False
+    lora_rank: int = 16
+    lora_alpha: int = 32
+    lora_target_modules: str = "q_proj_moe_gen,k_proj_moe_gen,v_proj_moe_gen,o_proj_moe_gen"
+    # Optional base-policy teacher anchor for LoRA post-training. The teacher
+    # is the same frozen network with adapters temporarily disabled, so this
+    # does not allocate a second copy of the model. The coefficient is applied
+    # in the same units as action_loss_weight; zero preserves legacy behavior.
+    lora_teacher_anchor_weight: float = 0.0
+    lora_teacher_anchor_data_key: str = "teacher_anchor_weight"
+
+    # Frozen-base action correction expert.  When enabled, ``build_net``
+    # freezes every base parameter and trains only the zero-initialized
+    # ``action_residual_head``.  The head is never merged into the base.
+    residual_expert: dict[str, Any] = attrs.Factory(dict)
+
+    # R3 correction-policy controls.  These are intentionally kept in the
+    # model config (rather than in the serving request) so an exported policy
+    # is a single always-on model: no phase/predicate/task router is required
+    # at inference time.  ``event_weighted`` only changes the training loss.
+    r3_event_weighted: bool = False
+    r3_event_weight: float = 2.0
+    r3_task_weight_key: str = "r3_task_weight"
+    r3_event_weight_key: str = "residual_release_event"
+
+    # Rectified flow configs
+    rectified_flow_training_config: RectifiedFlowTrainingConfig = RectifiedFlowTrainingConfig()
+    rectified_flow_inference_config: RectifiedFlowInferenceConfig = RectifiedFlowInferenceConfig()
+    # The field name is serialized in released configs; the type and all new
+    # public APIs use the Zeva paper terminology.
+    behavior_stage2: ZevaPolicyConfig = ZevaPolicyConfig()
+    proprio_condition: ProprioConditionConfig = ProprioConditionConfig()
+
+    # Optional fixed-step sampler for distilled models (None for base models).
+    fixed_step_sampler_config: FixedStepSamplerConfig | None = None
+
+    # Model configs
+    vlm_config: VLMConfig = VLMConfig()
+    diffusion_expert_config: DiffusionExpertConfig = DiffusionExpertConfig()
+    # Training data keys
+    input_video_key: str = "video"
+    input_image_key: str = "images"  # key to fetch input image from data_batch
+    input_caption_key: str = "ai_caption"  # Key used to fetch input captions
+
+    # State and sequence shapes
+    state_ch: int = 16  # for latent model, ref to the latent channel number
+    state_t: int = 8  # for latent model, ref to the latent number of frames
+    latent_downsample_factor: int = 8
+    resolution: str = "512"
+    max_num_tokens_after_packing: int = 13312  # Final num tokens after sequence packing
+
+    # Attention implementation for joint understanding + generation
+    # Note "two_way" and "three_way" disallow and remove "End-of-Vision" or other text token in the generation tower.
+    # "three_way" must only be used when introducing sparsity
+    joint_attn_implementation: str = "two_way"  # "two_way" or "three_way"
+
+    # Per-layer NATTEN parameters
+    # Must use "three_way" attention if used.
+    # If None, all attention layers remain dense.
+    # If not None, must be a list exactly the size of number of layers, and each layer can be either
+    # None (dense) or a dictionary, with at least 'kernel_size' or 'kernel_size_float' keys
+    # specifying sparsity. NATTEN parameters 'dilation' and 'stride' may also be specified either as
+    # static integers, or as floating point values that will be mapped to their domain during
+    # runtime. Integer parameters should never be mixed with floating point ones.
+    #
+    # Floating point parameters are highly recommended, unless the use case will have a fixed token
+    # layout (input resolution).
+    #
+    # Examples:
+    #   Interleaved sliding window layers, "GPT-OSS"-style, with static window size:
+    #     natten_parameter_list = [None if layer_idx % 2 != 0 else {"kernel_size": (8, 8)}]
+    #   Layers with odd indices ("None"s) will use dense attention, and layers with an even indices
+    #   will use a static sliding window size of 8x8.
+    #
+    #   Interleaved sliding window layers, "GPT-OSS"-style, with input-dependent window size:
+    #     natten_parameter_list = [None if layer_idx % 2 != 0 else {"kernel_size_float": (0.5, 0.5)}]
+    #   Layers with odd indices ("None"s) will use dense attention, and layers with an even indices
+    #   will use a dynamic window size that is 50% of the input along each of the two dimensions.
+    #
+    #   Interleaved sliding window and dilated layers, "DiNAT"-style:
+    #     natten_parameter_list = [
+    #       {
+    #           "kernel_size_float": (0.5, 0.5),
+    #           "dilation_float": (1.0, 1.0),
+    #       } if layer_idx % 2 != 0 else {
+    #           "kernel_size_float": (0.5, 0.5),
+    #       }
+    #     ]
+    #   All layers will use a dynamic window size that is 50% of the input along each of the two
+    #   dimensions. Layers with odd indices will also dilate to the maximum level possible.
+    #
+    natten_parameter_list: list | None = None
+
+    # Temporal causality for training autoregressive video generation models.
+    # When enabled, applies temporal causal attention to generation supertokens.
+    # Each supertoken is num_action_tokens_per_supertoken action tokens followed
+    # by H*W vision tokens; the value is stamped onto the packed sequence by the
+    # temporal-causal packer and read by attention/KV-cache code unchanged.
+    # Only supports image2video modes (with or without actions).
+    # Requires joint_attn_implementation="three_way".
+    video_temporal_causal: bool = False
+    # "none":             standard joint denoising (shared σ, no clean context)
+    # "teacher_forcing":  all frames noised with shared σ; clean history via cross-attention
+    # "diffusion_forcing": each latent frame gets independent σ ~ Uniform[0,1]
+    # "teacher_forcing_dcm": replayed teacher-forcing discrete-time consistency distillation
+    causal_training_strategy: str = attrs.field(
+        default="none",
+        validator=attrs.validators.in_({"none", "teacher_forcing", "diffusion_forcing", "teacher_forcing_dcm"}),
+    )
+
+    # Load balancing loss config.
+    lbl: LBLConfig = LBLConfig()
+
+    # vision configs
+    vision_gen: bool = True  # whether to use vision related parameters and condition/generate vision tokens
+
+    # action configs
+    action_gen: bool = False  # whether to use action related parameters and condition/generate action tokens
+    max_action_dim: int = 32  # maximum dimension of the action space, we need to pad the data to this dimension.
+    num_embodiment_domains: int = 32  # number of domains for the domain-aware linear layer
+
+    # sound configs
+    sound_gen: bool = False  # whether to use sound related parameters and condition/generate sound tokens
+    sound_tokenizer: LazyDict | None = None  # Sound tokenizer config (e.g., AVAE)
+    sound_dim: int | None = None  # Sound latent channel size (e.g., 64 for AVAE 48kHz)
+    sound_latent_fps: int = 25  # Sound tokenizer's latent rate (e.g., 48kHz / 1920 hop = 25 Hz)
+
+    # When False, removes bias from vae2llm, sound2llm, and the two Linear layers inside
+    # time_embedder.  These biases seem to inject token-constant DC offsets that dominate
+    # the MoE router input and create prompt invariant routing.  This is observed empirically
+    # in the 30B-A3B checkpoint.
+    enable_input_bias: bool = True
+
+    log_enc_time_every_n: int = 100  # Frequency of logging encoding time to W&B
+
+    # When True, ``OmniMoTModel.state_dict`` / ``load_state_dict`` skip the
+    # reasoner (und) pathway weights under ``language_model`` — i.e. every key
+    # WITHOUT a ``_moe_gen`` suffix (including ``visual`` / ``lm_head`` /
+    # ``embed_tokens``).  These are not written to checkpoints and are left
+    # untouched on load (typically already populated from the HF pretrained
+    # backbone).  Generation-pathway (``_moe_gen``) and VFM heads are saved /
+    # loaded as usual.
+    exclude_reasoner_weights_from_checkpoint: bool = False

releases/bridge-native-v3/code/cosmos_framework/configs/base/defaults/vlm.py

@@ -0,0 +1,17 @@
+"""Compatibility shim: behavior branch renamed defaults/vlm -> defaults/reasoner.
+Hydra/export may still locate class paths under defaults.vlm.* for checkpoints
+serialized before the rename; forward them to the reasoner definitions.
+"""
+from cosmos_framework.configs.base.defaults.reasoner import (
+    PretrainedWeightsConfig,
+    VLMConfig,
+    create_qwen2_tokenizer_with_download,
+    create_vlm_config,
+)
+
+__all__ = [
+    "PretrainedWeightsConfig",
+    "VLMConfig",
+    "create_qwen2_tokenizer_with_download",
+    "create_vlm_config",
+]

releases/bridge-native-v3/code/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_robocasa365_atomic5_zeva.py

@@ -0,0 +1,65 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""Inference recipe for the fixed-base RoboCasa Atomic-5 Zeva policy."""
+
+from __future__ import annotations
+
+import copy
+
+from hydra.core.config_store import ConfigStore
+
+from cosmos_framework.configs.base.experiment.action.posttrain_config.action_policy_droid_nano import (
+    action_policy_droid_nano,
+)
+
+cs = ConfigStore.instance()
+
+action_policy_robocasa365_atomic5_zeva = copy.deepcopy(action_policy_droid_nano)
+action_policy_robocasa365_atomic5_zeva["job"].update(
+    project="zeva",
+    group="zeva_robocasa_atomic5",
+    name="action_policy_robocasa365_atomic5_zeva",
+)
+action_policy_robocasa365_atomic5_zeva["model"]["config"]["behavior_stage2"] = dict(
+    enabled=True,
+    global_dim=256,
+    phase_dim=128,
+    effect_dim=128,
+    effect_history_length=4,
+    action_dim=7,
+    horizon=32,
+    num_anchors=8,
+    hidden_dim=256,
+    num_heads=4,
+    prior_loss_weight=0.01,
+    prior_dropout_rate=0.4,
+    prior_inference_guidance_scale=0.5,
+    global_prefix_tokens=1,
+    leading_condition_steps=0,
+)
+action_policy_robocasa365_atomic5_zeva["model"]["config"]["proprio_condition"] = dict(
+    enabled=True,
+    input_dim=9,
+    prefix_tokens=1,
+)
+# This is a post-trained Zeva checkpoint, not a fresh DROID initialization.
+# Load every trained policy/action tensor; only the unused EMA namespace may be
+# absent. Inheriting DROID's action-head skip list silently changes inference.
+action_policy_robocasa365_atomic5_zeva["checkpoint"]["keys_to_skip_loading"] = ["net_ema."]
+# The server supplies Qwen explicitly; retain the same environment contract as
+# the training recipe so checkpoint metadata and manual startup agree.
+action_policy_robocasa365_atomic5_zeva["model"]["config"]["vlm_config"]["tokenizer"][
+    "pretrained_model_name"
+] = "${oc.env:QWEN_VLM_PATH}"
+
+# Serving does not require training dataloaders.
+action_policy_robocasa365_atomic5_zeva["dataloader_train"] = None
+action_policy_robocasa365_atomic5_zeva["dataloader_val"] = None
+
+cs.store(
+    group="experiment",
+    package="_global_",
+    name="action_policy_robocasa365_atomic5_zeva",
+    node=action_policy_robocasa365_atomic5_zeva,
+)

releases/bridge-native-v3/code/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_robocasa365_atomic5_zeva_pim.py

@@ -0,0 +1,61 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""Frozen-policy PIM training for RoboCasa Atomic-5 Zeva."""
+
+from __future__ import annotations
+
+import copy
+
+from hydra.core.config_store import ConfigStore
+
+from cosmos_framework.configs.base.experiment.action.posttrain_config.action_policy_robocasa365_atomic5_zeva import (
+    action_policy_robocasa365_atomic5_zeva,
+)
+
+cs = ConfigStore.instance()
+
+action_policy_robocasa365_atomic5_zeva_pim = copy.deepcopy(
+    action_policy_robocasa365_atomic5_zeva
+)
+action_policy_robocasa365_atomic5_zeva_pim["job"].update(
+    project="zeva",
+    group="pim_adapter_robocasa_atomic5",
+    name="action_policy_robocasa365_atomic5_zeva_pim",
+)
+zeva_policy = action_policy_robocasa365_atomic5_zeva_pim["model"]["config"]["behavior_stage2"]
+zeva_policy.update(
+    pim_memory_enabled=True,
+    pim_persistent_length=4,
+    pim_context_dim=256,
+    pim_gate_init=0.0,
+)
+
+# Train only the Causal Prompt encoder, projection, and scalar gate.
+keys = action_policy_robocasa365_atomic5_zeva_pim["optimizer"]["keys_to_select"]
+keys[:] = ["behavior_pim_encoder", "behavior_pim_projector", "behavior_pim_gate"]
+action_policy_robocasa365_atomic5_zeva_pim["optimizer"]["lr_multipliers"].update(
+    behavior_pim_encoder=5.0,
+    behavior_pim_projector=5.0,
+    behavior_pim_gate=1.0,
+)
+action_policy_robocasa365_atomic5_zeva_pim["checkpoint"]["keys_to_skip_loading"] += [
+    "behavior_pim_encoder",
+    "behavior_pim_projector",
+    "behavior_pim_gate",
+]
+
+train_loader = action_policy_robocasa365_atomic5_zeva_pim.get("dataloader_train")
+if train_loader is not None:
+    dataset = train_loader["dataloader"]["datasets"]["robocasa365"]["dataset"]
+    dataset["behavior_pim_training_bank"] = "${oc.env:ZEVA_PIM_TRAINING_BANK}"
+    dataset["behavior_pim_history_length"] = 4
+    dataset["behavior_pim_context_dropout"] = 0.2
+    dataset["behavior_pim_support_dropout"] = 0.2
+
+cs.store(
+    group="experiment",
+    package="_global_",
+    name="action_policy_robocasa365_atomic5_zeva_pim",
+    node=action_policy_robocasa365_atomic5_zeva_pim,
+)

releases/bridge-native-v3/code/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_robocasa365_atomic5_zeva_pim_inference.py

@@ -0,0 +1,36 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""Inference config that loads trained PIM weights instead of warm-start skipping them."""
+
+from __future__ import annotations
+
+import copy
+
+from hydra.core.config_store import ConfigStore
+
+from cosmos_framework.configs.base.experiment.action.posttrain_config.action_policy_robocasa365_atomic5_zeva_pim import (
+    action_policy_robocasa365_atomic5_zeva_pim,
+)
+
+cs = ConfigStore.instance()
+
+action_policy_robocasa365_atomic5_zeva_pim_inference = copy.deepcopy(
+    action_policy_robocasa365_atomic5_zeva_pim
+)
+action_policy_robocasa365_atomic5_zeva_pim_inference["job"]["name"] = (
+    "action_policy_robocasa365_atomic5_zeva_pim_inference"
+)
+# The training recipe skips these keys only because its input checkpoint is the
+# old Stage-2 model.  A trained PIM DCP contains them and inference must load
+# them.  Keep the established net_ema exclusion used by the Stage-2 release.
+action_policy_robocasa365_atomic5_zeva_pim_inference["checkpoint"][
+    "keys_to_skip_loading"
+] = ["net_ema."]
+
+cs.store(
+    group="experiment",
+    package="_global_",
+    name="action_policy_robocasa365_atomic5_zeva_pim_inference",
+    node=action_policy_robocasa365_atomic5_zeva_pim_inference,
+)

releases/bridge-native-v3/code/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_robocasa365_atomic5_zeva_test.py

@@ -0,0 +1,14 @@
+from cosmos_framework.configs.base.experiment.action.posttrain_config.action_policy_robocasa365_atomic5_zeva import (
+    action_policy_robocasa365_atomic5_zeva,
+)
+
+
+def test_released_stage2_loads_trained_action_heads() -> None:
+    assert action_policy_robocasa365_atomic5_zeva["checkpoint"][
+        "keys_to_skip_loading"
+    ] == ["net_ema."]
+
+
+def test_inference_recipe_does_not_construct_training_loaders() -> None:
+    assert action_policy_robocasa365_atomic5_zeva["dataloader_train"] is None
+    assert action_policy_robocasa365_atomic5_zeva["dataloader_val"] is None

releases/bridge-native-v3/code/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_robocasa365_atomic5_zeva_transition_memory.py

@@ -0,0 +1,66 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""Experimental transition-memory adaptation for the RoboCasa Atomic-5 Zeva policy.
+
+The CTE path and frozen diffusion policy are loaded from the verified
+Atomic-5 Stage-2 checkpoint and frozen.  The separately trained
+history-only ``TransitionMemoryEncoder`` is consumed through the cached
+``behavior_online_context`` token; this recipe trains the zero-init Cosmos
+transition-memory projector so that the new context affects only the action branch.
+"""
+
+from __future__ import annotations
+
+import copy
+
+from hydra.core.config_store import ConfigStore
+
+from cosmos_framework.configs.base.experiment.action.posttrain_config.action_policy_robocasa365_atomic5_zeva import (
+    action_policy_robocasa365_atomic5_zeva,
+)
+
+cs = ConfigStore.instance()
+
+action_policy_robocasa365_atomic5_zeva_transition_memory = copy.deepcopy(
+    action_policy_robocasa365_atomic5_zeva
+)
+action_policy_robocasa365_atomic5_zeva_transition_memory["job"].update(
+    project="zeva",
+    group="transition_memory_robocasa_atomic5",
+    name="action_policy_robocasa365_atomic5_zeva_transition_memory",
+)
+zeva_policy = action_policy_robocasa365_atomic5_zeva_transition_memory["model"]["config"]["behavior_stage2"]
+zeva_policy.update(
+    online_memory_enabled=True,
+    online_context_dim=256,
+    online_prefix_tokens=1,
+)
+
+# The already trained policy-injection prior and task-context projector are deliberately not added
+# again.  Only the new zero-init Cosmos action-branch projector is optimized;
+# the query/key/value/history encoder is trained by the preceding
+# history-only recipe and loaded into the context cache.
+keys = action_policy_robocasa365_atomic5_zeva_transition_memory["optimizer"]["keys_to_select"]
+# Stage-2 checkpoint parameters and the Cosmos backbone must remain frozen in
+# this adaptation.  Replacing (rather than subtracting from) the inherited
+# selector also prevents the base RoboCasa recipe's action/proprio modules from
+# being silently optimized.
+keys[:] = ["behavior_online_projector"]
+action_policy_robocasa365_atomic5_zeva_transition_memory["optimizer"]["lr_multipliers"].update(
+    behavior_online_projector=5.0,
+)
+
+train_loader = action_policy_robocasa365_atomic5_zeva_transition_memory.get("dataloader_train")
+if train_loader is not None:
+    dataset = train_loader["dataloader"]["datasets"]["robocasa365"]["dataset"]
+    dataset["behavior_memory_bank"] = "${oc.env:ZEVA_TASK_CONTEXT_BANK}"
+    dataset["behavior_phase_cache"] = "${oc.env:ZEVA_CTE_FEATURE_CACHE}"
+    dataset["behavior_online_context_cache"] = "${oc.env:ZEVA_TRANSITION_CONTEXT_CACHE}"
+
+cs.store(
+    group="experiment",
+    package="_global_",
+    name="action_policy_robocasa365_atomic5_zeva_transition_memory",
+    node=action_policy_robocasa365_atomic5_zeva_transition_memory,
+)

releases/bridge-native-v3/code/cosmos_framework/data/__init__.py

@@ -0,0 +1,2 @@
+from pkgutil import extend_path
+__path__ = extend_path(__path__, __name__)

releases/bridge-native-v3/code/cosmos_framework/data/generator/__init__.py

@@ -0,0 +1,2 @@
+from pkgutil import extend_path
+__path__ = extend_path(__path__, __name__)

releases/bridge-native-v3/code/cosmos_framework/data/generator/action/__init__.py

@@ -0,0 +1,2 @@
+from pkgutil import extend_path
+__path__ = extend_path(__path__, __name__)

releases/bridge-native-v3/code/cosmos_framework/data/generator/action/datasets/__init__.py

@@ -0,0 +1,15 @@
+from pkgutil import extend_path
+__path__ = extend_path(__path__, __name__)
+
+# Preserve the upstream package's public dataset exports.
+from .agibotworld_beta_lerobot_dataset import AgiBotWorldBetaLeRobotDataset
+from .base_dataset import ActionBaseDataset
+from .bridge_orig_lerobot_dataset import BridgeOrigLeRobotDataset
+from .droid_merged_lerobot_dataset import DROIDMergedLeRobotDataset
+from .droid_lerobot_dataset import DROIDLeRobotDataset
+from .fractal_lerobot_dataset import FractalLeRobotDataset
+from .human_hand_pose_lerobot_dataset import HumanHandPoseLeRobotDataset
+from .libero_lerobot_dataset import LIBEROLeRobotDataset
+from .robomind_franka_dataset import RoboMINDFrankaDataset
+from .robomind_ur_dataset import RoboMINDURDataset
+from .umi_lerobot_dataset import UMILeRobotDataset

releases/bridge-native-v3/code/cosmos_framework/data/generator/action/datasets/robocasa365_lerobot_dataset.py

@@ -0,0 +1,334 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""Read-only RoboCasa365 LeRobot-v2.1 adapter for Cosmos3 Action.
+
+The target set is a directory of 50 independent LeRobot roots.  Current
+LeRobot rejects v2.1 metadata, so this reader indexes the immutable JSONL and
+per-episode parquet files directly while using LeRobot's torchcodec decoder.
+"""
+
+from __future__ import annotations
+
+import bisect
+import json
+import random
+import re
+from collections import OrderedDict
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Sequence
+
+import pyarrow.parquet as pq
+import torch
+import torch.nn.functional as F
+from lerobot.datasets.video_utils import decode_video_frames
+from torch.utils.data import Dataset
+
+from cosmos_framework.data.generator.action.action_processing import ActionNormalizer, resolve_action_normalization
+from cosmos_framework.data.generator.action.datasets.cosmos3_action_lerobot import split_episode_ids
+from cosmos_framework.data.generator.action.domain_utils import get_domain_id
+
+_CAMERAS = {
+    "left": "observation.images.robot0_agentview_left",
+    "right": "observation.images.robot0_agentview_right",
+    "wrist": "observation.images.robot0_eye_in_hand",
+}
+
+_CAMERA_LAYOUTS = {
+    "three_view_grid": ("left", "right", "wrist"),
+    "left_wrist_horizontal": ("left", "wrist"),
+}
+
+_ACTION_INDICES = {
+    "full12": tuple(range(12)),
+    # RoboCasa365 layout: base xy-yaw-height (0:4), control mode (4),
+    # end-effector delta pose + gripper (5:12).
+    "arm7": tuple(range(5, 12)),
+}
+
+_STATE_INDICES = {
+    "full16": tuple(range(16)),
+    # RoboCasa365 state: base pose (0:7), relative end-effector pose (7:14),
+    # gripper qpos (14:16). Fixed-base policies keep only the latter 9 values.
+    "arm9": tuple(range(7, 16)),
+}
+
+
+def _humanize_task(task: str) -> str:
+    words = re.sub(r"(?<!^)(?=[A-Z])", " ", task).strip()
+    return (words[:1].upper() + words[1:].lower() + ".") if words else task
+
+
+@dataclass(frozen=True)
+class RoboCasa365Episode:
+    root: Path
+    category: str
+    task_name: str
+    episode_id: int
+    length: int
+    command: str
+    parquet_path: Path
+    video_paths: dict[str, Path]
+
+
+class RoboCasa365LeRobotDataset(Dataset):
+    """Cosmos policy windows over the 50-task RoboCasa365 target dataset."""
+
+    EMBODIMENT_TYPE = "robocasa-panda-omron"
+
+    def __init__(
+        self,
+        root: str,
+        *,
+        fps: float = 20.0,
+        chunk_length: int = 32,
+        split: str = "train",
+        split_seed: int = 42,
+        split_val_ratio: float = 0.03,
+        mode: str = "wam",
+        viewpoint: str = "concat_view",
+        use_state: bool = False,
+        use_image_augmentation: bool = False,
+        emit_behavior_metadata: bool = False,
+        task_set: str = "target50",
+        task_names: Sequence[str] | None = None,
+        camera_layout: str = "three_view_grid",
+        action_mode: str = "full12",
+        state_mode: str = "full16",
+        action_normalization: str | None = "minmax",
+        action_stats_path: str | None = None,
+        parquet_cache_size: int = 8,
+        max_roots: int | None = None,
+    ) -> None:
+        if viewpoint not in {"concat_view", "third_person_view"}:
+            raise ValueError(f"Unsupported RoboCasa365 viewpoint={viewpoint!r}")
+        if split not in {"train", "val", "full"}:
+            raise ValueError("split must be train, val or full")
+        if task_set not in {"target50", "atomic", "composite", "composite_seen", "composite_unseen"}:
+            raise ValueError(f"Unsupported RoboCasa365 task_set={task_set!r}")
+        if camera_layout not in _CAMERA_LAYOUTS:
+            raise ValueError(f"Unsupported RoboCasa365 camera_layout={camera_layout!r}")
+        if action_mode not in _ACTION_INDICES:
+            raise ValueError(f"Unsupported RoboCasa365 action_mode={action_mode!r}")
+        if state_mode not in _STATE_INDICES:
+            raise ValueError(f"Unsupported RoboCasa365 state_mode={state_mode!r}")
+        self.root = Path(root)
+        self.fps = float(fps)
+        self.chunk_length = int(chunk_length)
+        self.split = split
+        self.split_seed = int(split_seed)
+        self.split_val_ratio = float(split_val_ratio)
+        self.mode = mode
+        self.viewpoint = viewpoint
+        self.use_state = bool(use_state)
+        self.use_image_augmentation = bool(use_image_augmentation)
+        self.emit_behavior_metadata = bool(emit_behavior_metadata)
+        self.task_set = task_set
+        self.task_names = frozenset(task_names) if task_names is not None else None
+        self.camera_layout = camera_layout
+        self.action_mode = action_mode
+        self.action_indices = _ACTION_INDICES[action_mode]
+        self.state_mode = state_mode
+        self.state_indices = _STATE_INDICES[state_mode]
+        self.camera_names = _CAMERA_LAYOUTS[camera_layout]
+        self.domain_id = get_domain_id(self.EMBODIMENT_TYPE)
+        self.action_dim = len(self.action_indices)
+        self._parquet_cache_size = int(parquet_cache_size)
+        self._parquet_cache: OrderedDict[Path, dict[str, torch.Tensor]] = OrderedDict()
+        self._normalizer: ActionNormalizer | None = None
+        if action_normalization is not None:
+            stats_path = (
+                Path(action_stats_path)
+                if action_stats_path
+                else Path(__file__).parents[1] / "normalizer_stats" / "robocasa365_target_action_stats.json"
+            )
+            if not stats_path.is_file():
+                raise FileNotFoundError(
+                    "RoboCasa365 action normalization requires the target50 global action statistics; "
+                    f"missing {stats_path}"
+                )
+            raw = json.loads(stats_path.read_text())
+            raw = raw.get("action", raw)
+            index = torch.tensor(self.action_indices, dtype=torch.long)
+            stats = {
+                key: torch.tensor(raw[key], dtype=torch.float32).index_select(0, index)
+                for key in ("min", "max", "mean", "std", "q01", "q99")
+            }
+            self._normalizer = resolve_action_normalization(action_normalization, stats)
+
+        roots = sorted(path.parent.parent for path in self.root.glob("*/*/*/lerobot/meta/info.json"))
+        roots = [path for path in roots if self._include_root(path)]
+        if max_roots is not None:
+            roots = roots[:max_roots]
+        if not roots:
+            raise FileNotFoundError(f"No RoboCasa365 LeRobot-v2.1 roots under {self.root}")
+        self.episodes: list[RoboCasa365Episode] = []
+        for root_path in roots:
+            self._index_root(root_path)
+        # A length-H command needs H action rows and H+1 observation frames.
+        self._window_lengths = [max(0, episode.length - self.chunk_length) for episode in self.episodes]
+        self._cum_ends: list[int] = []
+        total = 0
+        for length in self._window_lengths:
+            total += length
+            self._cum_ends.append(total)
+        self._num_windows = total
+
+    def _include_root(self, path: Path) -> bool:
+        category, task_name = path.relative_to(self.root).parts[:2]
+        if self.task_names is not None and task_name not in self.task_names:
+            return False
+        if self.task_set == "atomic":
+            return category == "atomic"
+        if self.task_set == "composite":
+            return category == "composite"
+        if self.task_set == "composite_seen":
+            return category == "composite" and task_name not in _COMPOSITE_UNSEEN
+        if self.task_set == "composite_unseen":
+            return category == "composite" and task_name in _COMPOSITE_UNSEEN
+        return True
+
+    def _index_root(self, path: Path) -> None:
+        info = json.loads((path / "meta" / "info.json").read_text())
+        if info["codebase_version"] != "v2.1" or info["fps"] != self.fps:
+            raise ValueError(f"Unexpected RoboCasa365 metadata in {path}")
+        category, task_name = path.relative_to(self.root).parts[:2]
+        records = [json.loads(line) for line in (path / "meta" / "episodes.jsonl").read_text().splitlines()]
+        selected = split_episode_ids(len(records), self.split_seed, self.split_val_ratio, self.split)
+        selected.sort()
+        for episode_id in selected:
+            record = records[episode_id]
+            length = int(record["length"])
+            if length <= self.chunk_length:
+                continue
+            command = str(record.get("tasks", [""])[0]).strip() or _humanize_task(task_name)
+            video_paths = {
+                name: path / f"videos/chunk-{episode_id // 1000:03d}/{key}/episode_{episode_id:06d}.mp4"
+                for name, key in _CAMERAS.items()
+                if name in self.camera_names
+            }
+            parquet_path = path / f"data/chunk-{episode_id // 1000:03d}/episode_{episode_id:06d}.parquet"
+            if not parquet_path.is_file() or any(not video.is_file() for video in video_paths.values()):
+                raise FileNotFoundError(f"Incomplete RoboCasa365 episode {path}#{episode_id}")
+            self.episodes.append(
+                RoboCasa365Episode(path, category, task_name, episode_id, length, command, parquet_path, video_paths)
+            )
+
+    def __len__(self) -> int:
+        return self._num_windows
+
+    def get_shuffle_blocks(self) -> list[tuple[int, int]]:
+        blocks: list[tuple[int, int]] = []
+        start = 0
+        for length in self._window_lengths:
+            blocks.append((start, length))
+            start += length
+        return blocks
+
+    def _resolve_index(self, index: int) -> tuple[int, int]:
+        if index < 0:
+            index += len(self)
+        if not 0 <= index < len(self):
+            raise IndexError(index)
+        episode_index = bisect.bisect_right(self._cum_ends, index)
+        start = 0 if episode_index == 0 else self._cum_ends[episode_index - 1]
+        return episode_index, index - start
+
+    def _load_low_dim(self, episode: RoboCasa365Episode) -> dict[str, torch.Tensor]:
+        cached = self._parquet_cache.pop(episode.parquet_path, None)
+        if cached is None:
+            table = pq.read_table(episode.parquet_path, columns=["action", "observation.state"])
+            payload = table.to_pydict()
+            cached = {
+                "action": torch.tensor(payload["action"], dtype=torch.float32),
+                "state": torch.tensor(payload["observation.state"], dtype=torch.float32),
+            }
+        self._parquet_cache[episode.parquet_path] = cached
+        while len(self._parquet_cache) > self._parquet_cache_size:
+            self._parquet_cache.popitem(last=False)
+        return cached
+
+    def _decode(self, path: Path, start: int) -> torch.Tensor:
+        timestamps = [(start + offset) / self.fps for offset in range(self.chunk_length + 1)]
+        return decode_video_frames(path, timestamps, tolerance_s=2e-4, backend="torchcodec")
+
+    def _compose_video(self, episode: RoboCasa365Episode, start: int) -> torch.Tensor:
+        left = self._decode(episode.video_paths["left"], start)
+        if self.viewpoint == "third_person_view":
+            return left
+        wrist = self._decode(episode.video_paths["wrist"], start)
+        if self.camera_layout == "left_wrist_horizontal":
+            if self.use_image_augmentation:
+                brightness = random.uniform(0.85, 1.15)
+                contrast = random.uniform(0.85, 1.15)
+                combined = torch.cat((left, wrist), dim=0)
+                mean = combined.mean(dim=(-2, -1), keepdim=True)
+                combined = ((combined - mean) * contrast + mean).mul(brightness).clamp(0, 1)
+                left, wrist = combined.chunk(2, dim=0)
+            # Preserve both native 256x256 views: agent-left | wrist.
+            return torch.cat((left, wrist), dim=-1)
+        right = self._decode(episode.video_paths["right"], start)
+        if self.use_image_augmentation:
+            brightness = random.uniform(0.85, 1.15)
+            contrast = random.uniform(0.85, 1.15)
+            combined = torch.cat((left, right, wrist), dim=0)
+            mean = combined.mean(dim=(-2, -1), keepdim=True)
+            combined = ((combined - mean) * contrast + mean).mul(brightness).clamp(0, 1)
+            n = left.shape[0]
+            left, right, wrist = combined[:n], combined[n : 2 * n], combined[2 * n :]
+        right = F.interpolate(right, size=(128, 128), mode="bilinear", align_corners=False)
+        wrist = F.interpolate(wrist, size=(128, 128), mode="bilinear", align_corners=False)
+        # Left agent view on top, right agent + wrist views in the bottom row.
+        return torch.cat((left, torch.cat((right, wrist), dim=-1)), dim=-2)
+
+    def get_action_normalizer(self, _sample: dict[str, Any] | None = None) -> ActionNormalizer | None:
+        return self._normalizer
+
+    def __getitem__(self, index: int) -> dict[str, Any]:
+        episode_index, frame_offset = self._resolve_index(index)
+        episode = self.episodes[episode_index]
+        low_dim = self._load_low_dim(episode)
+        action = low_dim["action"][frame_offset : frame_offset + self.chunk_length, self.action_indices]
+        if self.use_state:
+            # State and action have different physical layouts (16-D vs 12-D),
+            # so expose state separately instead of pretending it is an action.
+            initial_state = low_dim["state"][frame_offset, self.state_indices]
+        else:
+            initial_state = None
+        video = self._compose_video(episode, frame_offset)
+        result: dict[str, Any] = {
+            "ai_caption": episode.command,
+            "video": (video * 255.0).clamp(0, 255).to(torch.uint8).permute(1, 0, 2, 3),
+            "action": action,
+            "conditioning_fps": torch.tensor(int(self.fps), dtype=torch.long),
+            "mode": self.mode,
+            "domain_id": torch.tensor(self.domain_id, dtype=torch.long),
+            "viewpoint": self.viewpoint,
+            "additional_view_description": (
+                "The left panel is the left agent view and the right panel is the wrist view."
+                if self.camera_layout == "left_wrist_horizontal"
+                else "The top panel is the left agent view. The bottom-left panel is the right agent view, "
+                "and the bottom-right panel is the wrist view."
+            ),
+            "task_cluster": episode.task_name,
+            "task_category": episode.category,
+        }
+        if initial_state is not None:
+            result["proprio"] = initial_state
+        if self.emit_behavior_metadata:
+            # Keep ordinal source_index for diagnostics; resolve the canonical
+            # Stage-1 key from task + episode_id in the behavior store.
+            result["behavior_source_index"] = torch.tensor(episode_index, dtype=torch.long)
+            result["behavior_episode_id"] = torch.tensor(episode.episode_id, dtype=torch.long)
+            result["behavior_task_cluster"] = episode.task_name
+            result["behavior_frame_offset"] = torch.tensor(frame_offset, dtype=torch.long)
+        return result
+
+
+_COMPOSITE_UNSEEN = {
+    "ArrangeBreadBasket", "ArrangeTea", "BreadSelection", "CategorizeCondiments",
+    "CuttingToolSelection", "GarnishPancake", "GatherTableware", "HeatKebabSandwich",
+    "MakeIceLemonade", "PanTransfer", "PortionHotDogs", "RecycleBottlesByType",
+    "SeparateFreezerRack", "WaffleReheat", "WashFruitColander", "WeighIngredients",
+}

releases/bridge-native-v3/code/cosmos_framework/data/generator/action/datasets/widowx_bridge_v3_dataset.py

@@ -0,0 +1,163 @@
+"""Small, self-contained WidowX/Bridge LeRobot-v3 dataset adapter.
+
+The smoke export is deliberately kept independent from the legacy LeRobot
+parquet layout used by the upstream Cosmos loaders.  This adapter is also
+useful for quick Zeva policy experiments and returns normalized 7-D Bridge
+actions plus the 9-D proprioception contract used by the WidowX server.
+"""
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+from typing import Dict
+from collections import OrderedDict
+
+import av
+import numpy as np
+import pyarrow.parquet as pq
+import torch
+from torch.utils.data import Dataset
+
+Q01 = np.asarray([-0.028752606511116028, -0.041702210046350954, -0.026096698231995105,
+                  -0.08052822157740593, -0.09249736212193965, -0.20738411962985992], np.float32)
+Q99 = np.asarray([0.02830690816044803, 0.04089860741049051, 0.04018005654215808,
+                  0.08173405691981314, 0.07760896608233431, 0.20384809583425495], np.float32)
+
+
+def normalize_actions(a: np.ndarray) -> np.ndarray:
+    out = np.asarray(a, np.float32).copy()
+    out[..., :6] = np.clip(2 * ((out[..., :6] - Q01) / (Q99 - Q01)) - 1, -1, 1)
+    out[..., 6] = np.where(out[..., 6] > 0.5, 1.0, -1.0)
+    return out
+
+
+class WidowXBridgeV3Dataset(Dataset):
+    def __init__(self, root: str, chunk_length: int = 32, max_episodes: int | None = None,
+                 image_size: int = 256, feature_cache: str | None = None,
+                 window_mode: str | None = None, validate_cache: bool = True):
+        self.root = Path(root)
+        self.chunk_length = int(chunk_length)
+        self.window_mode = window_mode or os.environ.get('WIDOWX_WINDOW_MODE', 'full')
+        if self.window_mode not in ('full', 'legacy'):
+            raise ValueError('window_mode must be full or legacy')
+        self.image_size = int(image_size)
+        self.info = json.loads((self.root / 'meta/info.json').read_text())
+        self.video_key = 'observation.images.image_0'
+        if self.video_key not in self.info['features']:
+            raise ValueError('Bridge front camera image_0 is missing')
+        self.feature_cache = Path(feature_cache) if feature_cache else None
+        episodes = [json.loads(x) for x in (self.root / "meta/episodes.jsonl").read_text().splitlines()]
+        tasks = [json.loads(x) for x in (self.root / "meta/tasks.jsonl").read_text().splitlines()]
+        self.tasks = {int(x['task_index']): x['task'] for x in tasks}
+        if max_episodes is not None:
+            episodes = episodes[:int(max_episodes)]
+        self.episodes = episodes
+        self.windows = [(e["episode_index"], s) for e in episodes
+                        for s in range(0, int(e['length']) if self.window_mode == 'full'
+                                       else max(0, int(e['length']) - self.chunk_length), 4)]
+        self.feature_fingerprint = None
+        if self.feature_cache is not None:
+            if self.window_mode == 'full':
+                from scripts.widowx_contract import read_manifest
+                self.feature_fingerprint = read_manifest(self.root, self.feature_cache)
+            # Fail during dataset construction instead of much later in a
+            # worker, where a missing frozen CTE feature looks like a random
+            # dataloader crash after training has already started.
+            required_episodes = sorted({episode for episode, _ in self.windows})
+            missing = [episode for episode in required_episodes
+                       if not (self.feature_cache / f"episode_{episode:06d}.pt").is_file()]
+            if missing and validate_cache:
+                preview = ", ".join(str(x) for x in missing[:8])
+                raise FileNotFoundError(
+                    f"WIDOWX_FEATURE_CACHE is missing {len(missing)} episode feature files "
+                    f"(first: {preview}); run scripts/prepare_widowx_features.py before training"
+                )
+        self._frames = OrderedDict()
+        self._rows: Dict[int, Dict[str, np.ndarray]] = {}
+        self._feature_episode = None
+        self._features = None
+
+    def __len__(self):
+        return len(self.windows)
+
+    def _load(self, episode: int):
+        if episode in self._rows:
+            return
+        chunk = episode // int(self.info.get('chunks_size', 1000))
+        fmt = dict(episode_chunk=chunk, episode_index=episode, video_key=self.video_key)
+        dp = self.root / self.info['data_path'].format(**fmt)
+        tab = pq.read_table(dp)
+        rows = {k: np.asarray(tab[k].to_pylist())
+                for k in ("observation.state", "action", "task_index")}
+        vp = self.root / self.info['video_path'].format(**fmt)
+        with av.open(str(vp)) as container:
+            frames = [f.to_ndarray(format="rgb24") for f in container.decode(video=0)]
+        if len(frames) != len(rows['action']):
+            raise ValueError(f'Episode {episode}: video/action frame count mismatch')
+        if len(self._frames) >= 2:
+            old, _ = self._frames.popitem(last=False)
+            del self._rows[old]
+        self._rows[episode] = rows
+        self._frames[episode] = np.stack(frames)
+
+    def __getitem__(self, index: int):
+        episode, start = self.windows[index]
+        self._load(episode)
+        rows, frames = self._rows[episode], self._frames[episode]
+        end = start + self.chunk_length
+        # A terminal recorded action remains a valid policy label even where
+        # the following image was not exported. Mask those modalities separately.
+        n_action = min(self.chunk_length, len(rows['action']) - start)
+        n_frames = min(self.chunk_length + 1, len(frames) - start)
+        image_indices = np.minimum(np.arange(start, end + 1), len(frames) - 1)
+        imgs = torch.from_numpy(frames[image_indices]).permute(3, 0, 1, 2).contiguous()
+        action = np.zeros((self.chunk_length, 7), np.float32)
+        action[:n_action] = normalize_actions(rows["action"][start:start+n_action])
+        action[n_action:, 6] = action[n_action-1, 6]
+        state = rows["observation.state"][start]
+        proprio = np.concatenate([state, state[-1:]], axis=0).astype(np.float32)
+        task_index = int(rows['task_index'][start])
+        sample = {"video": imgs, "action": torch.from_numpy(action),
+                "proprio": torch.from_numpy(proprio),
+                "ai_caption": self.tasks[task_index], 'mode': 'wam',
+                'viewpoint': 'ego_view', 'domain_id': torch.tensor(7),
+                'conditioning_fps': torch.tensor(int(self.info['fps']))}
+        sample['widowx_action_valid'] = torch.arange(self.chunk_length) < n_action
+        # Wan's causal latents represent raw boundaries 0,4,8,...; require
+        # an entirely real four-frame group for each supervised future latent.
+        sample['widowx_vision_valid'] = torch.arange(0, self.chunk_length + 1, 4) < n_frames
+        sample['widowx_episode_start'] = torch.tensor([episode, start])
+        if self.feature_cache:
+            if self._feature_episode != episode:
+                features = torch.load(self.feature_cache / f'episode_{episode:06d}.pt', weights_only=True)
+                if self.feature_fingerprint is not None:
+                    if features.get('fingerprint') != self.feature_fingerprint:
+                        raise ValueError(f'Episode {episode}: feature fingerprint mismatch')
+                    if features['starts'].tolist() != list(range(0, len(frames), 4)):
+                        raise ValueError(f'Episode {episode}: incomplete feature boundaries')
+                self._features, self._feature_episode = features, episode
+            features = self._features
+            for key in ('behavior_global', 'behavior_phase', 'behavior_effect', 'behavior_effect_valid'):
+                sample[key] = features[key][start // 4]
+        return sample
+
+    def get_shuffle_blocks(self):
+        blocks = []
+        for i, (episode, _) in enumerate(self.windows):
+            if i == 0 or self.windows[i-1][0] != episode:
+                blocks.append([i, 1])
+            else:
+                blocks[-1][1] += 1
+        return blocks
+
+
+def get_action_widowx_sft_dataset(root: str, tokenizer_config=None, max_action_dim=64,
+                                resolution='256', **kwargs):
+    from cosmos_framework.data.generator.action.datasets.action_sft_dataset import ActionSFTDataset
+    from scripts.widowx_data_resume import WidowXIterableDataset
+    from cosmos_framework.data.generator.action.transforms import ActionTransformPipeline
+    raw = WidowXBridgeV3Dataset(root=root, **kwargs)
+    transform = ActionTransformPipeline(tokenizer_config=tokenizer_config, max_action_dim=max_action_dim,
+                                        format_prompt_as_json=True)
+    return WidowXIterableDataset(ActionSFTDataset(raw, transform, resolution))

releases/bridge-native-v3/code/cosmos_framework/data/generator/action/domain_utils.py

@@ -0,0 +1,83 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""Domain ID helpers for cross-embodiment action datasets."""
+
+EMBODIMENT_TO_DOMAIN_ID: dict[str, int] = {
+    "no_action": 0,
+    "av": 1,
+    "camera_pose": 2,
+    "hand_pose": 3,
+    "pusht": 4,
+    "libero": 5,
+    "umi": 6,
+    "bridge_orig_lerobot": 7,
+    "droid_lerobot": 8,
+    "robomind-franka": 8,
+    "embodiment_b": 9,
+    "robomind-franka-dual": 12,
+    "robomind-ur": 13,
+    "agibotworld": 15,
+    "embodiment_c_gripper": 15,
+    "embodiment_c_gripper_ext": 15,
+    "xdof_yam": 16,
+    "molmoact2_yam": 16,
+    "abc_yam": 16,
+    "fractal": 20,
+    "drawanything": 21,
+    # Dedicated Panda-Omron embedding used when the released policy was trained.
+    # Do not alias this to DROID or to an unused numeric slot.
+    "robocasa-panda-omron": 22,
+}
+
+
+EMBODIMENT_TO_RAW_ACTION_DIM: dict[str, int] = {
+    "av": 9,
+    "camera_pose": 9,
+    "pusht": 2,
+    "umi": 10,
+    "bridge_orig_lerobot": 10,
+    "droid_lerobot": 10,
+    "robomind-franka": 10,
+    "robomind-franka-dual": 20,
+    "robomind-ur": 10,
+    "embodiment_b": 30,
+    "agibotworld": 29,
+    "embodiment_c_gripper": 29,
+    "embodiment_c_gripper_ext": 29,
+    "xdof_yam": 20,
+    "molmoact2_yam": 20,
+    "abc_yam": 20,
+    "fractal": 10,
+    "drawanything": 3,
+    # Full source command before the Atomic-5 arm-only [5:12] selection.
+    "robocasa-panda-omron": 12,
+}
+
+
+def get_domain_id(embodiment_type: str) -> int:
+    """Get the model domain embedding ID for an embodiment."""
+    key = embodiment_type.lower().strip()
+    if key not in EMBODIMENT_TO_DOMAIN_ID:
+        raise KeyError(
+            f"Unknown embodiment type: {embodiment_type!r}. "
+            f"Available embodiments: {sorted(EMBODIMENT_TO_DOMAIN_ID.keys())}"
+        )
+    return EMBODIMENT_TO_DOMAIN_ID[key]
+
+
+def get_action_dim(embodiment_type: str) -> int:
+    """Get the raw action width before dataset-specific channel selection."""
+    key = embodiment_type.lower().strip()
+    if key not in EMBODIMENT_TO_RAW_ACTION_DIM:
+        raise KeyError(
+            f"Unknown embodiment type: {embodiment_type!r}. "
+            f"Available embodiments: {sorted(EMBODIMENT_TO_RAW_ACTION_DIM.keys())}"
+        )
+    return EMBODIMENT_TO_RAW_ACTION_DIM[key]
+
+
+def is_valid_domain_name(embodiment_type: str) -> bool:
+    """Check whether an embodiment has a canonical raw action width."""
+    key = embodiment_type.lower().strip()
+    return key in EMBODIMENT_TO_RAW_ACTION_DIM

releases/bridge-native-v3/code/cosmos_framework/data/generator/action/normalizer_stats/robocasa365_target_action_stats.json

@@ -0,0 +1,11 @@
+{
+  "action": {
+    "min": [-1.0, -1.0, -1.0, 0.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0],
+    "max": [1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
+    "mean": [0.02590581175350872, 0.00202585452064382, -0.000603911072971857, 0.0, -0.7797204007304583, 0.007985740558506635, -0.009527682245632989, -0.02404516475177401, 0.0024829546270168945, -0.009724957579023025, -0.00014074740094141753, -0.17212704805076373],
+    "std": [0.19966604450417424, 0.18601724713237694, 0.12950011360377808, 0.0, 0.6261278595805986, 0.4100137479669798, 0.3704388299067406, 0.3502567436993047, 0.10230268334142267, 0.11492962521591003, 0.11579242435141789, 0.9850747576645584],
+    "q01": [-0.3200703506810163, -0.34884918400617965, -0.21919116581311202, 3.9995784462121255e-16, -0.9968481276386474, -0.9422251311439209, -0.8442310960485607, -0.7521983922010141, -0.2722379479403091, -0.36124760554870705, -0.31036559990487256, -0.9993595477431751],
+    "q99": [0.39805452510884953, 0.33510666883940965, 0.21445504450811365, 3.9595826617500056e-14, 0.03382406019703607, 0.8090603250091402, 0.8343534480833622, 0.8754301710780941, 0.2783829802274566, 0.28977084389036595, 0.32830698508606554, 0.9179061147815708],
+    "count": [14957899]
+  }
+}

releases/bridge-native-v3/code/cosmos_framework/data/generator/action/transforms.py

@@ -0,0 +1,710 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""Dataset transform wrappers for the Action project.
+
+This module provides the ``ActionTransformPipeline`` and spatial padding utilities.
+
+The reflection padding snaps each sample to the closest predefined resolution from
+``VIDEO_RES_SIZE_INFO`` (matching VFM's approach), guaranteeing a bounded set of
+output shapes that are all multiples of 16.
+
+See :func:`~.unified_dataset.wrap_dataset` for the convenience factory that
+combines datasets with transforms, and :class:`~.unified_dataset.MapToIterableAdapter`
+for the map-to-iterable wrapper.
+"""
+
+from __future__ import annotations
+
+import torch
+import torchvision.transforms.functional as transforms_F
+
+from cosmos_framework.utils import log
+from cosmos_framework.data.generator.action.action_processing import (
+    ActionNormalizer,
+    ActionProcessor,
+)
+from cosmos_framework.data.generator.action.json_formatter import ActionPromptJsonFormatter
+from cosmos_framework.data.generator.action.viewpoint_utils import ViewpointTextInfo
+from cosmos_framework.data.generator.augmentors.duration_fps_text_timestamps import DurationFPSTextTimeStamps
+from cosmos_framework.data.generator.augmentors.idle_frames_text_info import IdleFramesTextInfo
+from cosmos_framework.data.generator.augmentors.resolution_text_info import ResolutionTextInfo
+from cosmos_framework.data.generator.augmentors.text_tokenizer import TextTokenizerTransform
+from cosmos_framework.data.generator.utils import VIDEO_RES_SIZE_INFO
+from cosmos_framework.data.generator.sequence_packing import SequencePlan
+from cosmos_framework.utils.generator.data_utils import get_vision_data_resolution
+
+
+def _should_append_idle_frame_info(mode: object) -> bool:
+    """Return whether idle-frame prompt metadata should be surfaced."""
+    return mode != "inverse_dynamics"
+
+
+def find_closest_target_size(h: int, w: int, resolution: str | int) -> tuple[int, int]:
+    """Find the closest predefined target size for a given input resolution.
+
+    Looks up ``VIDEO_RES_SIZE_INFO[resolution]`` and selects the aspect ratio
+    whose ``H/W`` ratio is closest to the input ``h/w``.
+
+    Args:
+        h: Input height in pixels.
+        w: Input width in pixels.
+        resolution: Resolution tier key (e.g. ``"256"``, ``"480"``, ``"720"``).
+
+    Returns:
+        ``(target_w, target_h)`` from the predefined table.
+
+    Raises:
+        ValueError: If *resolution* is not a key in ``VIDEO_RES_SIZE_INFO``.
+    """
+    if isinstance(resolution, int):
+        resolution = str(resolution)
+    if resolution not in VIDEO_RES_SIZE_INFO:
+        raise ValueError(
+            f"Resolution '{resolution}' not found in VIDEO_RES_SIZE_INFO. Available: {list(VIDEO_RES_SIZE_INFO.keys())}"
+        )
+
+    candidates = VIDEO_RES_SIZE_INFO[resolution]
+    input_ratio = h / w
+
+    best_key: str | None = None
+    best_diff = float("inf")
+    for aspect_key, (cand_w, cand_h) in candidates.items():
+        cand_ratio = cand_h / cand_w
+        diff = abs(input_ratio - cand_ratio)
+        if diff < best_diff:
+            best_diff = diff
+            best_key = aspect_key
+
+    assert best_key is not None
+    target_w, target_h = candidates[best_key]
+    return target_w, target_h
+
+
+def reflection_pad_to_target(
+    data_dict: dict,
+    keys: list[str],
+    keep_aspect_ratio: bool,
+    target_w: int,
+    target_h: int,
+) -> dict:
+    """Resize (aspect-preserving) and reflection-pad tensors to exact target size.
+
+    For each key in *keys*, the tensor is:
+
+    1. Resized so its spatial dimensions fit within ``(target_h, target_w)``
+       while preserving the aspect ratio (matching VFM's
+       ``ResizeLargestSideAspectPreserving``).
+    2. Reflection-padded (or edge-padded when the padding exceeds the spatial
+       dimension) to reach exactly ``(target_h, target_w)`` (matching VFM's
+       ``ReflectionPadding``).
+
+    After processing, the following entries are added to *data_dict*:
+
+    - ``"image_size"``: ``torch.Tensor`` of shape ``(4,)`` containing
+      ``[target_h, target_w, orig_h_resized, orig_w_resized]`` where
+      ``target_h/w`` is the padded canvas size and ``orig_h/w_resized``
+      is the original spatial size after aspect-preserving resize (i.e.
+      the content region before padding).  After ``default_collate``
+      this becomes ``(B, 4)``;  the ``IterativeJointDataLoader`` then
+      splits it into per-sample ``(1, 4)`` tensors so the model can
+      index as ``data_batch["image_size"][i][0][0]``.
+
+    Args:
+        data_dict: The sample dictionary (mutated in-place).
+        keys: Data-dict keys whose tensors should be resized and padded.
+            Tensors must have shape ``(C, H, W)`` or ``(C, T, H, W)``.
+        keep_aspect_ratio: Whether to keep the aspect ratio of the input tensor.
+        target_w: Target width in pixels.
+        target_h: Target height in pixels.
+
+    Returns:
+        The mutated *data_dict*.
+    """
+    orig_h_resized: int = 0
+    orig_w_resized: int = 0
+
+    for key in keys:
+        if key not in data_dict:
+            continue
+        tensor = data_dict[key]
+        if not isinstance(tensor, torch.Tensor):
+            continue
+
+        # Extract spatial dims
+        if tensor.ndim == 3:
+            orig_h, orig_w = tensor.shape[-2:]
+        elif tensor.ndim == 4:
+            orig_h, orig_w = tensor.shape[-2:]
+        else:
+            raise ValueError(f"Unexpected tensor ndim={tensor.ndim} for key '{key}', expected 3 or 4")
+
+        # Step 1: aspect-preserving resize to fit within (target_h, target_w)
+        if keep_aspect_ratio:
+            # Prevent upscaling the video by setting the upper bound of scaling_ratio to 1.0.
+            scaling_ratio = min(target_w / orig_w, target_h / orig_h, 1.0)
+            orig_h_resized = int(scaling_ratio * orig_h + 0.5)
+            orig_w_resized = int(scaling_ratio * orig_w + 0.5)
+            assert orig_h_resized <= target_h and orig_w_resized <= target_w, (
+                f"Resize error: orig ({orig_h}, {orig_w}) target ({target_h}, {target_w}) "
+                f"computed ({orig_h_resized}, {orig_w_resized})"
+            )
+        else:
+            orig_h_resized = target_h
+            orig_w_resized = target_w
+
+        if orig_h_resized != orig_h or orig_w_resized != orig_w:
+            tensor = transforms_F.resize(
+                tensor,
+                size=[orig_h_resized, orig_w_resized],
+                interpolation=transforms_F.InterpolationMode.BICUBIC,
+                antialias=True,
+            )
+
+        # Step 2: padding to exact target size (bottom and right only)
+        if orig_w_resized != target_w or orig_h_resized != target_h:
+            padding_right = target_w - orig_w_resized
+            padding_bottom = target_h - orig_h_resized
+            padding = [0, 0, padding_right, padding_bottom]
+
+            if padding_right >= orig_w_resized or padding_bottom >= orig_h_resized:
+                tensor = transforms_F.pad(tensor, padding, padding_mode="edge")
+            else:
+                tensor = transforms_F.pad(tensor, padding, padding_mode="reflect")
+
+        data_dict[key] = tensor
+
+    # image_size: shape (4,) — [target_h, target_w, orig_h_resized, orig_w_resized].
+    # Matches VFM's item_dataset convention.  default_collate stacks to (B, 4);
+    # IterativeJointDataLoader._get_next_sample slices to (1, 4) per sample so
+    # the model can index [i][0][0].
+    data_dict["image_size"] = torch.tensor(
+        [target_h, target_w, orig_h_resized, orig_w_resized], dtype=torch.float
+    )  # [4]
+
+    return data_dict
+
+
+def remove_reflection_padding(
+    tensor: torch.Tensor,
+    image_size: torch.Tensor | list[torch.Tensor] | None,
+) -> torch.Tensor:
+    """Remove reflection padding added by :func:`reflection_pad_to_target`.
+
+    Content is at top-left; crops to ``(orig_h_resized, orig_w_resized)``.
+
+    Args:
+        tensor: Tensor whose last two dimensions are the padded spatial dims.
+            Supports any leading dimensions, e.g. ``(C, T, H, W)`` or
+            ``(C, H, W)``.
+        image_size: Spatial metadata using the convention produced by
+            :func:`reflection_pad_to_target`. Accepted forms are ``None`` (no
+            crop), a tensor with shape ``(4,)`` or ``(1, 4)``, or a non-empty
+            list whose first element has one of those tensor shapes. The four
+            values are ``[target_h, target_w, orig_h_resized,
+            orig_w_resized]``, where ``orig_h/w_resized`` is the original
+            spatial size after aspect-preserving resize (i.e. the content
+            region before padding). This matches the convention stored by
+            :func:`reflection_pad_to_target` and VFM's ``ReflectionPadding``.
+
+    Returns:
+        Cropped tensor of shape ``(..., orig_h_resized, orig_w_resized)``.
+    """
+    if image_size is None:
+        return tensor
+    if isinstance(image_size, list):
+        if not image_size:
+            raise ValueError("Expected at least one image_size entry")
+        image_size = image_size[0]  # [1,4] or [4]
+    if image_size.ndim == 2 and image_size.shape[0] == 1:
+        image_size = image_size[0]  # [4]
+    if image_size.ndim != 1:
+        raise ValueError(f"Expected image_size shape [4] or [1,4], got {tuple(image_size.shape)}")
+
+    target_h = int(image_size[0].item())
+    target_w = int(image_size[1].item())
+    orig_h_resized = int(image_size[2].item())
+    orig_w_resized = int(image_size[3].item())
+
+    if orig_h_resized == target_h and orig_w_resized == target_w:
+        return tensor
+
+    return tensor[..., :orig_h_resized, :orig_w_resized].contiguous()
+
+
+def build_sequence_plan_from_mode(
+    mode: str,
+    video_length: int,
+    action_length: int,
+    has_text: bool = True,
+    video_temporal_downsample: int = 4,
+    num_history_actions: int = 0,
+) -> SequencePlan:
+    """Build a SequencePlan based on the training mode.
+
+    This function determines whether action should be included and computes the
+    appropriate condition frame indexes for vision and action based on the mode.
+
+    Args:
+        mode: Training mode. One of:
+            - "forward_dynamics": Predict video given first frame and all actions
+            - "inverse_dynamics": Predict actions given all video frames
+            - "wam": Jointly denoise video and actions given the first frame
+        video_length: Number of video frames (including the conditioning frame).
+        action_length: Number of action steps (typically video_length - 1).
+        has_text: Whether text conditioning is available. Defaults to True.
+        video_temporal_downsample: Temporal downsampling factor of the video
+            tokenizer. Used to compute condition frame indexes for inverse
+            dynamics mode. Defaults to 4.
+
+    Returns:
+        SequencePlan instance with appropriate settings.
+        Use ``sequence_plan.has_action`` to check if action should be included.
+
+    Raises:
+        ValueError: If mode is not one of the supported modes.
+
+    Example:
+        >>> sequence_plan = build_sequence_plan_from_mode(
+        ...     mode="wam",
+        ...     video_length=5,
+        ...     action_length=4,
+        ... )
+        >>> sequence_plan.has_action
+        True
+        >>> sequence_plan.as_dict()
+        {'has_text': True, 'has_vision': True, 'has_action': True,
+         'condition_frame_indexes_vision': [0], 'condition_frame_indexes_action': []}
+    """
+    valid_modes = ["forward_dynamics", "inverse_dynamics", "wam"]
+    if mode not in valid_modes:
+        raise ValueError(f"Invalid mode: {mode!r}. Must be one of {valid_modes}")
+
+    # Determine condition frame indexes based on mode
+    # forward_dynamics/wam: first frame is clean (conditioning)
+    # inverse_dynamics: all frames are provided as context
+    if mode in ["forward_dynamics", "wam"]:
+        condition_frame_indexes_vision = [0]
+    elif mode == "inverse_dynamics":
+        # All frames are observed for inverse dynamics
+        condition_frame_indexes_vision = list(range(0, (video_length - 1) // video_temporal_downsample + 1))
+    else:
+        condition_frame_indexes_vision = []
+
+    # For action conditioning indexes:
+    # forward_dynamics: all action steps are clean (conditioning)
+    # inverse_dynamics/wam: action is supervised (predicted)
+    # History frames (prepended) are always conditioning.
+    #
+    # `base_action_length` is the number of *predicted* action steps, i.e.
+    # `action_length` minus the prepended history block. There are three
+    # supported alignments between the action stream and the video stream
+    # (`video_length` counts the leading conditioning frame at index 0):
+    #
+    #   Case A (dense video, action = video_length - 1)
+    #       Standard forward-dynamics / WAM layout: one action per gap
+    #       between consecutive video frames, no explicit initial-state
+    #       action. Conditioning actions are only the prepended history.
+    #
+    #   Case B (dense video, action = video_length)
+    #       `use_state=True` layout: one extra action is prepended to
+    #       encode the initial state, so history + one initial-state
+    #       action are conditioning.
+    #
+    #   Case C (subsampled video, general action_length > video_length)
+    #       Video is subsampled by an integer factor while the action
+    #       stream stays dense. `has_initial_state` mirrors Case B: it is
+    #       True iff `(base_action_length - 1)` divides evenly by
+    #       `(video_length - 1)`, i.e. the extra initial-state action is
+    #       present.
+    base_action_length = action_length - num_history_actions
+    has_initial_state = video_length > 1 and (base_action_length - 1) % (video_length - 1) == 0
+    if mode == "forward_dynamics":
+        condition_frame_indexes_action = list(range(action_length))
+    elif base_action_length == video_length - 1:
+        # Case A: only the history block is conditioning.
+        condition_frame_indexes_action = list(range(num_history_actions))
+    elif base_action_length == video_length:
+        # Case B: history block + one initial-state action.
+        condition_frame_indexes_action = list(range(num_history_actions + 1))
+    else:
+        # Case C: same rule as A/B based on whether an initial-state action exists.
+        condition_frame_indexes_action = list(
+            range(num_history_actions + 1 if has_initial_state else num_history_actions)
+        )
+
+    # `action_start_frame_offset` shifts the action stream so that the first
+    # non-conditioning action aligns with video frame index 1 (the first
+    # frame the model must predict). The `-num_history_actions` term
+    # reserves room for the history block; the additional `+1` (i.e.
+    # `1 - num_history_actions`) is applied only when there is no
+    # initial-state action, so the offset skips past the first action
+    # slot that is instead consumed by the initial state.
+    if base_action_length == video_length - 1:
+        action_start_frame_offset = 1 - num_history_actions  # Case A
+    elif base_action_length == video_length:
+        action_start_frame_offset = -num_history_actions  # Case B
+    else:
+        # Case C: same rule as A/B based on whether an initial-state action exists.
+        action_start_frame_offset = -num_history_actions if has_initial_state else 1 - num_history_actions
+
+    return SequencePlan(
+        has_text=has_text,
+        has_vision=True,
+        has_action=True,
+        condition_frame_indexes_vision=condition_frame_indexes_vision,
+        condition_frame_indexes_action=condition_frame_indexes_action,
+        action_start_frame_offset=action_start_frame_offset,
+    )
+
+
+class VideoResize:
+    """Resize and reflection-pad video-aligned tensors for a single sample.
+
+    Resolution is supplied at call time. When ``resolution`` is ``None``, the
+    tier is auto-detected from the sample's ``"video"`` spatial dimensions.
+
+    Args:
+        pad_keys: Data-dict keys whose values should be resized and padded.
+            Pass an empty list to disable padding entirely. Defaults to
+            ``["video"]``.
+        keep_aspect_ratio: Whether to resize aspect-preservingly to the closest
+            predefined target size before padding. Defaults to ``True``.
+        log_prefix: Prefix used in debug logging.
+    """
+
+    def __init__(
+        self,
+        pad_keys: list[str] | None = None,
+        keep_aspect_ratio: bool = True,
+        log_prefix: str = "VideoResize",
+    ) -> None:
+        self.pad_keys = pad_keys if pad_keys is not None else ["video"]
+        self.keep_aspect_ratio = keep_aspect_ratio
+        self.log_prefix = log_prefix
+
+    def __call__(self, data_dict: dict, resolution: str | int | None) -> dict:
+        """Resize and pad a sample in-place.
+
+        Args:
+            data_dict: Sample dictionary containing a ``"video"`` entry.
+            resolution: Resolution tier key (e.g. ``"256"``, ``"480"``,
+                ``"720"``). When ``None``, auto-detected from video dimensions.
+
+        Returns:
+            The same dictionary, mutated in-place with padded tensors and an
+            ``"image_size"`` entry.
+        """
+        video = data_dict.get("video")
+        assert isinstance(video, torch.Tensor), "video is required for reflection padding"
+        h, w = video.shape[-2:]
+
+        if resolution is None:
+            resolution = get_vision_data_resolution((h, w))
+
+        if self.keep_aspect_ratio:
+            target_w, target_h = find_closest_target_size(h, w, resolution)
+        else:
+            target_w = int(resolution)
+            target_h = int(resolution)
+        reflection_pad_to_target(data_dict, self.pad_keys, self.keep_aspect_ratio, target_w, target_h)
+
+        return data_dict
+
+    def _log_shapes(self, data_dict: dict, when: str) -> None:
+        """Log tensor shapes for the configured pad keys."""
+        for key in self.pad_keys:
+            val = data_dict.get(key)
+            if isinstance(val, torch.Tensor):
+                log.debug(f"{self.log_prefix}: {when} padding '{key}' shape = {tuple(val.shape)}")
+
+
+class ActionTransformPipeline:
+    """A composable transform pipeline that chains ``VideoResize``, text
+    tokenization, and automatic sequence plan construction.
+
+    Reflection padding snaps each sample to the closest predefined aspect
+    ratio from ``VIDEO_RES_SIZE_INFO[resolution]``, resizes
+    (aspect-preserving) to fit within the target, then reflection-pads to
+    the exact target size.  This guarantees a bounded set of output shapes
+    (5 per resolution tier), all multiples of 16.  Resolution is supplied
+    at call time via the required ``resolution`` argument to ``__call__``;
+    when ``resolution`` is ``None``, the tier is auto-detected from the
+    video's spatial dimensions via ``get_vision_data_resolution``.
+
+    Text tokenization is enabled when ``tokenizer_config`` is provided.
+
+    When the data dictionary contains a ``"mode"`` key, the pipeline automatically
+    builds a ``SequencePlan`` via :func:`build_sequence_plan_from_mode` and attaches
+    it as ``data_dict["sequence_plan"]``.
+
+    Args:
+        pad_keys: Data-dict keys whose values should be resized and padded. Pass
+            an empty list to disable padding entirely. Defaults to ``["video"]``.
+        tokenizer_config: A lazy-instantiable config dict for the VLM tokenizer. When
+            ``None``, text tokenization is skipped. Defaults to ``None``.
+        cfg_dropout_rate: Probability of replacing the caption with an empty string for
+            classifier-free guidance. Only used when text tokenization is enabled.
+            Defaults to ``0.0``.
+        caption_key: The data-dict key that contains the input caption string.
+            Defaults to ``"ai_caption"``.
+        text_token_key: The data-dict key where tokenized text IDs will be stored.
+            Defaults to ``"text_token_ids"``.
+        video_temporal_downsample: Temporal downsampling factor of the video tokenizer.
+            Used when building a ``SequencePlan`` for ``"inverse_dynamics"`` mode.
+            Defaults to 4.
+        max_action_dim: Target action dimension to pad to.  The ``"action"`` tensor
+            in every sample is padded along its last dimension via
+            :func:`pad_action_to_max_dim`.  Defaults to 32.
+        action_channel_masking: When ``True`` (default), the original action
+            dimension is stored in ``"raw_action_dim"`` so that the model masks
+            loss/noise/velocity on zero-padded action channels.  When ``False``,
+            ``"raw_action_dim"`` is set to ``None`` and the model treats all
+            ``max_action_dim`` channels equally (original main-branch behavior).
+        append_viewpoint_info: Whether to append viewpoint type metadata to the
+            caption (via ``ViewpointTextInfo`` augmentor).  Requires that
+            samples contain a ``"viewpoint"`` key.  Defaults to ``True``.
+        append_duration_fps_timestamps: Whether to append duration and FPS metadata to the
+            caption (matching VFM's ``DurationFPSTextTimeStamps`` augmentor).
+            Defaults to ``True``.
+        append_resolution_info: Whether to append resolution metadata to the
+            caption (matching VFM's ``ResolutionTextInfo`` augmentor).
+            Defaults to ``True``.
+        append_idle_frames: Whether to append the idle-frame count out of the
+            total action frames to the caption (Pi0.7-style metadata, via
+            ``IdleFramesTextInfo`` augmentor).  The dataset is responsible for
+            populating ``data_dict["idle_frames"]``; samples without it are
+            silently skipped.  Idle-frame text is skipped only for
+            ``"inverse_dynamics"`` mode.  Defaults to ``False`` so existing
+            experiments are unaffected.
+        idle_frames_dropout: Per-field dropout rate for the idle-frame segment.
+            With this probability the augmentor leaves the caption unchanged
+            (matching Pi0.7's ~5% per-component dropout).  Independent of the
+            global ``cfg_dropout_rate``, which empties the whole caption.
+            Defaults to 0.05.
+        format_prompt_as_json: Whether to replace the plain text prompt with a
+            structured JSON-compatible dictionary before tokenization.  When
+            enabled, legacy string metadata appenders are skipped and the JSON
+            formatter owns viewpoint, action, resolution, duration, FPS, and
+            idle-frame fields.  Defaults to ``False``.
+    """
+
+    def __init__(
+        self,
+        pad_keys: list[str] | None = None,
+        keep_aspect_ratio: bool = True,
+        tokenizer_config: dict | None = None,
+        cfg_dropout_rate: float = 0.0,
+        caption_key: str = "ai_caption",
+        text_token_key: str = "text_token_ids",
+        video_temporal_downsample: int = 4,
+        max_action_dim: int = 32,
+        action_channel_masking: bool = True,
+        append_viewpoint_info: bool = True,
+        append_duration_fps_timestamps: bool = True,
+        append_resolution_info: bool = True,
+        append_idle_frames: bool = False,
+        idle_frames_dropout: float = 0.05,
+        format_prompt_as_json: bool = False,
+    ) -> None:
+        self.caption_key: str = caption_key
+        self.video_temporal_downsample: int = video_temporal_downsample
+        self.max_action_dim: int = max_action_dim
+        self.action_channel_masking: bool = action_channel_masking
+        self.action_processor: ActionProcessor = ActionProcessor(
+            max_action_dim=max_action_dim,
+            action_channel_masking=action_channel_masking,
+        )
+
+        # --- Spatial resize/padding stage (resolution supplied at call time) ---
+        self.video_resize: VideoResize = VideoResize(
+            pad_keys=pad_keys,
+            keep_aspect_ratio=keep_aspect_ratio,
+            log_prefix="ActionTransformPipeline",
+        )
+        self.pad_keys: list[str] = self.video_resize.pad_keys
+        self.keep_aspect_ratio: bool = self.video_resize.keep_aspect_ratio
+
+        self.prompt_json_formatter: ActionPromptJsonFormatter | None = None
+        if format_prompt_as_json:
+            self.prompt_json_formatter = ActionPromptJsonFormatter(caption_key=caption_key)
+
+        # --- Viewpoint text augmentor (runs after ai_caption, before duration/FPS) ---
+        self.viewpoint_augmentor: ViewpointTextInfo | None = None
+        if append_viewpoint_info and self.prompt_json_formatter is None:
+            self.viewpoint_augmentor = ViewpointTextInfo(
+                input_keys=[caption_key, "viewpoint"],
+                output_keys=[caption_key],
+                args={"caption_key": caption_key, "viewpoint_key": "viewpoint", "enabled": True},
+            )
+
+        # --- Duration/FPS text augmentor (runs before tokenization) ---
+        self.duration_fps_augmentor: DurationFPSTextTimeStamps | None = None
+        if append_duration_fps_timestamps and self.prompt_json_formatter is None:
+            self.duration_fps_augmentor = DurationFPSTextTimeStamps(
+                input_keys=[caption_key, "video", "conditioning_fps"],
+                output_keys=[caption_key],
+                args={"caption_key": caption_key, "video_key": "video", "fps_key": "conditioning_fps"},
+            )
+
+        # --- Resolution text augmentor (runs before tokenization) ---
+        self.resolution_info_augmentor: ResolutionTextInfo | None = None
+        if append_resolution_info and self.prompt_json_formatter is None:
+            self.resolution_info_augmentor = ResolutionTextInfo(
+                input_keys=[caption_key, "video", "image_size"],
+                output_keys=[caption_key],
+                args={"caption_key": caption_key, "video_key": "video", "enabled": True},
+            )
+
+        # --- IdleFrames text augmentor (Pi0.7-style episode metadata) ---
+        # Runs after resolution info, before tokenization. Per-field dropout is
+        # independent from the tokenizer's global cfg_dropout_rate.
+        self.idle_frames_augmentor: IdleFramesTextInfo | None = None
+        if append_idle_frames and self.prompt_json_formatter is None:
+            self.idle_frames_augmentor = IdleFramesTextInfo(
+                input_keys=[caption_key, "idle_frames", "action"],
+                output_keys=[caption_key],
+                args={
+                    "caption_key": caption_key,
+                    "idle_frames_key": "idle_frames",
+                    "action_key": "action",
+                    "dropout_rate": idle_frames_dropout,
+                    "enabled": True,
+                },
+            )
+
+        # --- Text tokenizer augmentor ---
+        self.text_tokenizer: TextTokenizerTransform | None = None
+        if tokenizer_config is not None:
+            self.text_tokenizer = TextTokenizerTransform(
+                input_keys=[caption_key],
+                output_keys=[text_token_key],
+                args={
+                    "tokenizer_config": tokenizer_config,
+                    "cfg_dropout_rate": cfg_dropout_rate,
+                },
+            )
+
+    def __call__(
+        self,
+        data_dict: dict,
+        resolution: str | None,
+        action_normalizer: ActionNormalizer | None = None,
+    ) -> dict:
+        """Apply the transform pipeline to a single data dictionary.
+
+        Resolution is required at call time and is the only source of truth
+        for this sample. When ``resolution`` is ``None``, the tier is
+        auto-detected from the video's spatial dimensions.
+
+        The pipeline runs in order:
+
+        1. Resize + reflection-pad spatial dimensions to the closest
+           predefined target from ``VIDEO_RES_SIZE_INFO[resolution]``.
+        2. Format the caption as a structured JSON prompt (if enabled).
+        3. Otherwise, append viewpoint type metadata to caption (if enabled).
+        4. Append duration/FPS metadata to caption (if enabled).
+        5. Append resolution metadata to caption (if enabled).
+        6. Append idle-frame metadata (Pi0.7-style) to caption unless the
+           sample is in inverse dynamics mode (if enabled).
+        7. Tokenize caption text (if enabled).
+        8. Build a ``SequencePlan`` from the ``"mode"`` key (if present).
+        9. Preserve the canonical unnormalized and unpadded action as
+           ``"action_raw"``, normalize real channels, pad ``"action"`` to
+           ``max_action_dim``, and attach ``"action_processing_record"``.
+
+        Args:
+            data_dict: A sample dictionary as returned by a Action dataset.
+            resolution: Resolution tier key (e.g. ``"256"``, ``"480"``, ``"720"``)
+                for this sample. When ``None``, auto-detected from video dimensions.
+            action_normalizer: Optional source-provided action normalizer. When
+                present, only unpadded real action channels are normalized
+                before model-space channel padding.
+
+        Returns:
+            The same dictionary, mutated in-place with padded tensors,
+            ``image_size``, tokenized text IDs, a ``"sequence_plan"`` entry,
+            and action processing metadata added.
+        """
+        mode = data_dict.get("mode")
+        assert mode is not None, "mode is required"
+
+        # 1. Resize + reflection-pad spatial dimensions to the closest predefined target from ``VIDEO_RES_SIZE_INFO[resolution]``.
+        data_dict = self.video_resize(data_dict, resolution)
+
+        # 2. Format the caption as structured JSON when requested; otherwise run the legacy string appenders.
+        if self.prompt_json_formatter is not None:
+            data_dict = self.prompt_json_formatter(data_dict)
+        else:
+            # The relabeling is only rendered by the JSON formatter.  Drop it
+            # here so that enabling ``annotation_mode`` without
+            # ``format_prompt_as_json`` degrades to the plain caption instead of
+            # leaking a dict into the collate function.
+            data_dict.pop("relabel_prompt", None)
+
+            # 3. Append viewpoint type metadata to caption (if enabled).
+            if self.viewpoint_augmentor is not None:
+                result = self.viewpoint_augmentor(data_dict)
+                if result is not None:
+                    data_dict = result
+
+            # 4. Append duration/FPS metadata to caption (if enabled).
+            if self.duration_fps_augmentor is not None:
+                result = self.duration_fps_augmentor(data_dict)
+                if result is not None:
+                    data_dict = result
+
+            # 5. Append resolution metadata to caption (if enabled).
+            if self.resolution_info_augmentor is not None:
+                result = self.resolution_info_augmentor(data_dict)
+                if result is not None:
+                    data_dict = result
+
+            # 6. Append idle-frame metadata to caption (if enabled for this mode).
+            if self.idle_frames_augmentor is not None and _should_append_idle_frame_info(mode):
+                result = self.idle_frames_augmentor(data_dict)
+                if result is not None:
+                    data_dict = result
+
+        # 7. Tokenize caption text (if enabled).
+        if self.text_tokenizer is not None:
+            data_dict = self.text_tokenizer(data_dict)
+
+        # 8. Build a ``SequencePlan`` from the ``"mode"`` key (if present).
+        video = data_dict.get("video")
+        action = data_dict.get("action")
+        assert video is not None, "video is required"
+        video_length = video.shape[1]  # [C,T,H,W] -> T
+        action_length = action.shape[0] if isinstance(action, torch.Tensor) else max(video_length - 1, 0)
+
+        # Prepend history action frames (ground-truth conditioning) if present.
+        history_action = data_dict.pop("history_action", None)
+        residual_gate = data_dict.get("residual_gate")
+        num_history_actions = 0
+        if history_action is not None and isinstance(action, torch.Tensor):
+            num_history_actions = history_action.shape[0]
+            action = torch.cat([history_action, action], dim=0)
+            action_length += num_history_actions
+            if isinstance(residual_gate, torch.Tensor):
+                residual_gate = torch.cat(
+                    [torch.zeros((num_history_actions, 1), dtype=residual_gate.dtype), residual_gate], dim=0
+                )
+                data_dict["residual_gate"] = residual_gate
+
+        sequence_plan = build_sequence_plan_from_mode(
+            mode=mode,
+            video_length=video_length,
+            action_length=action_length,
+            video_temporal_downsample=self.video_temporal_downsample,
+            num_history_actions=num_history_actions,
+        )
+        data_dict["sequence_plan"] = sequence_plan
+
+        assert isinstance(action, torch.Tensor), "action tensor is required for action modes"
+        data_dict = self.action_processor.preprocess_action(
+            data_dict,
+            action,
+            action_normalizer=action_normalizer,
+        )
+
+        return data_dict

releases/bridge-native-v3/code/cosmos_framework/data/generator/joint_dataloader.py

@@ -0,0 +1,1271 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+import math
+import queue
+import threading
+from collections import deque
+from collections.abc import Iterator, Mapping
+from dataclasses import dataclass
+from typing import Any, ClassVar, Dict, Union
+
+import numpy as np
+import torch
+import webdataset
+from torch.utils.data.dataloader import default_collate
+
+from cosmos_framework.utils.lazy_config import instantiate
+from cosmos_framework.utils import log
+from cosmos_framework.model.generator.tokenizers.uniae.frame_math import (
+    get_uniae_chunk_frames,
+    get_uniae_latent_num_frames,
+    normalize_uniae_chunk_frames,
+)
+from cosmos_framework.utils.generator.data_utils import read_positive_int_metadata
+
+_TIMING_KEYS = {"_sample_time", "_aug_time", "_pre_aug_time", "_aug_step_times"}
+_BATCH_TIMING_KEYS = {
+    "_worker_batch_time",
+    "_worker_aug_time",
+    "_worker_io_time",
+    "_worker_aug_step_times",
+    "_worker_id",
+}
+
+
+def custom_collate_fn(batch):
+    """
+    Collate function that works like default_collate for all keys other than "text_token_ids", "images", and "video".
+    For "text_token_ids", "images", and "video" it simply returns them in a list, instead of stacking them as a tensor.
+    """
+    list_collate_keys = {
+        "text_token_ids",
+        "images",
+        "video",
+        "action",
+        "action_raw",
+        "domain_id",
+        "sequence_plan",
+        "sound",
+        "raw_action_dim",
+        "image_size",
+        "action_processing_record",
+        "residual_gate",
+        "residual_release_event",
+    }
+
+    # Data keys where a per-sample value of ``None`` is a meaningful signal
+    # (e.g. audio extraction failed for that sample → ``sound=None`` paired
+    # with ``plan.has_sound=False``).  These keys must be kept as a list with
+    # ``None`` placeholders so the model can align per-sample data 1:1 with
+    # per-sample plans.  Dropping the entire key on any None would leave the
+    # remaining sound tensors mis-aligned with the plans whose ``has_sound``
+    # flag was set BEFORE collation, causing ``sequence_packing`` to index
+    # past the end of ``x0_tokens_sound``.
+    sparse_data_keys = {"sound"}
+
+    # Handle the case where the batch is already a dictionary (e.g. column-wise batching)
+    if isinstance(batch, dict):
+        return {key: (value if key in list_collate_keys else default_collate(value)) for key, value in batch.items()}
+
+    # Handle standard list of samples
+    elem = batch[0]
+    if isinstance(elem, dict):
+        # Some Action datasets add optional metadata keys (for example
+        # ``additional_view_description`` for concat-view captions) only for a
+        # subset of samples.  PyTorch can batch such samples together when
+        # DataLoader batch_size > 1; collating only elem's keys and indexing
+        # every sample by that key turns the optional field into a fatal
+        # KeyError.  Use the union of keys and skip optional keys that are not
+        # present in every sample.  Required training keys still fail loudly via
+        # downstream assertions if actually missing.
+        result = {}
+        keys = set().union(*(d.keys() for d in batch))
+        for key in keys:
+            if key in _TIMING_KEYS:
+                continue
+            values = [d.get(key) for d in batch]
+            if key == "action_processing_record":
+                result[key] = values
+                continue
+            if any(value is None for value in values):
+                # Sparse data keys keep their None placeholders to preserve
+                # 1:1 alignment with sequence_plan.  Other (optional metadata)
+                # keys not present in every sample are dropped.
+                if key in sparse_data_keys:
+                    result[key] = values
+                continue
+            if key in list_collate_keys:
+                result[key] = values
+            else:
+                result[key] = default_collate(values)
+        result.update(_aggregate_worker_timing(batch))
+        return result
+    else:
+        return default_collate(batch)
+
+
+def _aggregate_worker_timing(samples: list[dict]) -> dict:
+    """Extract per-sample timing keys, aggregate into per-batch scalars."""
+    info: dict[str, float | int] = {}
+    if "_sample_time" in samples[0]:
+        info["_worker_batch_time"] = sum(s.get("_sample_time", 0.0) for s in samples)
+    if "_aug_time" in samples[0]:
+        aug_total = sum(s.get("_aug_time", 0.0) for s in samples)
+        info["_worker_aug_time"] = aug_total
+        if "_worker_batch_time" in info:
+            info["_worker_io_time"] = info["_worker_batch_time"] - aug_total
+    if "_aug_step_times" in samples[0]:
+        agg: dict[str, float] = {}
+        for s in samples:
+            for step_name, t in s.get("_aug_step_times", {}).items():
+                agg[step_name] = agg.get(step_name, 0.0) + t
+        info["_worker_aug_step_times"] = agg
+    worker_info = torch.utils.data.get_worker_info()
+    info["_worker_id"] = worker_info.id if worker_info is not None else 0
+    return info
+
+
+@dataclass
+class _PackingMetrics:
+    """Per-batch packing statistics collected during the packing loop.
+
+    Also serves as the single source of truth for packing-related metric names
+    via ``STATS_SPEC``, which the dataloading monitor callback consumes to
+    drive accumulation and logging.
+    """
+
+    current_sequence_length: int = 0
+    num_samples: int = 0
+    dropped_count: int = 0
+    from_buffer: int = 0
+    from_workers: int = 0
+
+    STATS_SPEC: ClassVar[list[tuple[str, str, str]]] = [
+        # (batch_key, wandb_suffix, aggregation_type)
+        ("_num_tokens", "token_fraction", "scalar"),
+        ("_num_samples", "samples_per_batch", "list"),
+        ("_from_buffer", "from_buffer", "list"),
+        ("_from_workers", "from_workers", "list"),
+        ("_buffer_size", "buffer_size", "list"),
+        ("_dropped_count", "dropped", "scalar"),
+    ]
+
+    def attach_to(self, output_batch: dict, buffer_size: int) -> None:
+        """Write packing statistics into the output batch dict."""
+        output_batch["_num_tokens"] = self.current_sequence_length
+        output_batch["_num_samples"] = self.num_samples
+        output_batch["_from_buffer"] = self.from_buffer
+        output_batch["_from_workers"] = self.from_workers
+        output_batch["_buffer_size"] = buffer_size
+        output_batch["_dropped_count"] = self.dropped_count
+
+
+@dataclass
+class _AsyncBatchBuildResult:
+    """The outcome of building one packed batch."""
+
+    sequence_id: int
+    batch: dict[str, Any] | None = None
+    error: BaseException | None = None
+    end_of_stream: bool = False
+
+
+class _AsyncBatchBuilder:
+    """Build one requested batch at a time on a persistent daemon thread.
+
+    The caller submits one sequence ID, the daemon advances the existing
+    packing iterator once, and the caller receives that exact batch object
+    from the result queue. A second request is rejected until the first result
+    is consumed, so at most one future batch is built at a time.
+    """
+
+    def __init__(self, batch_iterator: Iterator[dict[str, Any]], timeout_s: float) -> None:
+        self._batch_iterator: Iterator[dict[str, Any]] = batch_iterator
+        self._timeout_s: float = timeout_s
+        self._requests: queue.Queue[int | None] = queue.Queue()
+        self._results: queue.Queue[_AsyncBatchBuildResult | None] = queue.Queue()
+        self._outstanding_sequence_id: int | None = None
+        self._closed: bool = False
+        self._thread: threading.Thread = threading.Thread(target=self._run, name="async-batch-builder", daemon=True)
+        self._thread.start()
+
+    @property
+    def is_alive(self) -> bool:
+        return self._thread.is_alive()
+
+    def request_batch(self, sequence_id: int) -> None:
+        """Ask the daemon thread to build exactly one batch."""
+        if self._closed:
+            raise RuntimeError("Cannot request a batch after the async batch builder is closed.")
+        if self._outstanding_sequence_id is not None:
+            raise RuntimeError(
+                "Cannot request another batch while async batch building is already in progress "
+                f"for sequence_id={self._outstanding_sequence_id}."
+            )
+        self._outstanding_sequence_id = sequence_id
+        self._requests.put(sequence_id)
+
+    def get_batch(self) -> dict[str, Any]:
+        """Wait for and return the exact batch object produced for the outstanding request."""
+        sequence_id = self._outstanding_sequence_id
+        if sequence_id is None:
+            raise RuntimeError("No async batch building request is outstanding.")
+        if self._closed:
+            raise RuntimeError("Async batch builder was closed while waiting for a batch.")
+
+        try:
+            result = self._results.get(timeout=self._timeout_s)
+        except queue.Empty as error:
+            producer_alive = self.is_alive
+            # A timed-out request may still be running in the daemon thread, so it cannot be
+            # safely abandoned or retried. Close the builder and make the timeout terminal.
+            self.close()
+            raise TimeoutError(
+                f"Timed out waiting for async batch building sequence_id={sequence_id}; "
+                f"producer_alive={producer_alive}."
+            ) from error
+
+        if self._closed or result is None:
+            raise RuntimeError("Async batch builder was closed while waiting for a batch.")
+        if result.sequence_id != sequence_id:
+            raise RuntimeError(
+                "Async batch builder returned an unexpected result: "
+                f"expected sequence_id={sequence_id}, got sequence_id={result.sequence_id}."
+            )
+        self._outstanding_sequence_id = None
+
+        if result.error is not None:
+            raise RuntimeError(f"Async batch building failed for sequence_id={sequence_id}.") from result.error
+        if result.end_of_stream:
+            raise StopIteration
+        if result.batch is None:
+            raise RuntimeError(f"Async batch builder returned no batch for sequence_id={sequence_id}.")
+        return result.batch
+
+    def _run(self) -> None:
+        """Build one batch for each request and publish its result by reference."""
+        while True:
+            sequence_id = self._requests.get()
+            if sequence_id is None or self._closed:
+                return
+
+            result = _AsyncBatchBuildResult(sequence_id=sequence_id)
+            try:
+                result.batch = next(self._batch_iterator)
+            except StopIteration:
+                result.end_of_stream = True
+            except BaseException as error:
+                result.error = error
+
+            if self._closed:
+                return
+            # queue.Queue stores object references within this process. The batch is neither
+            # copied nor serialized here; the consumer receives the exact dictionary above.
+            self._results.put(result)
+            if result.end_of_stream or result.error is not None:
+                return
+
+    def close(self) -> None:
+        """Permanently close the builder and perform a bounded join of its daemon thread."""
+        if not self._closed:
+            self._closed = True
+            self._requests.put(None)
+            self._results.put(None)
+        if self._thread is not threading.current_thread():
+            self._thread.join(timeout=min(1.0, self._timeout_s))
+
+
+class JointDataLoader(webdataset.WebLoader):
+    r"""
+    A joint dataloader that supports loading both images and videos.
+    """
+
+    _DEFAULT_LOOKAHEAD_LIMIT: ClassVar[int] = 10
+
+    def __init__(
+        self,
+        dataloaders: Dict[str, Dict[str, Union[torch.utils.data.DataLoader, webdataset.WebLoader, int]]],
+        tokenizer_spatial_compression_factor: int,
+        tokenizer_temporal_compression_factor: int,
+        patch_spatial: int,
+        max_sequence_length: int | None,
+        max_samples_per_batch: int | None,
+        sound_latent_fps: float = 0,
+        audio_sample_rate: int = 48000,
+        prewarm: bool = True,
+        default_lookahead_limit: int = _DEFAULT_LOOKAHEAD_LIMIT,
+        lookahead_limits: Dict[str, int] | None = None,
+        uniae_chunk_frames: int | Mapping[str, int] | None = None,
+        uniae_pad_frames: int | None = None,
+    ):
+        """
+        Initialize the JointDataLoader with multiple datasets.
+
+        The effective mini-batch size can be controlled with either max_sequence_length or
+        max_samples_per_batch. To use max_sequence_length, max_samples_per_batch needs to be None.
+        Vice versa, to use max_samples_per_batch, max_sequence_length needs to be None.
+        max_sequence_length and max_samples_per_batch cannot both be None simultaneously.
+
+        Args:
+            dataloaders: key - dataset_name; value - {"dataloader": dataloader, "ratio": data_ratio}
+            tokenizer_spatial_compression_factor: The spatial compression factor of the tokenizer.
+            tokenizer_temporal_compression_factor: The temporal compression factor of the tokenizer.
+            patch_spatial: Spatial pathification factor.
+            max_samples_per_batch: Max number of samples per packed batch (alternative to max_sequence_length).
+            sound_latent_fps: Sound tokenizer latent rate in Hz (e.g. 25). If 0, sound tokens are not counted.
+            audio_sample_rate: Audio sample rate in Hz (e.g. 48000). Used with sound_latent_fps to estimate
+                sound token count.
+            default_lookahead_limit: Packing-loop look-ahead fallback for dataloaders not in
+                ``lookahead_limits``.
+            lookahead_limits: Optional ``{dataset_name: int}`` per-dataloader override.
+            uniae_chunk_frames: Optional UniAE full chunk size, or resolution-keyed chunk sizes.
+            uniae_pad_frames: Optional UniAE boundary padding frames per chunk.
+
+        Example:
+            joint_loader = IterativeJointDataLoader(
+                dataloaders{
+                    "image_data": {
+                        "dataloader": webdataset.WebLoader(...),
+                        "ratio": 4,
+                    },
+                    "video_data": {
+                        "dataloader": torch.utils.data.DataLoader(...),
+                        "ratio": 1,
+                    },
+                }
+            )
+        """
+        self.dataloader_list, self.dataset_name_list, self.data_ratios = [], [], []
+        self.lookahead_limits: list[int] = []
+        self.tokenizer_spatial_compression_factor = tokenizer_spatial_compression_factor
+        self.tokenizer_temporal_compression_factor = tokenizer_temporal_compression_factor
+        self.patch_spatial = patch_spatial
+        self.max_sequence_length = max_sequence_length
+        self.max_samples_per_batch = max_samples_per_batch
+        self.sound_latent_fps = sound_latent_fps
+        self.audio_sample_rate = audio_sample_rate
+        self.default_lookahead_limit = int(default_lookahead_limit)
+        self.uniae_pad_frames = int(uniae_pad_frames) if uniae_pad_frames is not None else None
+        self.uniae_chunk_frames = self._normalize_uniae_chunk_frames(uniae_chunk_frames)
+
+        assert (self.max_sequence_length is None) != (self.max_samples_per_batch is None), (
+            "Exactly one of max_sequence_length or max_samples_per_batch must be None, but not both."
+        )
+
+        _lookahead_overrides: Dict[str, int] = dict(lookahead_limits) if lookahead_limits else {}
+        unknown = set(_lookahead_overrides) - set(dataloaders)
+        assert not unknown, f"lookahead_limits references unknown dataloaders {unknown}; valid: {sorted(dataloaders)}"
+
+        for dataset_name, dataloader_data in dataloaders.items():
+            if dataloader_data is None:
+                continue
+            assert set(dataloader_data.keys()) == {"dataloader", "ratio"}, f"Invalid config: {dataloader_data}"
+            if dataloader_data["ratio"] <= 0:
+                continue
+            self.dataset_name_list.append(dataset_name)
+            self.dataloader_list.append(instantiate(dataloader_data["dataloader"], collate_fn=custom_collate_fn))
+            self.data_ratios.append(dataloader_data["ratio"])
+            self.lookahead_limits.append(int(_lookahead_overrides.get(dataset_name, self.default_lookahead_limit)))
+
+        self.global_id = 0
+        self.ratio_sum = sum(self.data_ratios)
+
+        total = self.ratio_sum if self.ratio_sum > 0 else 1.0
+        lines = [f"JointDataLoader: {len(self.dataset_name_list)} streams"]
+        for name, ratio in zip(self.dataset_name_list, self.data_ratios):
+            lines.append(f"  {name}: ratio={ratio:.4g} ({ratio / total:.1%})")
+        log.info("\n".join(lines))
+
+        self.data_len = 0
+        self.dataloaders = [iter(dataloader) for dataloader in self.dataloader_list]
+        self.buffers = [deque() for _ in range(len(self.dataloader_list))]
+        for data in self.dataloader_list:
+            self.data_len += len(data)
+
+        # Pre-warm all dataloaders: force worker process spawning and first
+        # batch loading so that slow dataset initialisation (e.g. action
+        # datasets with spawn workers) happens here rather than mid-training
+        # where it would cause NCCL collective timeouts.
+        if prewarm:
+            self._prewarm_dataloaders()
+        else:
+            log.info(
+                "JointDataLoader: prewarm DISABLED (debug mode); first iteration may incur per-stream cold-load cost"
+            )
+
+    def _normalize_uniae_chunk_frames(
+        self, uniae_chunk_frames: int | Mapping[str, int] | None
+    ) -> int | dict[str, int] | None:
+        return normalize_uniae_chunk_frames(
+            uniae_chunk_frames,
+            pad_frames=self.uniae_pad_frames,
+            temporal_compression_factor=self.tokenizer_temporal_compression_factor,
+            temporal_divisibility_name="tokenizer_temporal_compression_factor",
+        )
+
+    def _get_uniae_chunk_frames(self, spatial_shape: tuple[int, int]) -> int:
+        assert self.uniae_chunk_frames is not None
+        return get_uniae_chunk_frames(self.uniae_chunk_frames, spatial_shape=spatial_shape)
+
+    def _compute_vision_latent_t_shape(self, T: int, H: int, W: int) -> int:
+        if T < 1:
+            raise ValueError(f"Vision media must contain at least one frame, got {T}.")
+        if T == 1 or self.uniae_chunk_frames is None:
+            return 1 + (T - 1) // self.tokenizer_temporal_compression_factor
+
+        assert self.uniae_pad_frames is not None
+        return get_uniae_latent_num_frames(
+            T,
+            self.uniae_chunk_frames,
+            pad_frames=self.uniae_pad_frames,
+            temporal_compression_factor=self.tokenizer_temporal_compression_factor,
+            spatial_shape=(H, W),
+        )
+
+    def _prewarm_dataloaders(self) -> None:
+        """Force all dataloader iterators to spawn workers and produce one batch.
+
+        The first ``next()`` call on an ``InfiniteDataLoader`` iterator triggers
+        ``DataLoader.__iter__()`` which spawns worker processes.  For action
+        dataloaders using ``multiprocessing_context='spawn'``, each worker must
+        fully initialise heavy datasets (BridgeOrigLeRobotDataset, embodiment_a, etc.)
+        from scratch.  If this happens lazily during training, the resulting
+        delay (potentially minutes) causes NCCL collective timeouts when faster
+        ranks enter the forward pass while slower ranks are still loading data.
+
+        By pulling one batch from every dataloader here — before any training
+        iteration — we ensure all workers are alive and warmed up.  The fetched
+        samples are pushed into the per-dataloader buffer so they are consumed
+        normally by the first iteration that selects that dataloader.
+
+        A ``dist.barrier()`` at the end synchronises all ranks so that training
+        only begins once every rank has finished pre-warming.
+        """
+        import time
+
+        for i, (name, dl_iter) in enumerate(zip(self.dataset_name_list, self.dataloaders)):
+            t0 = time.monotonic()
+            try:
+                batch = next(dl_iter)
+            except StopIteration:
+                log.warning(f"Pre-warm: dataloader {name!r} is empty, skipping")
+                continue
+            elapsed = time.monotonic() - t0
+
+            # Split the collated batch into individual samples and push them
+            # into the buffer — identical to the splitting logic in
+            # _get_next_sample — so the samples are not wasted.
+            is_image_batch = "images" in batch
+            input_images_or_videos = batch["images" if is_image_batch else "video"]
+            batch_size = len(input_images_or_videos)
+
+            for j in range(batch_size):
+                sample = {}
+                for k, v in batch.items():
+                    if k in _BATCH_TIMING_KEYS:
+                        sample[k] = v
+                    elif isinstance(v, list) and k in self._MULTI_ITEM_KEYS:
+                        elem = v[j]
+                        if isinstance(elem, list):
+                            sample[k] = elem
+                        else:
+                            sample[k] = v[j : j + 1]
+                    elif isinstance(v, list):
+                        sample[k] = v[j]
+                    elif isinstance(v, torch.Tensor) and v.dim() > 0:
+                        sample[k] = v[j : j + 1]
+                    else:
+                        sample[k] = v[j : j + 1]
+                self.buffers[i].append(sample)
+
+            log.info(
+                f"Pre-warm: dataloader {name!r} ready — {batch_size} samples buffered in {elapsed:.1f}s",
+                rank0_only=False,
+            )
+
+        # Synchronise so training only starts once every rank is warmed up.
+        if torch.distributed.is_initialized():
+            log.info("Pre-warm: waiting at barrier for all ranks …")
+            torch.distributed.barrier()
+            log.info("Pre-warm: all ranks ready")
+
+    def _compute_num_tokens_per_sample(self, data_batch: dict) -> int:
+        """
+        This function computes the number of tokens per sample in the data batch.
+        This includes text + vision generation tokens + action tokens.
+
+        Args:
+            data_batch (dict): The data batch containing the text tokens.
+
+        Returns:
+            int: The number of tokens per sample.
+        """
+
+        # The token sequence we have is
+        # <text tokens> <eos> <vision_start> <image tokens> <vision_end> [<action tokens>]
+        # The spatial dimension of image tokens is compressed by
+        # vae spatial downsampling factor + pathification
+        # The temporal dimension of image tokens is compressed by
+        # vae temporal downsampling factor
+        # Action tokens have 1 token per time step (no spatial dimension)
+
+        has_text_tokens = "text_token_ids" in data_batch
+        num_tokens = 0
+        if has_text_tokens:
+            text_token_ids = data_batch["text_token_ids"]
+            if isinstance(text_token_ids, list):
+                num_text_tokens = text_token_ids[0].shape[0]
+            else:
+                num_text_tokens = text_token_ids.shape[1]
+            num_tokens += num_text_tokens + 1
+
+        # Vision part
+        is_image_batch = "images" in data_batch
+        input_images_or_videos = data_batch["images" if is_image_batch else "video"]
+        if "enable_per_camera_vae_encoding" in data_batch:
+            sample_n_views_values = read_positive_int_metadata(data_batch, "sample_n_views", expected_count=1)
+            frames_per_view_values = read_positive_int_metadata(
+                data_batch,
+                "num_video_frames_per_view",
+                expected_count=1,
+            )
+            if sample_n_views_values is None or frames_per_view_values is None:
+                raise ValueError(
+                    "sample_n_views and num_video_frames_per_view are required when per-camera VAE encoding is enabled."
+                )
+            sample_n_views = sample_n_views_values[0]
+            frames_per_view = frames_per_view_values[0]
+        else:
+            sample_n_views = None
+            frames_per_view = None
+
+        # iterate over all the media in the batch
+        for media in input_images_or_videos if isinstance(input_images_or_videos, list) else [input_images_or_videos]:
+            if is_image_batch:
+                _, H, W = media.shape
+                T = 1
+            else:
+                _, T, H, W = media.shape
+
+            latent_h_shape = H // self.tokenizer_spatial_compression_factor
+            latent_w_shape = W // self.tokenizer_spatial_compression_factor
+            patch_h_shape = math.ceil(latent_h_shape / self.patch_spatial)
+            patch_w_shape = math.ceil(latent_w_shape / self.patch_spatial)
+            if sample_n_views is not None and frames_per_view is not None and not is_image_batch:
+                # The multiview dataset resizes every camera to the same H/W and
+                # selects the same synchronized frame count before concatenating
+                # the camera-major clips along T.
+                expected_frames = sample_n_views * frames_per_view
+                if T != expected_frames:
+                    raise ValueError(
+                        "Multiview vision length must equal sample_n_views * num_video_frames_per_view: "
+                        f"got T={T}, sample_n_views={sample_n_views}, num_video_frames_per_view={frames_per_view}."
+                    )
+                latent_t_shape = sample_n_views * self._compute_vision_latent_t_shape(frames_per_view, H, W)
+            else:
+                latent_t_shape = self._compute_vision_latent_t_shape(T, H, W)
+
+            num_vision_tokens = patch_h_shape * patch_w_shape * latent_t_shape
+            if has_text_tokens:
+                num_vision_tokens += 2
+            num_tokens += num_vision_tokens
+
+        # Action part: each action time step is 1 token.
+        # Action tensor shape is (T_action, D) per sample; stored as a single-element list.
+        if "action" in data_batch:
+            list_of_actions = data_batch["action"]
+            for action in list_of_actions:
+                # skip None actions
+                if action is None:
+                    continue
+                num_action_tokens = action.shape[0]
+                num_tokens += num_action_tokens
+
+        # Sound part — estimate sound tokens from audio waveform length
+        if self.sound_latent_fps > 0 and "sound" in data_batch:
+            sound_data = data_batch["sound"]
+            if isinstance(sound_data, list) and len(sound_data) > 0:
+                first_sound = sound_data[0]
+                # Unwrap nested list if needed
+                if isinstance(first_sound, list):
+                    first_sound = first_sound[0]
+                if first_sound is not None and isinstance(first_sound, torch.Tensor):
+                    num_audio_samples = first_sound.shape[-1]
+                    audio_duration = num_audio_samples / self.audio_sample_rate
+                    num_sound_tokens = int(audio_duration * self.sound_latent_fps)
+                    num_tokens += num_sound_tokens
+
+        return num_tokens
+
+    # Keys whose value per sample is a list of tensors to be flattened into one list in the batch
+    _FLATTEN_LIST_KEYS = {"image_size"}
+
+    def _update_output_batch(self, output_batch: dict, output: dict):
+        for key, value in output.items():
+            if key in _BATCH_TIMING_KEYS:
+                if key not in output_batch:
+                    output_batch[key] = value
+            elif key in self._FLATTEN_LIST_KEYS and isinstance(value, list):
+                if key not in output_batch:
+                    output_batch[key] = value
+                else:
+                    output_batch[key].extend(value)
+            elif key not in output_batch:
+                output_batch[key] = [value]
+            else:
+                output_batch[key].append(value)
+
+    def __len__(self) -> int:
+        return self.data_len
+
+    # Keys where each sample may hold multiple tensors (e.g. multiple video
+    # clips in a packed sequence).  Kept as single-element lists per sample
+    # via v[i:i+1] so that _update_output_batch yields list[list[Tensor]].
+    _MULTI_ITEM_KEYS = {"text_token_ids", "images", "video", "action", "action_raw", "sound"}
+
+    def _get_next_sample(self, index_id: int) -> dict:
+        """Pop the next single-sample dict from the buffer for the given dataloader.
+
+        If the buffer is empty, fetches the next collated batch from the inner
+        dataloader and splits it into individual samples.
+
+        Splitting rules:
+            - Multi-item list values (keys in ``_MULTI_ITEM_KEYS``): sliced
+              via ``v[i:i+1]`` to yield a single-element list ``[tensor]``.
+              A packed sequence can contain multiple items per key.
+            - Per-sequence metadata list values (all other list keys, e.g.
+              ``sequence_plan``, ``domain_id``): direct-indexed via ``v[i]``
+              to yield the bare element.
+            - Tensor values ``(B, ...)``: sliced to ``(1, ...)`` via
+              ``v[i : i + 1]`` to preserve the batch dimension.
+
+        After ``_update_output_batch`` accumulates samples, the packed output
+        batch has the following shapes:
+            - Multi-item keys (``text_token_ids``, ``video``, ``images``,
+              ``action_raw``, ``action``): ``list[list[Tensor]]`` — each inner
+              list has one element from one sub-sample.
+            - Per-sequence metadata keys (``sequence_plan``, ``domain_id``,
+              ``dataset_name``, etc.): ``list[element]`` — flat list.
+            - Tensor-origin keys: ``list[Tensor(1, ...)]``.
+
+        Args:
+            index_id: Index of the dataloader to fetch from.
+
+        Returns:
+            A single-sample dictionary.
+        """
+        buffer = self.buffers[index_id]
+        if not buffer:
+            try:
+                batch = next(self.dataloaders[index_id])
+            except StopIteration:
+                raise
+
+            is_image_batch = "images" in batch
+            input_images_or_videos = batch["images" if is_image_batch else "video"]
+            batch_size = len(input_images_or_videos)
+
+            for i in range(batch_size):
+                sample = {}
+                for k, v in batch.items():
+                    if k in _BATCH_TIMING_KEYS:
+                        sample[k] = v
+                    elif isinstance(v, list) and k in self._MULTI_ITEM_KEYS:
+                        # For multi-item keys (images, video, etc.), the collated
+                        # value is a list with one element per sample.  If the element
+                        # is itself a list (e.g. image editing: [src, tgt]), use v[i]
+                        # directly to avoid wrapping it in a redundant single-element
+                        # list.  Otherwise keep the v[i:i+1] slice so that
+                        # _update_output_batch produces list[list[Tensor]].
+                        elem = v[i]
+                        if isinstance(elem, list):
+                            sample[k] = elem
+                        else:
+                            sample[k] = v[i : i + 1]
+                    elif isinstance(v, list):
+                        sample[k] = v[i]
+                    else:
+                        sample[k] = v[i : i + 1]
+                buffer.append(sample)
+
+        return buffer.popleft()
+
+    def set_start_iteration(self, iteration: int):
+        self.global_id = iteration
+
+    def __iter__(self):
+        raise NotImplementedError("__iter__ function is not implemented yet")
+
+
+class IterativeJointDataLoader(JointDataLoader):
+    r"""
+    An iterative joint dataloader that supports loading multiple modalities.
+
+    The behavior depends on the ``seed`` parameter:
+
+    - **seed is not None** (Default):
+      The modality is randomly selected at each iteration based on the probability distribution
+      derived from the ratios. The random state is seeded with ``seed + global_id``, ensuring
+      that all ranks select the same modality at the same iteration (assuming synchronized global_id).
+      This prevents load imbalance due to mixed resolutions across ranks.
+
+    - **seed is None**:
+      The modality selection follows a deterministic round-robin pattern based on the ratios.
+      For example, with 2 modalities (image and video) and ratio 2:1:
+        - Iterations 0, 1: all ranks process images
+        - Iteration 2: all ranks process videos
+        - ... and so on.
+      This also ensures all ranks process the same modality at the same iteration.
+
+    ``enable_async_batch_building=True`` advances this same batch iterator on
+    one daemon thread and builds exactly one future batch. The default keeps
+    the original calling-thread behavior.
+    """
+
+    def __init__(
+        self,
+        dataloaders: Dict[str, Dict[str, Union[torch.utils.data.DataLoader, webdataset.WebLoader, int]]],
+        tokenizer_spatial_compression_factor: int,
+        tokenizer_temporal_compression_factor: int,
+        patch_spatial: int,
+        max_sequence_length: int | None = None,
+        max_samples_per_batch: int | None = None,
+        sound_latent_fps: float = 0,
+        audio_sample_rate: int = 48000,
+        seed: int | None = 42,
+        prewarm: bool = True,
+        default_lookahead_limit: int = JointDataLoader._DEFAULT_LOOKAHEAD_LIMIT,
+        lookahead_limits: Dict[str, int] | None = None,
+        uniae_chunk_frames: int | Mapping[str, int] | None = None,
+        uniae_pad_frames: int | None = None,
+        enable_async_batch_building: bool = False,
+        async_batch_building_timeout_s: float = 1200.0,
+    ) -> None:
+        if async_batch_building_timeout_s <= 0:
+            raise ValueError(f"async_batch_building_timeout_s must be positive, got {async_batch_building_timeout_s}.")
+        # Keep PyTorch/WebDataset/Lance worker creation on the main thread.
+        # The packer thread starts only after constructor prewarm has completed.
+        if enable_async_batch_building and not prewarm:
+            raise ValueError("enable_async_batch_building=True requires prewarm=True.")
+
+        super().__init__(
+            dataloaders,
+            tokenizer_spatial_compression_factor,
+            tokenizer_temporal_compression_factor,
+            patch_spatial,
+            max_sequence_length,
+            max_samples_per_batch,
+            sound_latent_fps=sound_latent_fps,
+            audio_sample_rate=audio_sample_rate,
+            prewarm=prewarm,
+            default_lookahead_limit=default_lookahead_limit,
+            lookahead_limits=lookahead_limits,
+            uniae_chunk_frames=uniae_chunk_frames,
+            uniae_pad_frames=uniae_pad_frames,
+        )
+        self.seed = seed
+        # Calculate probabilities for random sampling
+        total_ratio = sum(self.data_ratios)
+        self.data_probs = np.array([ratio / total_ratio for ratio in self.data_ratios])
+        # Async batch builder state
+        self.enable_async_batch_building: bool = enable_async_batch_building
+        self.async_batch_building_timeout_s: float = float(async_batch_building_timeout_s)
+        self._async_batch_builder: _AsyncBatchBuilder | None = None
+
+    def __iter__(self) -> Iterator[dict[str, Any]]:
+        if not self.enable_async_batch_building:
+            return self._iter_synchronous()
+        return self._iter_asynchronous()
+
+    def _iter_synchronous(self) -> Iterator[dict[str, Any]]:
+        """Run the existing packing path on the calling thread."""
+        while True:
+            if self.seed is not None:
+                rng = np.random.RandomState(self.seed + self.global_id)
+                index_id = rng.choice(len(self.dataloader_list), p=self.data_probs)
+            else:
+                data_id = self.global_id % self.ratio_sum
+                index_id = self._get_dataloader_index(data_id)
+
+            metrics = _PackingMetrics()
+            output_batch = dict()
+            skipped_samples = deque()
+            lookahead_limit = self.lookahead_limits[index_id]
+            lookahead_count = 0
+
+            while True:
+                # Check max samples limit first
+                if self.max_samples_per_batch is not None and metrics.num_samples >= self.max_samples_per_batch:
+                    break
+
+                # If we have started packing and tried lookahead_limit times to find a fitting sample but failed, stop.
+                if len(output_batch) > 0 and lookahead_count >= lookahead_limit:
+                    break
+
+                had_buffer = len(self.buffers[index_id]) > 0
+                try:
+                    output = self._get_next_sample(index_id)
+                except StopIteration:
+                    break  # No more data in this dataloader
+
+                if had_buffer:
+                    metrics.from_buffer += 1
+                else:
+                    metrics.from_workers += 1
+
+                num_tokens_in_current_sample = self._compute_num_tokens_per_sample(output)
+
+                if (
+                    self.max_sequence_length is not None
+                    and metrics.current_sequence_length + num_tokens_in_current_sample >= self.max_sequence_length
+                ):
+                    if len(output_batch) == 0:
+                        # This case happens when current_sequence_length = 0 and num_tokens_in_current_sample > self.max_sequence_length
+                        # In this case, we should simply discard the current sample and get the next sample.
+                        log.info(
+                            f"Discarding oversized sample with {num_tokens_in_current_sample} tokens. Max sequence length: {self.max_sequence_length}",
+                            rank0_only=False,
+                        )
+                        metrics.dropped_count += 1
+                        continue
+
+                    # current_sequence_length > 0 and selected sample is too large to fit in the remaining space.
+                    # Instead of stopping immediately (creating large padding), we buffer this large sample
+                    # and try to find a smaller one that fits in the remaining space.
+                    skipped_samples.append(output)
+                    lookahead_count += 1
+                    continue
+
+                metrics.current_sequence_length += num_tokens_in_current_sample
+                metrics.num_samples += 1
+                output["dataset_name"] = self.dataset_name_list[index_id]
+                self._update_output_batch(output_batch, output)
+
+            # Add back skipped samples to the buffer for the next batch.
+            # appendleft puts item at HEAD. So we insert S3, then S2, then S1.
+            for sample in reversed(skipped_samples):
+                self.buffers[index_id].appendleft(sample)
+
+            if len(output_batch) == 0:
+                return
+
+            metrics.attach_to(output_batch, buffer_size=len(self.buffers[index_id]))
+            self.global_id += 1
+            yield output_batch
+
+    def _start_async_batch_builder(self) -> _AsyncBatchBuilder:
+        """Create the async builder and request its first batch."""
+        if self._async_batch_builder is not None:
+            raise RuntimeError("Async batch building has already started.")
+        async_batch_builder = _AsyncBatchBuilder(
+            self._iter_synchronous(),
+            timeout_s=self.async_batch_building_timeout_s,
+        )
+        self._async_batch_builder = async_batch_builder
+        async_batch_builder.request_batch(self.global_id)
+        return async_batch_builder
+
+    def _iter_asynchronous(self) -> Iterator[dict[str, Any]]:
+        """Yield each completed batch while the async builder builds its replacement."""
+        async_batch_builder = self._async_batch_builder or self._start_async_batch_builder()
+
+        while True:
+            try:
+                batch = async_batch_builder.get_batch()
+            except StopIteration:
+                break
+
+            # Request one replacement before yielding. The request queue contains only a
+            # sequence ID, and the result queue returns this exact batch object by reference.
+            async_batch_builder.request_batch(self.global_id)
+            yield batch
+
+    def set_start_iteration(self, iteration: int) -> None:
+        """Set the batch sequence and eagerly request the first async batch."""
+        if self._async_batch_builder is not None:
+            raise RuntimeError("set_start_iteration must be called before async batch building starts.")
+        self.global_id = iteration
+        if self.enable_async_batch_building:
+            self._start_async_batch_builder()
+
+    def close(self) -> None:
+        """Permanently stop the optional async batch builder."""
+        if self._async_batch_builder is not None:
+            # Keep the closed builder reference: its bounded join may return while the daemon
+            # is still blocked in data loading, so starting a replacement could race with it.
+            self._async_batch_builder.close()
+
+    def _get_dataloader_index(self, data_id: float) -> int:
+        """Maps global id to the corresponding dataloader index based on ratio."""
+        for i, r in enumerate(self.data_ratios):
+            if data_id < r:
+                return i
+            data_id -= r
+        raise ValueError("Invalid data_id")
+
+
+class RankPartitionedDataLoader:
+    """Assigns each rank to exactly one dataset based on ratios.
+
+    For N GPUs with datasets having ratios r_1:r_2:...:r_k, the first
+    N * r_1 / sum(r) ranks are assigned dataset 1, the next N * r_2 / sum(r)
+    ranks are assigned dataset 2, etc.  Each rank instantiates a single
+    PyTorch DataLoader for its assigned dataset.
+
+    The sharding information (``shard_world_size`` and ``shard_rank``) is set
+    on each dataset so that it shards data only across ranks that share the
+    same dataset, rather than across the full world.
+
+    Example:
+        With 128 GPUs and datasets ``{"video": {"dataset": ..., "ratio": 3},
+        "image": {"dataset": ..., "ratio": 1}}``:
+
+        - Ranks   0-95  -> video  (shard_world_size=96, shard_rank=0..95)
+        - Ranks  96-127 -> image  (shard_world_size=32, shard_rank=0..31)
+    """
+
+    def __init__(
+        self,
+        datasets: dict[str, dict[str, Any]],
+        **dataloader_kwargs: Any,
+    ):
+        """
+        Args:
+            datasets: Mapping of dataset name to config dict with keys:
+
+                - ``"dataset"`` (required): a lazy config or dataset instance.
+                - ``"ratio"`` (required): positive int weight.
+                - ``"dataloader_kwargs"`` (optional): dict of keyword arguments
+                  that override the top-level ``**dataloader_kwargs`` for this
+                  dataset only (e.g. different ``num_workers`` or ``batch_size``).
+
+            **dataloader_kwargs: Default kwargs forwarded to
+                ``torch.utils.data.DataLoader``. ``collate_fn`` defaults to
+                ``custom_collate_fn`` if not given.
+        """
+        world_size = torch.distributed.get_world_size()
+        rank = torch.distributed.get_rank()
+        log.info(f"RankPartitionedDataLoader: world_size: {world_size} and rank: {rank}", rank0_only=False)
+
+        _VALID_KEYS = {"dataset", "ratio", "dataloader_kwargs"}
+        names: list[str] = []
+        dataset_configs: list[Any] = []
+        ratios: list[int] = []
+        per_dataset_kwargs: list[dict[str, Any]] = []
+        for name, cfg in datasets.items():
+            extra = set(cfg.keys()) - _VALID_KEYS
+            assert not extra, f"Dataset {name!r}: unexpected keys {extra}. Allowed: {_VALID_KEYS}"
+            if cfg["ratio"] <= 0:
+                log.warning(
+                    f"RankPartitionedDataLoader: Skipping dataset {name} with ratio {cfg['ratio']}", rank0_only=False
+                )
+                continue
+            names.append(name)
+            dataset_configs.append(cfg["dataset"])
+            ratios.append(cfg["ratio"])
+            per_dataset_kwargs.append(cfg.get("dataloader_kwargs", {}))
+
+        assert len(names) > 0, "No datasets with positive ratios provided."
+        assert world_size >= len(names), (
+            f"world_size ({world_size}) must be >= number of datasets ({len(names)}) "
+            f"so each dataset gets at least one rank."
+        )
+
+        total_ratio = sum(ratios)
+        ideal = [r / total_ratio * world_size for r in ratios]
+        allocations = [max(1, int(q)) for q in ideal]
+        remaining = world_size - sum(allocations)
+        if remaining > 0:
+            remainders = sorted(range(len(ratios)), key=lambda i: ideal[i] - allocations[i], reverse=True)
+            for j in range(remaining):
+                allocations[remainders[j]] += 1
+        elif remaining < 0:
+            deficit = -remaining
+            while deficit > 0:
+                best = max(
+                    (i for i in range(len(allocations)) if allocations[i] > 1),
+                    key=lambda i: (allocations[i] - ideal[i], allocations[i]),
+                )
+                allocations[best] -= 1
+                deficit -= 1
+
+        expected_ratios = [r / total_ratio for r in ratios]
+        actual_ratios = [a / world_size for a in allocations]
+        lines = [f"RankPartitionedDataLoader allocation ({world_size} GPUs):"]
+        start = 0
+        for i, (name, alloc) in enumerate(zip(names, allocations)):
+            end = start + alloc - 1
+            lines.append(
+                f"  {name} (ratio {ratios[i]}): ranks {start}-{end} ({alloc} GPUs) "
+                f"| expected {expected_ratios[i]:.2%}, actual {actual_ratios[i]:.2%}"
+            )
+            start += alloc
+        log.info("\n".join(lines), rank0_only=False)
+
+        cumulative = 0
+        my_dataset_idx = -1
+        for i, alloc in enumerate(allocations):
+            if rank < cumulative + alloc:
+                my_dataset_idx = i
+                break
+            cumulative += alloc
+        assert my_dataset_idx >= 0
+
+        shard_rank = rank - cumulative
+        shard_world_size = allocations[my_dataset_idx]
+
+        dataset: Any = instantiate(dataset_configs[my_dataset_idx])
+        dataset.shard_world_size = shard_world_size
+        dataset.shard_rank = shard_rank
+        dataset.shard_id = my_dataset_idx
+
+        merged_kwargs = {**dataloader_kwargs, **per_dataset_kwargs[my_dataset_idx]}
+        merged_kwargs.setdefault("collate_fn", custom_collate_fn)
+        self.dataloader = torch.utils.data.DataLoader(dataset, **merged_kwargs)
+        self.dataset_name = names[my_dataset_idx]
+        self.dataset = dataset
+
+    def __iter__(self):
+        return iter(self.dataloader)
+
+    def __len__(self) -> int:
+        return len(self.dataloader)
+
+
+class PackingDataLoader(JointDataLoader):
+    """Packs multiple samples from a single dataloader into token-budget-constrained batches.
+
+    Unlike the other ``JointDataLoader`` subclasses which manage multiple
+    dataloaders with configurable ratios, this class wraps a single dataloader
+    and greedily packs consecutive samples until the token budget
+    (``max_sequence_length``) or sample count limit (``max_samples_per_batch``)
+    is reached.
+    """
+
+    def __init__(
+        self,
+        dataloader: torch.utils.data.DataLoader | webdataset.WebLoader,
+        tokenizer_spatial_compression_factor: int,
+        tokenizer_temporal_compression_factor: int,
+        patch_spatial: int,
+        max_sequence_length: int | None = None,
+        max_samples_per_batch: int | None = None,
+        sound_latent_fps: float = 0,
+        audio_sample_rate: int = 48000,
+        dataset_name: str = "default",
+        lookahead_limit: int = JointDataLoader._DEFAULT_LOOKAHEAD_LIMIT,
+        uniae_chunk_frames: int | Mapping[str, int] | None = None,
+        uniae_pad_frames: int | None = None,
+    ):
+        """
+        Args:
+            dataloader: A single dataloader (or lazy config) to draw samples from.
+            tokenizer_spatial_compression_factor: Spatial compression factor of the tokenizer.
+            tokenizer_temporal_compression_factor: Temporal compression factor of the tokenizer.
+            patch_spatial: Spatial patchification factor.
+            max_sequence_length: Max total tokens per packed batch. Mutually exclusive with
+                ``max_samples_per_batch``.
+            max_samples_per_batch: Max number of samples per packed batch. Mutually exclusive
+                with ``max_sequence_length``.
+            sound_latent_fps: Sound tokenizer latent rate in Hz. If 0, sound tokens are not counted.
+            audio_sample_rate: Audio sample rate in Hz.
+            dataset_name: Name tag attached to every sample in the output batch.
+            lookahead_limit: Packing-loop look-ahead for the wrapped dataloader.
+            uniae_chunk_frames: Optional UniAE full chunk size, or resolution-keyed chunk sizes.
+            uniae_pad_frames: Optional UniAE boundary padding frames per chunk.
+        """
+        wrapped = {dataset_name: {"dataloader": dataloader, "ratio": 1}}
+        super().__init__(
+            dataloaders=wrapped,
+            tokenizer_spatial_compression_factor=tokenizer_spatial_compression_factor,
+            tokenizer_temporal_compression_factor=tokenizer_temporal_compression_factor,
+            patch_spatial=patch_spatial,
+            max_sequence_length=max_sequence_length,
+            max_samples_per_batch=max_samples_per_batch,
+            sound_latent_fps=sound_latent_fps,
+            audio_sample_rate=audio_sample_rate,
+            lookahead_limits={dataset_name: int(lookahead_limit)},
+            uniae_chunk_frames=uniae_chunk_frames,
+            uniae_pad_frames=uniae_pad_frames,
+        )
+
+    def __iter__(self):
+        inner = self.dataloader_list[0]
+        ds_name = getattr(inner, "dataset_name", self.dataset_name_list[0])
+
+        while True:
+            current_sequence_length = 0
+            num_samples = 0
+            output_batch: dict = {}
+
+            skipped_samples: deque = deque()
+            # PackingDataLoader wraps a single dataloader, so lookahead_limits has one entry.
+            lookahead_limit = self.lookahead_limits[0]
+            lookahead_count = 0
+
+            while True:
+                if self.max_samples_per_batch is not None and num_samples >= self.max_samples_per_batch:
+                    break
+
+                if len(output_batch) > 0 and lookahead_count >= lookahead_limit:
+                    break
+
+                try:
+                    output = self._get_next_sample(0)
+                except StopIteration:
+                    break
+
+                num_tokens_in_current_sample = self._compute_num_tokens_per_sample(output)
+
+                if (
+                    self.max_sequence_length is not None
+                    and current_sequence_length + num_tokens_in_current_sample >= self.max_sequence_length
+                ):
+                    if len(output_batch) == 0:
+                        # This case happens when current_sequence_length = 0 and num_tokens_in_current_sample > self.max_sequence_length
+                        # In this case, we should simply discard the current sample and get the next sample.
+                        log.error(
+                            f"PackingDataLoader: Discarding oversized sample with {num_tokens_in_current_sample} tokens. Max sequence length: {self.max_sequence_length}",
+                            rank0_only=False,
+                        )
+                        continue
+
+                    skipped_samples.append(output)
+                    lookahead_count += 1
+                    continue
+
+                current_sequence_length += num_tokens_in_current_sample
+                num_samples += 1
+                # Allows the dataset name to be overridden by the sample itself.
+                output["dataset_name"] = output.get("dataset_name", ds_name)
+                self._update_output_batch(output_batch, output)
+
+            for sample in reversed(skipped_samples):
+                self.buffers[0].appendleft(sample)
+
+            if len(output_batch) == 0:
+                return
+
+            self.global_id += 1
+            yield output_batch
+
+
+class RandomJointDataLoader(JointDataLoader):
+    r"""
+    A random joint dataloader that supports loading multiple modalities with stochastic sampling.
+
+    In this dataloader, the modality is randomly selected at each iteration based on the
+    probability distribution derived from the ratios. Each rank independently samples a
+    modality, so different ranks may process different modalities at the same iteration.
+
+    For example, with 2 modalities (image and video) and ratio 2:1:
+        - Each iteration has 66.7% probability of selecting images
+        - Each iteration has 33.3% probability of selecting videos
+        - The selection is independent across iterations and ranks
+
+    Note: Unlike IterativeJointDataLoader, this does not guarantee synchronized modality
+    selection across ranks.
+    """
+
+    def __init__(
+        self,
+        dataloaders: Dict[str, Dict[str, Union[torch.utils.data.DataLoader, webdataset.WebLoader, int]]],
+        tokenizer_spatial_compression_factor: int,
+        tokenizer_temporal_compression_factor: int,
+        patch_spatial: int,
+        max_sequence_length: int | None = None,
+        max_samples_per_batch: int | None = None,
+        sound_latent_fps: float = 0,
+        audio_sample_rate: int = 48000,
+        default_lookahead_limit: int = JointDataLoader._DEFAULT_LOOKAHEAD_LIMIT,
+        lookahead_limits: Dict[str, int] | None = None,
+        uniae_chunk_frames: int | Mapping[str, int] | None = None,
+        uniae_pad_frames: int | None = None,
+    ):
+        super().__init__(
+            dataloaders,
+            tokenizer_spatial_compression_factor,
+            tokenizer_temporal_compression_factor,
+            patch_spatial,
+            max_sequence_length,
+            max_samples_per_batch,
+            sound_latent_fps=sound_latent_fps,
+            audio_sample_rate=audio_sample_rate,
+            default_lookahead_limit=default_lookahead_limit,
+            lookahead_limits=lookahead_limits,
+            uniae_chunk_frames=uniae_chunk_frames,
+            uniae_pad_frames=uniae_pad_frames,
+        )
+
+        # Convert data ratios to probabilities
+        self.data_ratios = np.array([ratio / sum(self.data_ratios) for ratio in self.data_ratios])
+
+    def __iter__(self):
+        while True:
+            index_id = np.random.choice(len(self.dataloader_list), p=self.data_ratios)
+
+            metrics = _PackingMetrics()
+            output_batch = dict()
+            skipped_samples = deque()
+            lookahead_limit = self.lookahead_limits[index_id]
+            lookahead_count = 0
+
+            while True:
+                # Check max samples limit first
+                if self.max_samples_per_batch is not None and metrics.num_samples >= self.max_samples_per_batch:
+                    break
+
+                # If we have started packing and tried lookahead_limit times to find a fitting sample but failed, stop.
+                if len(output_batch) > 0 and lookahead_count >= lookahead_limit:
+                    break
+
+                had_buffer = len(self.buffers[index_id]) > 0
+                try:
+                    output = self._get_next_sample(index_id)
+                except StopIteration:
+                    break  # No more data in this dataloader
+
+                if had_buffer:
+                    metrics.from_buffer += 1
+                else:
+                    metrics.from_workers += 1
+
+                num_tokens_in_current_sample = self._compute_num_tokens_per_sample(output)
+
+                if (
+                    self.max_sequence_length is not None
+                    and metrics.current_sequence_length + num_tokens_in_current_sample >= self.max_sequence_length
+                ):
+                    if len(output_batch) == 0:
+                        # This case happens when current_sequence_length = 0 and num_tokens_in_current_sample > self.max_sequence_length
+                        # In this case, we should simply discard the current sample and get the next sample.
+                        log.info(
+                            f"Discarding oversized sample with {num_tokens_in_current_sample} tokens. Max sequence length: {self.max_sequence_length}",
+                            rank0_only=False,
+                        )
+                        metrics.dropped_count += 1
+                        continue
+
+                    # current_sequence_length > 0 and selected sample is too large to fit in the remaining space.
+                    # Instead of stopping immediately (creating large padding), we buffer this large sample
+                    # and try to find a smaller one that fits in the remaining space.
+                    skipped_samples.append(output)
+                    lookahead_count += 1
+                    continue
+
+                metrics.current_sequence_length += num_tokens_in_current_sample
+                metrics.num_samples += 1
+                output["dataset_name"] = self.dataset_name_list[index_id]
+                self._update_output_batch(output_batch, output)
+
+            # Add back skipped samples to the buffer for the next batch.
+            # appendleft puts item at HEAD. So we insert S3, then S2, then S1.
+            for sample in reversed(skipped_samples):
+                self.buffers[index_id].appendleft(sample)
+
+            if len(output_batch) == 0:
+                return
+
+            metrics.attach_to(output_batch, buffer_size=len(self.buffers[index_id]))
+            yield output_batch

releases/bridge-native-v3/code/cosmos_framework/data/generator/sequence_packing/modality.py

@@ -0,0 +1,256 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""Modality utility helpers for VFM sequence packing."""
+
+from dataclasses import dataclass, field
+
+import torch
+
+
+def _empty_long_tensor() -> torch.Tensor:
+    return torch.empty(0, dtype=torch.long)  # [0]
+
+
+def _empty_float_tensor() -> torch.Tensor:
+    return torch.empty(0, dtype=torch.float32)  # [0]
+
+
+@dataclass
+class ModalitySpan:
+    """One contiguous packed span paired with its logical modality payload slice.
+
+    Attributes:
+        sequence_start: First global packed-sequence index owned by this span.
+        sequence_len: Number of contiguous tokens in the packed sequence.
+        payload_index: Index into ``ModalityData.tokens`` for the backing payload tensor.
+        payload_start: First flattened token offset within the backing payload tensor.
+        payload_len: Number of flattened payload tokens covered by this span.
+        payload_shape: Logical payload slice shape. For example, vision frame spans use
+            ``(1, patch_h, patch_w)``, action spans use ``(tcf,)``, and sound spans use
+            ``(1, 1, 1)``.
+    """
+
+    sequence_start: int
+    sequence_len: int
+    payload_index: int
+    payload_start: int
+    payload_len: int
+    payload_shape: tuple[int, ...]
+
+
+@dataclass
+class ModalityDataBuilder:
+    """Append-only construction state for a single generation modality.
+
+    Attributes:
+        spans: Contiguous packed spans pointing back into grouped payload tensors.
+        sequence_indexes: Global packed-sequence indexes for all tokens of this modality.
+        timesteps: Diffusion timesteps for noised tokens only.
+        mse_loss_indexes: Global packed-sequence indexes where MSE loss should be computed.
+        token_shapes: Shape metadata for each payload tensor. Vision and sound use
+            ``(T, H, W)``-style tuples; action uses ``(T,)``.
+        tokens: Original modality payload tensors, kept grouped by sample/item.
+        condition_mask: Per-payload masks where 1 indicates clean/conditioning tokens
+            and 0 indicates noised/supervised tokens.
+        noisy_frame_indexes: Per-payload indexes of noised frames. These are constructed
+            during packing to avoid GPU-to-CPU synchronization later.
+    """
+
+    spans: list[ModalitySpan] = field(default_factory=list)
+    sequence_indexes: list[int] = field(default_factory=list)
+    timesteps: list[float] = field(default_factory=list)
+    mse_loss_indexes: list[int] = field(default_factory=list)
+    # list[tuple[int,int,int]] for vision, list[tuple[int]] for action, list[tuple[int,int,int]] for sound
+    token_shapes: list[tuple[int, ...]] = field(default_factory=list)
+    tokens: list[torch.Tensor] = field(default_factory=list)
+    condition_mask: list[torch.Tensor] = field(default_factory=list)
+    noisy_frame_indexes: list[torch.Tensor] = field(default_factory=list)
+
+
+@dataclass
+class ModalityData:
+    """Finalized model-facing data for a single generation modality.
+
+    Index-like fields are tensors after packing. Payload fields remain lists
+    because samples may have variable token shapes and because downstream code
+    mutates token payloads after noise injection and clean replay.
+
+    Attributes:
+        sequence_indexes: Tensor of global packed-sequence indexes for all tokens of
+            this modality.
+        timesteps: Tensor of diffusion timesteps for noised tokens only.
+        mse_loss_indexes: Tensor of global packed-sequence indexes where MSE loss
+            should be computed.
+        spans: Contiguous packed spans pointing back into grouped payload tensors.
+        token_shapes: Shape metadata for each payload tensor. Vision and sound use
+            ``(T, H, W)``-style tuples; action uses ``(T,)``.
+        tokens: Original modality payload tensors, kept grouped by sample/item.
+        condition_mask: Per-payload masks where 1 indicates clean/conditioning tokens
+            and 0 indicates noised/supervised tokens.
+        noisy_frame_indexes: Per-payload indexes of noised frames.
+        domain_id: Domain IDs for multi-domain training. Only used for action.
+        raw_action_dim: Raw action dimensions. Only used for action-channel masking.
+    """
+
+    sequence_indexes: torch.Tensor = field(default_factory=_empty_long_tensor)
+    timesteps: torch.Tensor = field(default_factory=_empty_float_tensor)
+    mse_loss_indexes: torch.Tensor = field(default_factory=_empty_long_tensor)
+    spans: list[ModalitySpan] = field(default_factory=list)
+    # list[tuple[int,int,int]] for vision, list[tuple[int]] for action, list[tuple[int,int,int]] for sound
+    token_shapes: list[tuple[int, ...]] = field(default_factory=list)
+    tokens: list[torch.Tensor] = field(default_factory=list)
+    condition_mask: list[torch.Tensor] = field(default_factory=list)
+    noisy_frame_indexes: list[torch.Tensor] = field(default_factory=list)
+    domain_id: list[torch.Tensor] = field(default_factory=list)
+    raw_action_dim: list[torch.Tensor | None] | None = field(default_factory=list)
+    # Optional per-token gate for an action residual expert.
+    residual_gate: list[torch.Tensor] | None = None
+
+    def __post_init__(self) -> None:
+        assert isinstance(self.sequence_indexes, torch.Tensor), "ModalityData.sequence_indexes must be finalized"
+        assert isinstance(self.timesteps, torch.Tensor), "ModalityData.timesteps must be finalized"
+        assert isinstance(self.mse_loss_indexes, torch.Tensor), "ModalityData.mse_loss_indexes must be finalized"
+
+    def to_cuda(self) -> None:
+        """Move all tensor fields to CUDA in-place."""
+        self.sequence_indexes = self.sequence_indexes.cuda()
+        self.timesteps = self.timesteps.cuda()
+        self.mse_loss_indexes = self.mse_loss_indexes.cuda()
+        self.tokens = [token.cuda() for token in self.tokens]
+        self.condition_mask = [cm.cuda() for cm in self.condition_mask]
+        self.noisy_frame_indexes = [ni.cuda() for ni in self.noisy_frame_indexes]
+        self.domain_id = [d.cuda() for d in self.domain_id]
+        # raw_action_dim is optional (e.g., when action-channel masking is disabled).
+        if self.raw_action_dim is not None:
+            self.raw_action_dim = [d.cuda() if d is not None else None for d in self.raw_action_dim]
+        if self.residual_gate is not None:
+            self.residual_gate = [gate.cuda() for gate in self.residual_gate]
+
+
+def prepare_attention_mask_per_sample(split_lens, attn_modes, device="cpu"):
+    """Prepare dense attention mask for a single sample with multiple splits.
+
+    Args:
+        split_lens: List of integers indicating length of each split within the sample
+        attn_modes: List of attention modes for each split ('causal', 'full', or 'noise')
+        device: Device to place the attention mask tensor on
+
+    Returns:
+        Attention mask tensor of shape (sample_len, sample_len) with -inf for masked positions
+    """
+    sample_len = sum(split_lens)
+    attention_mask = torch.zeros((sample_len, sample_len), dtype=torch.bool, device=device)  # [sample_len,sample_len]
+
+    # First pass: Set up basic attention patterns for each split
+    current_pos = 0
+    for split_len, attn_mode in zip(split_lens, attn_modes):
+        assert attn_mode in ["causal", "full", "noise"], f"Invalid attention mode: {attn_mode}"
+
+        split_start = current_pos
+        split_end = current_pos + split_len
+
+        if attn_mode == "causal":
+            # Causal: lower triangular within split + full attention to previous splits
+            attention_mask[split_start:split_end, split_start:split_end] = torch.ones(
+                (split_len, split_len), device=device
+            ).tril()  # [split_len,split_len]
+            attention_mask[split_start:split_end, :split_start] = 1
+        else:  # "full" or "noise"
+            # Full attention within split and to previous splits
+            attention_mask[split_start:split_end, split_start:split_end] = torch.ones(
+                (split_len, split_len), device=device
+            )  # [split_len,split_len]
+            attention_mask[split_start:split_end, :split_start] = 1
+
+        current_pos += split_len
+
+    # Second pass: Handle noise mode - mask out noise columns except within same split
+    current_pos = 0
+    for split_len, attn_mode in zip(split_lens, attn_modes):
+        if attn_mode == "noise":
+            split_start = current_pos
+            split_end = current_pos + split_len
+
+            # Zero out the entire column for noise tokens
+            attention_mask[:, split_start:split_end] = 0
+            # But allow self-attention within the noise split
+            attention_mask[split_start:split_end, split_start:split_end] = 1
+
+        current_pos += split_len
+
+    # Convert boolean mask to float with -inf for masked positions
+    attention_mask = torch.zeros_like(attention_mask, dtype=torch.float).masked_fill_(
+        ~attention_mask, float("-inf")
+    )  # [sample_len,sample_len]
+
+    return attention_mask
+
+
+# ============================================================================
+# Tokenizer utilities
+# ============================================================================
+
+
+def add_special_tokens(tokenizer):
+    """Add image-related special tokens to tokenizer if not already present.
+
+    Args:
+        tokenizer: Tokenizer to add special tokens to
+
+    Returns:
+        Tuple of (modified tokenizer, dict of new token IDs)
+    """
+    # Collect existing special tokens
+    existing_special_tokens = []
+    for key, value in tokenizer.special_tokens_map.items():
+        if isinstance(value, str):
+            existing_special_tokens.append(value)
+        elif isinstance(value, list):
+            existing_special_tokens.extend(value)
+
+    # Define image boundary tokens to add if missing
+    tokens_to_add = []
+    if "<|vision_start|>" not in existing_special_tokens:
+        tokens_to_add.append("<|vision_start|>")
+    if "<|vision_end|>" not in existing_special_tokens:
+        tokens_to_add.append("<|vision_end|>")
+
+    # Add new tokens to tokenizer vocabulary
+    if tokens_to_add:
+        tokenizer.add_tokens(tokens_to_add)
+
+    # Get token IDs for image boundary tokens
+    new_token_ids = {
+        "start_of_generation": tokenizer.convert_tokens_to_ids("<|vision_start|>"),
+        "end_of_generation": tokenizer.convert_tokens_to_ids("<|vision_end|>"),
+    }
+
+    return tokenizer, new_token_ids
+
+
+def compute_text_split_length(
+    num_caption_tokens: int,
+    special_tokens: dict[str, int],
+    has_generation: bool = True,
+) -> int:
+    """Compute the total text split length without mutating any state.
+
+    This is the number of token positions occupied by the text split in a
+    packed sequence: caption tokens + optional BOS + EOS + optional BOV.
+
+    Args:
+        num_caption_tokens: Number of raw caption token IDs (before special tokens).
+        special_tokens: Dictionary of special token IDs (checked for ``"bos_token_id"``).
+        has_generation: Whether a start-of-generation (BOV) token follows text.
+
+    Returns:
+        Total text split length (positions consumed in the packed sequence).
+    """
+    n = num_caption_tokens
+    if "bos_token_id" in special_tokens:
+        n += 1
+    n += 1  # EOS
+    if has_generation:
+        n += 1  # start-of-generation / BOV
+    return n

releases/bridge-native-v3/code/cosmos_framework/data/generator/sequence_packing/packers.py

@@ -0,0 +1,446 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""Top-level input sequence packing orchestration."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from cosmos_framework.data.generator.sequence_packing.sequence import PackedSequence, PackedSequenceBuilder, SequencePlan
+from cosmos_framework.data.generator.sequence_packing.temporal_causal import pack_supertokens_temporal_causal
+
+if TYPE_CHECKING:
+    from cosmos_framework.model.generator.utils.data_and_condition import GenerationDataClean
+
+
+def _get_optional_fps(
+    fps_values: torch.Tensor | list[torch.Tensor | float] | None,
+    index: int,
+) -> float | None:
+    """Return an optional FPS value as a float.
+
+    Args:
+        fps_values: Optional tensor/list of per-sample or per-item FPS values.
+        index: Entry to read from ``fps_values``.
+
+    Returns:
+        FPS value as ``float`` when present, otherwise ``None``.
+    """
+    if fps_values is None or index >= len(fps_values):
+        return None
+    fps_value = fps_values[index]  # []
+    if isinstance(fps_value, torch.Tensor):
+        return float(fps_value.item())
+    return float(fps_value)
+
+
+def pack_input_sequence(
+    sequence_plans: list[SequencePlan],
+    input_text_indexes: list[list[int]],
+    gen_data_clean: GenerationDataClean,
+    input_timesteps: torch.Tensor,
+    special_tokens: dict[str, int],
+    max_num_tokens: int | None = None,
+    latent_patch_size: int = 1,
+    skip_text_tokens: bool = False,
+    include_end_of_generation_token: bool = False,
+    unified_3d_mrope_reset_spatial_ids: bool = True,
+    unified_3d_mrope_temporal_modality_margin: int = 0,
+    enable_fps_modulation: bool = False,
+    base_fps: float = 24.0,
+    sound_base_temporal_compression_factor: int | None = None,
+    temporal_compression_factor: int = 4,
+    vision_temporal_position_mode: str = "latent_index",
+    video_temporal_causal: bool = False,
+    action_dim: int = 32,
+    initial_mrope_temporal_offset: int | float = 0,
+    num_behavior_prefix_tokens: int = 0,
+    num_online_memory_prefix_tokens: int = 0,
+    num_proprio_prefix_tokens: int = 0,
+) -> PackedSequence:
+    """
+    Pack a sequence of input strings and VAE latents into a packed tensor format.
+    Uses SequencePlan to determine which modalities are present for each sample,
+    and maintains separate indices for text, vision, action, and sound to handle variable modality presence.
+
+    Args:
+        sequence_plans: List of SequencePlan items describing which modalities are present.
+        input_text_indexes: List of text token ID sequences (only for samples where has_text=True).
+        gen_data_clean: GenerationDataClean containing vision, action, and sound tensors.
+            - x0_tokens_vision: Vision tensors for samples where has_vision=True
+            - x0_tokens_action: Action tensors for samples where has_action=True
+            - x0_tokens_sound: Sound tensors (list of [C, T]) for samples where has_sound=True
+        input_timesteps: Diffusion timesteps for each sample. Shape (B,) or (B, 1) for
+            teacher_forcing/none (all frames share the same sigma), or (B, T_max) for
+            diffusion_forcing (per-frame independent sigma). Entries are extracted per
+            sample as a float (numel==1) or Tensor(T_max,) for per-frame indexing.
+        special_tokens: Dictionary containing special token IDs (eos_token_id, start_of_generation, end_of_generation)
+        max_num_tokens: Maximum number of tokens in the packed sequence
+        latent_patch_size: Patch size used by the network to pack latents
+        skip_text_tokens: If True, skip packing text tokens
+        include_end_of_generation_token: If True, append end-of-generation token
+        unified_3d_mrope_reset_spatial_ids: If True (default), spatial (H, W) indices
+            start from 0 for each vision segment. If False, spatial indices are offset
+            by the temporal offset (Qwen2VL-style).
+        unified_3d_mrope_temporal_modality_margin: Extra temporal offset inserted between
+            text and generation modalities.
+        enable_fps_modulation: If True, scale temporal position IDs based on video FPS
+            to reflect real time. Requires fps_vision in gen_data_clean.
+            Uses the same flag as diffusion_expert_config.enable_fps_modulation.
+        base_fps: Base FPS for normalization (default 24.0).
+            Uses the same value as diffusion_expert_config.base_fps.
+        sound_base_temporal_compression_factor: Base temporal compression factor for sound FPS scaling.
+            ``None`` preserves the current behavior where sound advances at ``base_fps`` positions/sec.
+        temporal_compression_factor: VAE temporal compression factor (default 4).
+            Obtained from the VAE tokenizer at runtime.
+        vision_temporal_position_mode: Temporal coordinates used for unified_3d_mrope vision tokens.
+            "latent_index" keeps legacy positions; "uniae_source_right_edge" uses
+            per-latent positions from gen_data_clean.temporal_positions_vision.
+        video_temporal_causal: If True, pack vision and optional action as temporal-causal
+            supertokens instead of separate modality blocks.
+        action_dim: Action feature dimension used when temporal-causal packing creates
+            null action tokens.
+        initial_mrope_temporal_offset: Initial temporal cursor for each sample, used by
+            autoregressive inference to seed mRoPE positions.
+        num_behavior_prefix_tokens: Number of learned causal condition slots to
+            reserve before text for each action-bearing sample. Zero preserves
+            every existing Cosmos packing layout.
+
+    Returns:
+        PackedSequence containing all packed tensors and metadata. See PackedSequence for field details.
+    """
+    del max_num_tokens
+
+    assert special_tokens is not None, "Special tokens must be provided"
+    assert isinstance(input_timesteps, torch.Tensor), "input_timesteps must be a tensor"
+    if input_timesteps.is_cuda:
+        raise ValueError("input_timesteps must be on CPU, not CUDA")
+    if isinstance(input_text_indexes, torch.Tensor):
+        raise ValueError("input_text_tokens must be a list, not a tensor")
+
+    supported_vision_temporal_position_modes = {"latent_index", "uniae_source_right_edge"}
+    if vision_temporal_position_mode not in supported_vision_temporal_position_modes:
+        raise ValueError(
+            "Unsupported vision_temporal_position_mode: "
+            f"{vision_temporal_position_mode}. Supported modes: {supported_vision_temporal_position_modes}."
+        )
+    has_any_vision = any(plan.has_vision for plan in sequence_plans)
+    explicit_vision_temporal_positions_active = vision_temporal_position_mode != "latent_index" and has_any_vision
+    if explicit_vision_temporal_positions_active:
+        if gen_data_clean.temporal_positions_vision is None:
+            raise ValueError(
+                f"vision_temporal_position_mode={vision_temporal_position_mode} requires "
+                "gen_data_clean.temporal_positions_vision."
+            )
+        if gen_data_clean.x0_tokens_vision is not None and len(gen_data_clean.temporal_positions_vision) != len(
+            gen_data_clean.x0_tokens_vision
+        ):
+            raise ValueError(
+                "temporal_positions_vision must have one entry per x0_tokens_vision item, "
+                f"got {len(gen_data_clean.temporal_positions_vision)} positions for "
+                f"{len(gen_data_clean.x0_tokens_vision)} vision items."
+            )
+        if video_temporal_causal:
+            raise NotImplementedError(
+                "video_temporal_causal=True is not wired for explicit UniAE vision temporal positions yet."
+            )
+        if any(plan.has_action for plan in sequence_plans):
+            raise NotImplementedError("Action packing is not wired for explicit UniAE vision temporal positions yet.")
+        if initial_mrope_temporal_offset != 0:
+            raise NotImplementedError(
+                "Autoregressive mRoPE temporal offsets are not wired for explicit UniAE vision temporal positions yet."
+            )
+    use_float_mrope_positions = enable_fps_modulation or explicit_vision_temporal_positions_active
+
+    # Initialize mutable builder state for sequence construction.
+    seq_builder = PackedSequenceBuilder(uses_single_timestep=input_timesteps.numel() == 1)
+
+    # Configure 3D mRoPE on the builder.
+    seq_builder._mrope_reset_spatial = unified_3d_mrope_reset_spatial_ids
+
+    # Maintain separate indices for each modality
+    idx_text = 0
+    idx_vision = 0
+    idx_action = 0
+    idx_sound = 0
+    null_action_flags: list[bool] = []  # collected from TC path; asserted consistent after the loop
+
+    # Validate: all samples must have text (causal split is always required for two-way attention).
+    # CFG dropout only drops text *content*, not the structural text split.
+    if not skip_text_tokens:
+        for plan in sequence_plans:
+            assert plan.has_text, "All sequence plans must have has_text=True when skip_text_tokens=False"
+
+    # Pack each sample based on its sequence plan
+    for sample_idx, sequence_plan in enumerate(sequence_plans):
+        sample_len = 0
+
+        # mRoPE temporal offset resets per sample.
+        # initial_mrope_temporal_offset is non-zero only for AR inference (frame N seeds at N*tcf).
+        seq_builder.begin_sample(initial_mrope_temporal_offset)
+
+        _ts = input_timesteps[sample_idx]
+        input_timestep = _ts.item() if _ts.numel() == 1 else _ts  # float (TF) or Tensor(T_max,) (DF)
+
+        # Pack text tokens if has_text=True and not skipped
+        if sequence_plan.has_text and not skip_text_tokens:
+            text_ids = input_text_indexes[idx_text]
+            idx_text += 1
+
+            has_generation_for_sample = sequence_plan.has_vision or sequence_plan.has_action or sequence_plan.has_sound
+            text_sample_len = seq_builder.pack_text_tokens(
+                text_ids,
+                special_tokens,
+                has_generation=has_generation_for_sample,
+                use_float_positions=use_float_mrope_positions,
+                num_behavior_prefix_tokens=(
+                    num_behavior_prefix_tokens if sequence_plan.has_action else 0
+                ),
+                num_online_memory_prefix_tokens=(
+                    num_online_memory_prefix_tokens if sequence_plan.has_action else 0
+                ),
+                num_proprio_prefix_tokens=(
+                    num_proprio_prefix_tokens if sequence_plan.has_action else 0
+                ),
+            )
+            sample_len += text_sample_len
+
+            # End of text modality, add an offset as the boundary between text and vision.
+            seq_builder.advance_mrope_temporal_offset(unified_3d_mrope_temporal_modality_margin)
+
+        # Save temporal offset before vision for action tokens (action uses same offset as vision start)
+        vision_start_temporal_offset = seq_builder.mrope_temporal_offset
+
+        # Pack vision (and optionally action) tokens
+        if video_temporal_causal and sequence_plan.has_vision:
+            # Temporal causal path: when sequence_plan.has_action=True, interleaved supertokens
+            # [action_t, vision_t]; when False, supertokens are just vision patches.
+            input_vision_tokens = gen_data_clean.x0_tokens_vision[idx_vision]
+            idx_vision += 1
+
+            vision_fps = _get_optional_fps(gen_data_clean.fps_vision, idx_vision - 1)
+
+            input_action_tokens_tc: torch.Tensor | None = None
+            action_fps_tc = None
+            if sequence_plan.has_action:
+                input_action_tokens_tc = gen_data_clean.x0_tokens_action[idx_action]
+                action_fps_tc = _get_optional_fps(gen_data_clean.fps_action, idx_action)
+                idx_action += 1
+
+            supertoken_split_len, null_flag = pack_supertokens_temporal_causal(
+                seq_builder=seq_builder,
+                input_vision_tokens=input_vision_tokens,
+                input_action_tokens=input_action_tokens_tc,
+                condition_frame_indexes_vision=sequence_plan.condition_frame_indexes_vision,
+                input_timestep=input_timestep,
+                latent_patch_size=latent_patch_size,
+                temporal_compression_factor=temporal_compression_factor,
+                action_dim=action_dim,
+                vision_fps=vision_fps,
+                action_fps=action_fps_tc,
+                enable_fps_modulation=enable_fps_modulation,
+                base_fps=base_fps,
+                pack_action_tokens=sequence_plan.has_action,
+            )
+            null_action_flags.append(null_flag)
+            # We assume all samples in a batch share the same has_action layout, so
+            # stamp the supertoken layout constant directly here. This is the
+            # single source of truth read by downstream attention / KV-cache
+            # code (no recomputation in the network).
+            seq_builder.num_action_tokens_per_supertoken = (
+                temporal_compression_factor if sequence_plan.has_action else 0
+            )
+            sample_len += supertoken_split_len
+            vision_split_len = supertoken_split_len
+            action_split_len = 0  # Already absorbed into supertoken_split_len
+
+        else:
+            # Standard path: vision and action packed separately
+            if sequence_plan.has_vision:
+                # Determine how many vision items this sample owns.
+                # For multi-item samples (e.g. image editing), num_vision_items_per_sample
+                # records [2, 2, ...]; for standard T2I/T2V it is None (1 item per sample).
+                num_vis = (
+                    gen_data_clean.num_vision_items_per_sample[sample_idx]
+                    if gen_data_clean.num_vision_items_per_sample is not None
+                    else 1
+                )
+
+                vision_split_len = 0
+                # Per-item split lengths for multi-control attention routing.
+                # Only tracked when control_weights are present (inference-only);
+                # skipped during training to avoid unnecessary side effects.
+                track_item_split_lens = gen_data_clean.control_weights is not None
+                sample_item_split_lens: list[int] = []
+                # Controlnet-style transfer: when set, all vision items share the same
+                # temporal mRoPE grid. We snapshot the offset before the loop and
+                # rewind to it before each item, so every item produces identical
+                # temporal IDs. Each pack_vision_tokens call still advances the
+                # offset by latent_t internally; in shared-grid mode the post-loop
+                # offset equals snapshot + latent_t (single-clip semantics for
+                # downstream EOV / next-modality tokens).
+                shared_grid = sequence_plan.share_vision_temporal_positions and num_vis > 1
+                items_temporal_offset_snapshot = seq_builder.mrope_temporal_offset
+                shared_latent_t: int | None = None
+                shared_patch_h: int | None = None
+                shared_patch_w: int | None = None
+                shared_temporal_positions: torch.Tensor | None = None
+                # FPS is recorded per-sample (shape [B]); for multi-item samples
+                # (transfer / image-edit) every vision item in this sample shares
+                # the same conditioning FPS, so we read by sample_idx, not by the
+                # flat idx_vision counter (which would alias to a neighbor sample's
+                # fps and corrupt RoPE FPS modulation).
+                sample_vision_fps = _get_optional_fps(gen_data_clean.fps_vision, sample_idx)
+
+                for item_idx in range(num_vis):
+                    flat_vision_idx = idx_vision
+                    input_vision_tokens = gen_data_clean.x0_tokens_vision[flat_vision_idx]  # [1,C,T,H,W]
+                    vision_temporal_positions: torch.Tensor | None = None
+                    if explicit_vision_temporal_positions_active:
+                        assert gen_data_clean.temporal_positions_vision is not None
+                        vision_temporal_positions = gen_data_clean.temporal_positions_vision[flat_vision_idx]  # [T]
+                        if vision_temporal_positions.shape[0] != input_vision_tokens.shape[2]:
+                            raise ValueError(
+                                "vision_temporal_positions must match latent_t for each vision item, "
+                                f"got {vision_temporal_positions.shape[0]} positions and "
+                                f"latent_t={input_vision_tokens.shape[2]} for item {flat_vision_idx}."
+                            )
+                    idx_vision += 1
+
+                    # Determine conditioning for this vision item.
+                    # For multi-item mode: all items except the last are fully conditioned
+                    # (all frames are clean); the last item uses the SequencePlan's
+                    # condition_frame_indexes_vision (typically [] = fully generated).
+                    if num_vis > 1 and item_idx < num_vis - 1:
+                        # Conditioning item (e.g. source image): mark all frames as clean
+                        latent_t = input_vision_tokens.shape[2]
+                        item_condition_frames = list(range(latent_t))
+                    else:
+                        # Generation item (single-item mode or last item in multi-item)
+                        item_condition_frames = sequence_plan.condition_frame_indexes_vision
+
+                    if shared_grid:
+                        item_latent_t = input_vision_tokens.shape[2]
+                        item_latent_h = input_vision_tokens.shape[3]
+                        item_latent_w = input_vision_tokens.shape[4]
+                        if shared_latent_t is None:
+                            shared_latent_t = item_latent_t
+                            shared_patch_h = item_latent_h
+                            shared_patch_w = item_latent_w
+                        else:
+                            assert item_latent_t == shared_latent_t, (
+                                f"share_vision_temporal_positions requires equal latent_t across items, "
+                                f"got item {item_idx} latent_t={item_latent_t} vs first={shared_latent_t}"
+                            )
+                            assert item_latent_h == shared_patch_h and item_latent_w == shared_patch_w, (
+                                f"share_vision_temporal_positions requires equal spatial grid across items, "
+                                f"got item {item_idx} (H,W)=({item_latent_h},{item_latent_w}) "
+                                f"vs first=({shared_patch_h},{shared_patch_w})"
+                            )
+                        if vision_temporal_positions is not None:
+                            if shared_temporal_positions is None:
+                                shared_temporal_positions = vision_temporal_positions
+                            else:
+                                comparison_temporal_positions = vision_temporal_positions.to(
+                                    device=shared_temporal_positions.device
+                                )  # [T]
+                                assert torch.allclose(comparison_temporal_positions, shared_temporal_positions), (
+                                    "share_vision_temporal_positions requires equal explicit temporal positions "
+                                    f"across vision items, got item {item_idx} positions "
+                                    f"{vision_temporal_positions.tolist()} vs first "
+                                    f"{shared_temporal_positions.tolist()}."
+                                )
+                        # Rewind so this item starts at the same temporal offset as item 0.
+                        seq_builder.set_mrope_temporal_offset(items_temporal_offset_snapshot)
+
+                    item_split_len = seq_builder.pack_vision_tokens(
+                        input_vision_tokens=input_vision_tokens,
+                        condition_frame_indexes_vision=item_condition_frames,
+                        input_timestep=input_timestep,
+                        latent_patch_size=latent_patch_size,
+                        vision_fps=sample_vision_fps,
+                        enable_fps_modulation=enable_fps_modulation,
+                        base_fps=base_fps,
+                        temporal_compression_factor=temporal_compression_factor,
+                        vision_temporal_positions=vision_temporal_positions,
+                    )
+                    vision_split_len += item_split_len
+                    if track_item_split_lens:
+                        sample_item_split_lens.append(item_split_len)
+
+                if track_item_split_lens:
+                    seq_builder.vision_item_split_lens.append(sample_item_split_lens)
+                sample_len += vision_split_len
+
+            else:
+                vision_split_len = 0
+
+            # Pack action tokens if has_action=True
+            if sequence_plan.has_action:
+                input_action_tokens = gen_data_clean.x0_tokens_action[idx_action]
+                action_fps = _get_optional_fps(gen_data_clean.fps_action, idx_action)
+                idx_action += 1
+
+                action_split_len = seq_builder.pack_action_tokens(
+                    input_action_tokens=input_action_tokens,
+                    condition_frame_indexes_action=sequence_plan.condition_frame_indexes_action,
+                    input_timestep=input_timestep,
+                    action_temporal_offset=vision_start_temporal_offset,
+                    enable_fps_modulation=enable_fps_modulation,
+                    base_fps=base_fps,
+                    action_fps=action_fps,
+                    base_temporal_compression_factor=temporal_compression_factor,
+                    action_start_frame_offset=sequence_plan.action_start_frame_offset,
+                )
+                sample_len += action_split_len
+            else:
+                action_split_len = 0
+
+        # Pack sound tokens if has_sound=True
+        if sequence_plan.has_sound:
+            input_sound_tokens = gen_data_clean.x0_tokens_sound[idx_sound]
+            sound_fps = _get_optional_fps(gen_data_clean.fps_sound, idx_sound)
+            idx_sound += 1
+
+            sound_split_len = seq_builder.pack_sound_tokens(
+                input_sound_tokens=input_sound_tokens,
+                condition_frame_indexes_sound=sequence_plan.condition_frame_indexes_sound,
+                input_timestep=input_timestep,
+                sound_temporal_offset=vision_start_temporal_offset,
+                enable_fps_modulation=enable_fps_modulation,
+                base_fps=base_fps,
+                sound_fps=sound_fps,
+                sound_base_temporal_compression_factor=sound_base_temporal_compression_factor,
+            )
+            sample_len += sound_split_len
+        else:
+            sound_split_len = 0
+
+        # Add end-of-generation token if needed
+        eov_len = 0
+        has_any_generation = sequence_plan.has_vision or sequence_plan.has_action or sequence_plan.has_sound
+        if include_end_of_generation_token and has_any_generation:
+            eov_len = seq_builder.append_end_of_generation_token(
+                token_id=special_tokens["end_of_generation"],
+                use_float_mrope_positions=use_float_mrope_positions,
+            )
+            sample_len += eov_len
+
+        combined_split_len = vision_split_len + action_split_len + sound_split_len + eov_len
+        seq_builder.finish_sample(combined_split_len, sample_len)
+
+    # Assert consistent null_action_supertokens across all TC samples, then set once
+    if null_action_flags:
+        assert len(set(null_action_flags)) == 1, (
+            f"Inconsistent null_action_supertokens across samples: {null_action_flags}. "
+            "All samples in a batch must have the same structure (all training or all AR inference)."
+        )
+        seq_builder.null_action_supertokens = null_action_flags[0]
+
+    # Finalize and return packed data
+    return seq_builder.finalize(
+        gen_data_clean=gen_data_clean,
+    )

releases/bridge-native-v3/code/cosmos_framework/data/generator/sequence_packing/sequence.py

@@ -0,0 +1,1203 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""Sequence builder/output and plan helpers for VFM sequence packing."""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+import torch
+
+from cosmos_framework.data.generator.sequence_packing.modality import ModalityData, ModalityDataBuilder, ModalitySpan
+from cosmos_framework.data.generator.sequence_packing.mrope import (
+    get_3d_mrope_ids_text_tokens,
+    get_3d_mrope_ids_vae_tokens,
+)
+from cosmos_framework.data.generator.sequence_packing.runtime import (
+    SequencePackMetadata,
+    prepare_sequence_pack_metadata,
+)
+
+if TYPE_CHECKING:
+    from cosmos_framework.model.generator.utils.data_and_condition import GenerationDataClean
+
+
+def _empty_long_tensor() -> torch.Tensor:
+    return torch.empty(0, dtype=torch.long)  # [0]
+
+
+@dataclass
+class PackedSequenceBuilder:
+    """Mutable construction state for sequence packing.
+
+    Attributes:
+        sample_lens: Length of each sample in the packed sequence.
+        split_lens: Length of each attention split. Each sample contributes a causal
+            text split and a full generation split.
+        attn_modes: Attention mode for each split, such as ``"causal"`` or ``"full"``.
+        is_image_batch: Whether this batch contains images rather than videos.
+        uses_single_timestep: Whether all noised tokens share one input timestep scalar.
+        sequence_length: Total packed sequence length. Filled during finalization.
+        current_seq_index: Next global packed-sequence index to append.
+        text_ids: Text token IDs accumulated during packing, including special tokens.
+        text_indexes: Global packed-sequence indexes for text tokens.
+        position_ids: mRoPE position ID blocks accumulated as ``[3, N]`` tensors.
+        label_ids: Label IDs for text cross-entropy loss.
+        ce_loss_indexes: Global packed-sequence indexes for text cross-entropy loss.
+        ce_loss_weights: Per-token weights for text cross-entropy loss.
+        _mrope_temporal_offset: Running temporal cursor for per-sample mRoPE generation.
+        _mrope_reset_spatial: Whether spatial mRoPE IDs reset for each segment.
+        null_action_supertokens: Whether temporal-causal supertoken 0 contains null
+            action tokens.
+        num_action_tokens_per_supertoken: Number of action tokens prefixing each
+            temporal-causal vision supertoken.
+        vision: Vision modality construction state, or ``None`` if no vision was appended.
+        action: Action modality construction state, or ``None`` if no action was appended.
+        sound: Sound modality construction state, or ``None`` if no sound was appended.
+        vision_item_split_lens: Per-sample per-vision-item token counts for multi-control
+            transfer.
+        control_weights: Per-sample per-control weights for multi-control weighted V-scaling.
+    """
+
+    # Sequence structure
+    sample_lens: list[int] = field(default_factory=list)
+    split_lens: list[int] = field(default_factory=list)
+    attn_modes: list[str] = field(default_factory=list)
+    is_image_batch: bool = False
+    uses_single_timestep: bool = False
+    sequence_length: int = 0
+
+    # Build-time tracking (used during packing, not after finalize)
+    current_seq_index: int = 0
+
+    # Text modality append state
+    text_ids: list[int] = field(default_factory=list)
+    text_indexes: list[int] = field(default_factory=list)
+    # Learned non-text condition slots that share the causal/understanding
+    # attention split with text. Zeva reserves one such slot per action-bearing
+    # sample and overwrites it with the projected task context.
+    behavior_indexes: list[int] = field(default_factory=list)
+    online_memory_indexes: list[int] = field(default_factory=list)
+    proprio_indexes: list[int] = field(default_factory=list)
+    position_ids: list[torch.Tensor] = field(default_factory=list)
+
+    # Loss computation - Cross Entropy (text)
+    label_ids: list[int] = field(default_factory=list)
+    ce_loss_indexes: list[int] = field(default_factory=list)
+    ce_loss_weights: list[float] = field(default_factory=list)
+
+    # Build-time mRoPE tracking (used during packing, not after finalize).
+    # position_ids accumulates (3, N) tensors and finalize() produces a
+    # (3, total_seq_len) tensor.
+    # Running temporal index for mRoPE position ID generation within a single sample.
+    # Reset to 0 at the start of each sample, then advanced by text and vision helpers
+    # as segments are packed. Action reuses the pre-vision snapshot (parallel temporal
+    # range) without advancing it. Float when FPS modulation is enabled.
+    # E.g. offset=0 -> text(4 tokens) -> offset=4 -> vision(3 frames) -> offset=7.
+    _mrope_temporal_offset: int | float = 0
+    _mrope_reset_spatial: bool = True
+
+    # Temporal causal: whether supertoken 0's action slot contains null tokens.
+    # True for all training calls and AR frame 0; False for AR frame N>0 (real actions).
+    # Used by three_way_attention to zero out V for null action tokens (inline when attention_meta.null_action_supertokens=True).
+    null_action_supertokens: bool = False
+
+    # Temporal causal: number of action tokens prefixing each vision supertoken.
+    # Equals temporal_compression_factor when actions are packed inline; 0 when
+    # action_gen=False or for non-temporal-causal layouts. Single source of truth
+    # for downstream attention/KV-cache code (per-supertoken layout is
+    # num_action_tokens_per_supertoken + H_p * W_p).
+    num_action_tokens_per_supertoken: int = 0
+
+    # Generation modality construction state
+    vision: ModalityDataBuilder | None = None
+    action: ModalityDataBuilder | None = None
+    sound: ModalityDataBuilder | None = None
+
+    def ensure_vision(self) -> ModalityDataBuilder:
+        """Return the vision builder, creating it on first use.
+
+        Returns:
+            Vision ``ModalityDataBuilder`` for subsequent append operations.
+        """
+        if self.vision is None:
+            self.vision = ModalityDataBuilder()
+        return self.vision
+
+    def ensure_action(self) -> ModalityDataBuilder:
+        """Return the action builder, creating it on first use.
+
+        Returns:
+            Action ``ModalityDataBuilder`` for subsequent append operations.
+        """
+        if self.action is None:
+            self.action = ModalityDataBuilder()
+        return self.action
+
+    def ensure_sound(self) -> ModalityDataBuilder:
+        """Return the sound builder, creating it on first use.
+
+        Returns:
+            Sound ``ModalityDataBuilder`` for subsequent append operations.
+        """
+        if self.sound is None:
+            self.sound = ModalityDataBuilder()
+        return self.sound
+
+    # Multi-control transfer: per-sample list of per-vision-item token counts.
+    # For a multi-control transfer sample with N controls + 1 noisy target,
+    # vision_item_split_lens[i] = [L_ctrl0, L_ctrl1, ..., L_ctrlN-1, L_noisy].
+    # Used by cosmos3_vfm_network.py to derive gen-relative control/noisy ranges
+    # for multi_control_two_way_attention.
+    vision_item_split_lens: list[list[int]] = field(default_factory=list)
+
+    # Per-sample per-control weights for multi-control weighted V-scaling.
+    # Parallel to vision_item_split_lens[i][:-1] (excludes noisy item).
+    # None for non-transfer or standard single-control samples.
+    control_weights: list[list[float]] | None = None
+
+    @property
+    def mrope_temporal_offset(self) -> int | float:
+        """Current per-sample temporal cursor used for mRoPE position generation."""
+        return self._mrope_temporal_offset
+
+    def set_mrope_temporal_offset(self, temporal_offset: int | float) -> None:
+        """Set the per-sample temporal cursor used for mRoPE position generation.
+
+        Args:
+            temporal_offset: New mRoPE temporal cursor value.
+        """
+        self._mrope_temporal_offset = temporal_offset
+
+    def advance_mrope_temporal_offset(self, delta: int | float) -> None:
+        """Advance the per-sample temporal cursor by ``delta``.
+
+        Args:
+            delta: Amount to add to the current mRoPE temporal cursor.
+        """
+        self._mrope_temporal_offset += delta
+
+    def begin_sample(self, initial_mrope_temporal_offset: int | float) -> None:
+        """Reset per-sample position cursors before appending sample tokens.
+
+        Args:
+            initial_mrope_temporal_offset: Initial mRoPE temporal cursor for this sample.
+        """
+        self._mrope_temporal_offset = initial_mrope_temporal_offset
+
+    def pack_text_tokens(
+        self,
+        text_ids: list[int],
+        special_tokens: dict[str, int],
+        has_generation: bool,
+        use_float_positions: bool = False,
+        num_behavior_prefix_tokens: int = 0,
+        num_online_memory_prefix_tokens: int = 0,
+        num_proprio_prefix_tokens: int = 0,
+    ) -> int:
+        """Pack text tokens into the sequence.
+
+        Args:
+            text_ids: List of text token IDs (integers).
+            special_tokens: Dictionary of special token IDs.
+            has_generation: Whether there's media/action after text.
+            use_float_positions: If True, generate float position IDs for 3D mRoPE
+                (for consistency with FPS-modulated vision tokens).
+
+        Returns:
+            Text sample length.
+        """
+        if num_behavior_prefix_tokens < 0:
+            raise ValueError("num_behavior_prefix_tokens must be non-negative")
+        if num_online_memory_prefix_tokens < 0:
+            raise ValueError("num_online_memory_prefix_tokens must be non-negative")
+        if num_proprio_prefix_tokens < 0:
+            raise ValueError("num_proprio_prefix_tokens must be non-negative")
+
+        # Reserve learned causal-prompt slots before BOS, matching Zeva's
+        # [causal prompt | image | language] prefix order. They are deliberately not
+        # represented by tokenizer IDs and therefore do not participate in CE.
+        if num_behavior_prefix_tokens:
+            behavior_position_ids, self._mrope_temporal_offset = get_3d_mrope_ids_text_tokens(
+                num_tokens=num_behavior_prefix_tokens,
+                temporal_offset=self._mrope_temporal_offset,
+                use_float_positions=use_float_positions,
+            )
+            behavior_start = self.current_seq_index
+            self.behavior_indexes.extend(
+                range(behavior_start, behavior_start + num_behavior_prefix_tokens)
+            )
+            self._append_position_ids(behavior_position_ids, num_behavior_prefix_tokens)
+            self.current_seq_index += num_behavior_prefix_tokens
+
+        if num_online_memory_prefix_tokens:
+            online_position_ids, self._mrope_temporal_offset = get_3d_mrope_ids_text_tokens(
+                num_tokens=num_online_memory_prefix_tokens,
+                temporal_offset=self._mrope_temporal_offset,
+                use_float_positions=use_float_positions,
+            )
+            online_start = self.current_seq_index
+            self.online_memory_indexes.extend(
+                range(online_start, online_start + num_online_memory_prefix_tokens)
+            )
+            self._append_position_ids(online_position_ids, num_online_memory_prefix_tokens)
+            self.current_seq_index += num_online_memory_prefix_tokens
+
+        if num_proprio_prefix_tokens:
+            proprio_position_ids, self._mrope_temporal_offset = get_3d_mrope_ids_text_tokens(
+                num_tokens=num_proprio_prefix_tokens,
+                temporal_offset=self._mrope_temporal_offset,
+                use_float_positions=use_float_positions,
+            )
+            proprio_start = self.current_seq_index
+            self.proprio_indexes.extend(range(proprio_start, proprio_start + num_proprio_prefix_tokens))
+            self._append_position_ids(proprio_position_ids, num_proprio_prefix_tokens)
+            self.current_seq_index += num_proprio_prefix_tokens
+
+        # Prepend BOS token if available
+        if "bos_token_id" in special_tokens:
+            shifted_text_ids = [special_tokens["bos_token_id"]] + text_ids
+        else:
+            shifted_text_ids = text_ids
+
+        span_token_ids = shifted_text_ids + [special_tokens["eos_token_id"]]
+        # Add start-of-generation token, but only if there's media/action present.
+        if has_generation:
+            span_token_ids.append(special_tokens["start_of_generation"])
+
+        text_split_len = len(span_token_ids)
+        split_len = (
+            num_behavior_prefix_tokens
+            + num_online_memory_prefix_tokens
+            + num_proprio_prefix_tokens
+            + text_split_len
+        )
+        expected_split_len = len(text_ids)
+        if "bos_token_id" in special_tokens:
+            expected_split_len += 1
+        expected_split_len += 1  # EOS
+        if has_generation:
+            expected_split_len += 1  # start-of-generation / BOV
+        assert text_split_len == expected_split_len
+
+        position_ids, self._mrope_temporal_offset = get_3d_mrope_ids_text_tokens(
+            num_tokens=text_split_len,
+            temporal_offset=self._mrope_temporal_offset,
+            use_float_positions=use_float_positions,
+        )  # position_ids: [3,text_split_len]
+
+        span_start, _ = self.append_text_span(span_token_ids, position_ids)
+
+        # Configure loss computation for text tokens
+        self.ce_loss_indexes.extend(range(span_start, span_start + len(shifted_text_ids)))
+        self.ce_loss_weights.extend([1.0] * len(shifted_text_ids))
+        self.label_ids.extend(text_ids[1:] + [special_tokens["eos_token_id"]])
+
+        self.attn_modes.append("causal")
+        self.split_lens.append(split_len)
+
+        return split_len
+
+    def pack_vision_tokens(
+        self,
+        input_vision_tokens: torch.Tensor,
+        condition_frame_indexes_vision: list[int],
+        input_timestep: float | torch.Tensor,
+        latent_patch_size: int = 1,
+        vision_fps: float | None = None,
+        enable_fps_modulation: bool = False,
+        base_fps: float = 24.0,
+        temporal_compression_factor: int = 4,
+        vision_temporal_positions: torch.Tensor | None = None,
+    ) -> int:
+        """Pack vision tokens into the sequence.
+
+        Args:
+            input_vision_tokens: Vision latent tokens (C, T, H, W).
+            condition_frame_indexes_vision: Indexes of conditioning frames.
+            input_timestep: Diffusion timestep. Either a float (teacher_forcing/none — all frames
+                share the same sigma) or a Tensor(T_max,) (diffusion_forcing — per-frame sigma;
+                indexed as input_timestep[frame_idx] for each noisy frame).
+            latent_patch_size: Patch size for latent patchification.
+            vision_fps: Frames per second of the video. Used when enable_fps_modulation=True.
+            enable_fps_modulation: If True, scale temporal position IDs based on video FPS.
+            base_fps: Base FPS for normalization (default 24.0).
+            temporal_compression_factor: VAE temporal compression factor (default 4).
+            vision_temporal_positions: Optional explicit temporal coordinate per latent
+                frame, shape ``(T,)``. Used by UniAE to account for kept boundary latents.
+
+        Returns:
+            Vision split length.
+        """
+        vision = self.ensure_vision()
+
+        # Compute position IDs for image patches
+        _, _, latent_t, latent_h, latent_w = input_vision_tokens.shape
+        if latent_patch_size < 1:
+            raise ValueError(f"latent_patch_size must be >= 1, got {latent_patch_size}")
+        # Use ceil to support latent dims not divisible by patch size (padding handled in network)
+        patch_h = math.ceil(latent_h / latent_patch_size)
+        patch_w = math.ceil(latent_w / latent_patch_size)
+        vision.token_shapes.append((latent_t, patch_h, patch_w))
+        vision.tokens.append(input_vision_tokens)
+        vision_payload_index = len(vision.tokens) - 1
+
+        # Supervise vision tokens based on conditioning frames
+        condition_set = {idx for idx in condition_frame_indexes_vision if 0 <= idx < latent_t}
+
+        vision_condition_mask = torch.zeros(
+            (latent_t, 1, 1), device=input_vision_tokens.device, dtype=input_vision_tokens.dtype
+        )  # [T,1,1]
+        for frame_idx in condition_set:
+            vision_condition_mask[frame_idx, 0, 0] = 1.0
+        vision.condition_mask.append(vision_condition_mask)
+
+        vision_noisy_frame_indexes = torch.tensor(
+            [idx for idx in range(latent_t) if idx not in condition_set],
+            device=input_vision_tokens.device,
+            dtype=torch.long,
+        )  # [N_noisy_frames]
+        vision.noisy_frame_indexes.append(vision_noisy_frame_indexes)
+
+        frame_token_stride = patch_h * patch_w
+        effective_fps = vision_fps if enable_fps_modulation else None
+        if vision_temporal_positions is not None:
+            vision_temporal_positions = vision_temporal_positions.to(device="cpu", dtype=torch.float32)  # [T]
+
+        vision_mrope_ids, self._mrope_temporal_offset = get_3d_mrope_ids_vae_tokens(
+            grid_t=latent_t,
+            grid_h=patch_h,
+            grid_w=patch_w,
+            temporal_offset=self._mrope_temporal_offset,
+            reset_spatial_indices=self._mrope_reset_spatial,
+            fps=effective_fps,
+            base_fps=base_fps,
+            temporal_compression_factor=temporal_compression_factor,
+            temporal_positions=vision_temporal_positions,
+            actual_temporal_compression_factor=temporal_compression_factor,
+        )  # vision_mrope_ids: [3,N_vision_tokens]
+        vision_mrope_ids = vision_mrope_ids.reshape(3, latent_t, frame_token_stride)  # [3,T,H*W]
+
+        vision_split_len = 0
+        for frame_idx in range(latent_t):
+            position_ids = vision_mrope_ids[:, frame_idx, :]  # [3,H*W]
+            frame_indexes = self.append_vision_span(
+                frame_token_stride,
+                position_ids,
+                payload_index=vision_payload_index,
+                payload_start=frame_idx * frame_token_stride,
+                payload_shape=(1, patch_h, patch_w),
+            )
+            vision_split_len += frame_token_stride
+
+            if frame_idx in condition_set:
+                continue
+            vision.mse_loss_indexes.extend(frame_indexes)
+            if isinstance(input_timestep, torch.Tensor):
+                frame_ts = input_timestep[frame_idx].item()
+            else:
+                frame_ts = input_timestep
+            vision.timesteps.extend([frame_ts] * frame_token_stride)
+
+        return vision_split_len
+
+    def pack_action_tokens(
+        self,
+        input_action_tokens: torch.Tensor,
+        condition_frame_indexes_action: list[int],
+        input_timestep: float,
+        action_temporal_offset: int | float = 0,
+        enable_fps_modulation: bool = False,
+        base_fps: float = 24.0,
+        action_fps: float | None = None,
+        base_temporal_compression_factor: int | None = None,
+        action_start_frame_offset: int = 1,
+    ) -> int:
+        """Pack action tokens into the sequence.
+
+        Args:
+            input_action_tokens: Action latent tokens (T, D).
+            condition_frame_indexes_action: Indexes of conditioning action steps.
+            input_timestep: Diffusion timestep.
+            action_temporal_offset: Temporal offset for action mRoPE IDs (typically
+                the vision start offset so action aligns temporally with vision).
+            enable_fps_modulation: If True, scale temporal position IDs based on FPS.
+            base_fps: Base FPS for normalization (default 24.0).
+            action_fps: Frames per second of the action data. Used when enable_fps_modulation=True.
+            base_temporal_compression_factor: Base temporal compression factor for FPS scaling.
+                Should be set to the vision temporal compression factor (e.g. 4) so that action
+                tokens advance at frame rate (4x finer) relative to vision latent frames.
+                Only affects behavior when FPS modulation is enabled.
+            action_start_frame_offset: Frame offset for aligning action[0] with the
+                corresponding vision frame. Default 1 aligns action[0] with vision frame 1.
+
+        Returns:
+            Number of action tokens added.
+        """
+        action_split_len = input_action_tokens.shape[0]
+
+        action = self.ensure_action()
+
+        # Add token indexes and loss information
+        action.token_shapes.append((action_split_len,))
+        action.tokens.append(input_action_tokens)
+        action_payload_index = len(action.tokens) - 1
+
+        condition_set = {idx for idx in condition_frame_indexes_action if 0 <= idx < action_split_len}
+
+        action_condition_mask = torch.zeros(
+            (action_split_len, 1), device=input_action_tokens.device, dtype=input_action_tokens.dtype
+        )  # [T_action,1]
+        for frame_idx in condition_set:
+            action_condition_mask[frame_idx, 0] = 1.0
+        action.condition_mask.append(action_condition_mask)
+
+        action_noisy_frame_indexes = torch.tensor(
+            [idx for idx in range(action_split_len) if idx not in condition_set],
+            device=input_action_tokens.device,
+            dtype=torch.long,
+        )  # [N_noisy_action_frames]
+        action.noisy_frame_indexes.append(action_noisy_frame_indexes)
+
+        # Action tokens use a 1x1 spatial grid with start_frame_offset=1 by default,
+        # so action[0] (null token) aligns with vision frame 1, not frame 0.
+        effective_fps = action_fps if enable_fps_modulation else None
+
+        action_mrope_ids, _ = get_3d_mrope_ids_vae_tokens(
+            grid_t=action_split_len,
+            grid_h=1,
+            grid_w=1,
+            temporal_offset=action_temporal_offset,
+            reset_spatial_indices=self._mrope_reset_spatial,
+            fps=effective_fps,
+            base_fps=base_fps,
+            temporal_compression_factor=1,  # Action is at frame rate (no temporal compression)
+            base_temporal_compression_factor=base_temporal_compression_factor,
+            start_frame_offset=action_start_frame_offset,  # Align action[0] with vision frame action_start_frame_offset
+        )  # action_mrope_ids: [3,N_action_tokens]
+        # Note: we don't update _mrope_temporal_offset here because action tokens
+        # share the temporal space with vision tokens (they run in parallel).
+
+        for frame_idx in range(action_split_len):
+            position_ids = action_mrope_ids[:, frame_idx : frame_idx + 1]  # [3,1]
+            frame_indexes = self.append_action_span(
+                1,
+                position_ids,
+                payload_index=action_payload_index,
+                payload_start=frame_idx,
+                payload_shape=(1,),
+            )
+
+            if frame_idx in condition_set:
+                continue
+            action.mse_loss_indexes.extend(frame_indexes)
+            action.timesteps.extend([input_timestep])
+
+        return action_split_len
+
+    def pack_sound_tokens(
+        self,
+        input_sound_tokens: torch.Tensor,
+        condition_frame_indexes_sound: list[int],
+        input_timestep: float,
+        sound_temporal_offset: int | float = 0,
+        enable_fps_modulation: bool = False,
+        base_fps: float = 24.0,
+        sound_fps: float | None = None,
+        sound_base_temporal_compression_factor: int | None = None,
+    ) -> int:
+        """Pack sound/audio tokens into the sequence.
+
+        Sound latents have shape [C, T] where C is channels and T is temporal frames.
+        Sound tokens are added to the unified generation split to maintain SequencePack's
+        2-split invariant (causal + full).
+
+        Args:
+            input_sound_tokens: Sound latent tokens (C, T).
+            condition_frame_indexes_sound: Indexes of conditioning frames.
+                [] means all frames are noised/supervised.
+                All frames specified means all frames are clean (no MSE supervision).
+            input_timestep: Diffusion timestep.
+            sound_temporal_offset: Temporal offset for m-RoPE position IDs (aligned with vision start).
+            enable_fps_modulation: If True, scale temporal positions by FPS ratio.
+            base_fps: Base FPS for normalization (default 24.0).
+            sound_fps: Sound latent FPS (e.g., 25.0). Used for FPS-aware m-RoPE positions.
+            sound_base_temporal_compression_factor: Base temporal compression factor for sound FPS scaling.
+                ``None`` preserves the current behavior where sound advances at ``base_fps`` positions/sec.
+
+        Returns:
+            Number of sound tokens added.
+        """
+        # Sound latent shape: [C, T] → T tokens
+        _, sound_split_len = input_sound_tokens.shape
+
+        sound = self.ensure_sound()
+
+        # Add token indexes - sound uses (T, 1, 1) shape for compatibility with 3D RoPE
+        sound.token_shapes.append((sound_split_len, 1, 1))
+        sound.tokens.append(input_sound_tokens)
+        sound_payload_index = len(sound.tokens) - 1
+
+        # Supervise sound tokens based on conditioning frames
+        condition_set = {idx for idx in condition_frame_indexes_sound if 0 <= idx < sound_split_len}
+
+        # Condition mask: shape (T, 1) — 1 = clean/conditioning, 0 = noised/supervised
+        sound_condition_mask = torch.zeros(
+            (sound_split_len, 1), device=input_sound_tokens.device, dtype=input_sound_tokens.dtype
+        )  # [T_sound,1]
+        for frame_idx in condition_set:
+            sound_condition_mask[frame_idx, 0] = 1.0
+        sound.condition_mask.append(sound_condition_mask)
+
+        sound_noisy_frame_indexes = torch.tensor(
+            [idx for idx in range(sound_split_len) if idx not in condition_set],
+            device=input_sound_tokens.device,
+            dtype=torch.long,
+        )  # [N_noisy_sound_frames]
+        sound.noisy_frame_indexes.append(sound_noisy_frame_indexes)
+
+        # Sound tokens use a 1x1 spatial grid, aligned with vision temporal positions.
+        # sound[0] aligns with vision frame 0 (start_frame_offset=0, unlike action which offsets by 1).
+        effective_fps = sound_fps if enable_fps_modulation else None
+
+        sound_mrope_ids, _ = get_3d_mrope_ids_vae_tokens(
+            grid_t=sound_split_len,
+            grid_h=1,
+            grid_w=1,
+            temporal_offset=sound_temporal_offset,
+            reset_spatial_indices=self._mrope_reset_spatial,
+            fps=effective_fps,
+            base_fps=base_fps,
+            temporal_compression_factor=1,  # Sound latent is already at sound_latent_fps (no further compression)
+            base_temporal_compression_factor=sound_base_temporal_compression_factor,
+            start_frame_offset=0,  # Sound[0] aligns with vision frame 0
+        )  # sound_mrope_ids: [3,N_sound_tokens]
+        # Note: we don't update _mrope_temporal_offset here because sound tokens
+        # share the temporal space with vision tokens (they run in parallel).
+
+        # Add to MSE loss indexes and timesteps for non-conditioning frames
+        for frame_idx in range(sound_split_len):
+            position_ids = sound_mrope_ids[:, frame_idx : frame_idx + 1]  # [3,1]
+            frame_indexes = self.append_sound_span(
+                1,
+                position_ids,
+                payload_index=sound_payload_index,
+                payload_start=frame_idx,
+                payload_shape=(1, 1, 1),
+            )
+
+            if frame_idx in condition_set:
+                continue
+            sound.mse_loss_indexes.extend(frame_indexes)
+            sound.timesteps.extend([input_timestep])
+
+        return sound_split_len
+
+    def _append_position_ids(self, position_ids: torch.Tensor, span_len: int) -> None:
+        """Append one block of position IDs.
+
+        Args:
+            position_ids: Position ID tensor with shape ``[3, span_len]``.
+            span_len: Number of token positions described by ``position_ids``.
+        """
+        assert position_ids.shape[-1] == span_len, (
+            f"position_ids last dimension must match span_len={span_len}; got {tuple(position_ids.shape)}"
+        )
+        self.position_ids.append(position_ids)
+
+    def _append_modality_span(
+        self,
+        modality: ModalityDataBuilder,
+        span_len: int,
+        position_ids: torch.Tensor,
+        *,
+        payload_index: int,
+        payload_start: int,
+        payload_shape: tuple[int, ...],
+    ) -> list[int]:
+        """Append one modality span and return its global sequence indexes.
+
+        Args:
+            modality: Modality builder receiving sequence indexes and span metadata.
+            span_len: Number of contiguous tokens to append.
+            position_ids: Position ID tensor with shape ``[3, span_len]``.
+            payload_index: Index into ``modality.tokens`` for the backing payload tensor.
+            payload_start: First flattened token offset within the backing payload tensor.
+            payload_shape: Logical payload slice shape represented by this span.
+
+        Returns:
+            Global packed-sequence indexes covered by the appended span.
+        """
+        assert payload_index >= 0, f"payload_index must be non-negative, got {payload_index}"
+        assert payload_start >= 0, f"payload_start must be non-negative, got {payload_start}"
+        payload_len = 1
+        for dim in payload_shape:
+            assert dim > 0, f"payload_shape dimensions must be positive, got {payload_shape}"
+            payload_len *= dim
+        assert payload_len == span_len, (
+            f"payload_shape={payload_shape} describes {payload_len} tokens, but span_len={span_len}"
+        )
+
+        span_start = self.current_seq_index
+        span_indexes = list(range(span_start, span_start + span_len))
+        modality.spans.append(
+            ModalitySpan(
+                sequence_start=span_start,
+                sequence_len=span_len,
+                payload_index=payload_index,
+                payload_start=payload_start,
+                payload_len=payload_len,
+                payload_shape=payload_shape,
+            )
+        )
+        modality.sequence_indexes.extend(span_indexes)
+        self._append_position_ids(position_ids, span_len)
+        self.current_seq_index += span_len
+        return span_indexes
+
+    def append_text_span(
+        self,
+        token_ids: list[int],
+        position_ids: torch.Tensor,
+    ) -> tuple[int, int]:
+        """Append one contiguous text span.
+
+        Args:
+            token_ids: Text token IDs to append.
+            position_ids: Position ID tensor with shape ``[3, len(token_ids)]``.
+
+        Returns:
+            Tuple of ``(start_index, span_len)`` for the appended span.
+        """
+        span_start = self.current_seq_index
+        span_len = len(token_ids)
+        self.text_ids.extend(token_ids)
+        self.text_indexes.extend(range(span_start, span_start + span_len))
+        self._append_position_ids(position_ids, span_len)
+        self.current_seq_index += span_len
+        return span_start, span_len
+
+    def append_vision_span(
+        self,
+        span_len: int,
+        position_ids: torch.Tensor,
+        *,
+        payload_index: int,
+        payload_start: int,
+        payload_shape: tuple[int, int, int],
+    ) -> list[int]:
+        """Append one contiguous vision span.
+
+        Args:
+            span_len: Number of contiguous vision tokens to append.
+            position_ids: Position ID tensor with shape ``[3, span_len]``.
+            payload_index: Index into ``vision.tokens`` for the backing payload tensor.
+            payload_start: First flattened token offset within the backing vision payload.
+            payload_shape: Logical vision slice shape, usually ``(1, patch_h, patch_w)``.
+
+        Returns:
+            Global packed-sequence indexes covered by the appended vision span.
+        """
+        return self._append_modality_span(
+            self.ensure_vision(),
+            span_len,
+            position_ids,
+            payload_index=payload_index,
+            payload_start=payload_start,
+            payload_shape=payload_shape,
+        )
+
+    def append_action_span(
+        self,
+        span_len: int,
+        position_ids: torch.Tensor,
+        *,
+        payload_index: int,
+        payload_start: int,
+        payload_shape: tuple[int],
+    ) -> list[int]:
+        """Append one contiguous action span.
+
+        Args:
+            span_len: Number of contiguous action tokens to append.
+            position_ids: Position ID tensor with shape ``[3, span_len]``.
+            payload_index: Index into ``action.tokens`` for the backing payload tensor.
+            payload_start: First flattened token offset within the backing action payload.
+            payload_shape: Logical action slice shape, usually ``(span_len,)``.
+
+        Returns:
+            Global packed-sequence indexes covered by the appended action span.
+        """
+        return self._append_modality_span(
+            self.ensure_action(),
+            span_len,
+            position_ids,
+            payload_index=payload_index,
+            payload_start=payload_start,
+            payload_shape=payload_shape,
+        )
+
+    def append_sound_span(
+        self,
+        span_len: int,
+        position_ids: torch.Tensor,
+        *,
+        payload_index: int,
+        payload_start: int,
+        payload_shape: tuple[int, int, int],
+    ) -> list[int]:
+        """Append one contiguous sound span.
+
+        Args:
+            span_len: Number of contiguous sound tokens to append.
+            position_ids: Position ID tensor with shape ``[3, span_len]``.
+            payload_index: Index into ``sound.tokens`` for the backing payload tensor.
+            payload_start: First flattened token offset within the backing sound payload.
+            payload_shape: Logical sound slice shape, usually ``(1, 1, 1)``.
+
+        Returns:
+            Global packed-sequence indexes covered by the appended sound span.
+        """
+        return self._append_modality_span(
+            self.ensure_sound(),
+            span_len,
+            position_ids,
+            payload_index=payload_index,
+            payload_start=payload_start,
+            payload_shape=payload_shape,
+        )
+
+    def append_end_of_generation_token(
+        self,
+        token_id: int,
+        use_float_mrope_positions: bool,
+    ) -> int:
+        """Append an end-of-generation text token.
+
+        Args:
+            token_id: Token ID for the end-of-generation marker.
+            use_float_mrope_positions: Whether to write float position IDs for this token.
+
+        Returns:
+            Span length for the appended end-of-generation token.
+        """
+        eov_dtype = torch.float32 if use_float_mrope_positions else torch.long
+        position_ids = torch.full((3, 1), self._mrope_temporal_offset, dtype=eov_dtype)  # [3,1]
+        self._mrope_temporal_offset += 1
+
+        _, span_len = self.append_text_span([token_id], position_ids)
+        return span_len
+
+    def finish_sample(self, generation_split_len: int, sample_len: int) -> None:
+        """Record the non-causal generation split and total sample length.
+
+        Args:
+            generation_split_len: Length of the full-attention generation split.
+            sample_len: Total number of tokens appended for this sample.
+        """
+        self.attn_modes.append("full")
+        self.split_lens.append(generation_split_len)
+        self.sample_lens.append(sample_len)
+
+    def _finalize_modality(
+        self,
+        modality: ModalityDataBuilder | None,
+        *,
+        domain_id: list[torch.Tensor] | None = None,
+        raw_action_dim: list[torch.Tensor | None] | None = None,
+        residual_gate: list[torch.Tensor] | None = None,
+        include_raw_action_dim: bool = False,
+    ) -> ModalityData | None:
+        """Finalize one modality builder into model-facing modality data.
+
+        Args:
+            modality: Modality builder to finalize, or ``None`` if no tokens were appended.
+            domain_id: Optional action domain IDs to attach to finalized action data.
+            raw_action_dim: Optional raw action dimensions for action-channel masking.
+            include_raw_action_dim: Whether to include ``raw_action_dim`` in the finalized data.
+
+        Returns:
+            Finalized ``ModalityData`` or ``None`` when the modality has no sequence indexes.
+        """
+        if modality is None or len(modality.sequence_indexes) == 0:
+            return None
+
+        kwargs = {
+            "sequence_indexes": torch.tensor(modality.sequence_indexes, dtype=torch.long),  # [N_modality_tokens]
+            "timesteps": torch.tensor(modality.timesteps, dtype=torch.float32),  # [N_modality_noisy_tokens]
+            "mse_loss_indexes": torch.tensor(modality.mse_loss_indexes, dtype=torch.long),  # [N_modality_noisy_tokens]
+            "spans": list(modality.spans),
+            "token_shapes": list(modality.token_shapes),
+            "tokens": modality.tokens,
+            "condition_mask": list(modality.condition_mask),
+            "noisy_frame_indexes": list(modality.noisy_frame_indexes),
+        }
+        if domain_id is not None:
+            kwargs["domain_id"] = domain_id
+        if include_raw_action_dim:
+            kwargs["raw_action_dim"] = raw_action_dim
+        if residual_gate is not None:
+            kwargs["residual_gate"] = residual_gate
+        return ModalityData(**kwargs)
+
+    def finalize(
+        self,
+        gen_data_clean: GenerationDataClean,
+    ) -> "PackedSequence":
+        """Convert all lists to tensors and compute derived values.
+
+        Args:
+            gen_data_clean: GenerationDataClean for metadata (e.g., action domain IDs).
+
+        Returns:
+            New PackedSequence instance with tensors instead of lists.
+        """
+        # Compute sequence length
+        sequence_length = sum(self.sample_lens)
+        sample_lens = self.sample_lens.copy()
+        split_lens = self.split_lens.copy()
+        attn_modes = self.attn_modes.copy()
+
+        # Prepare loss-related tensors (cross-entropy)
+        label_ids: torch.Tensor | None = None
+        ce_loss_indexes: torch.Tensor | None = None
+        ce_loss_weights: torch.Tensor | None = None
+        if self.label_ids and len(self.label_ids) > 0:
+            label_ids = torch.tensor(self.label_ids)  # [N_ce_tokens]
+            ce_loss_indexes = torch.tensor(self.ce_loss_indexes)  # [N_ce_tokens]
+            ce_loss_weights = torch.tensor(self.ce_loss_weights)  # [N_ce_tokens]
+
+        # The condition_mask and noisy_frame_indexes are kept as lists to support variable shapes.
+
+        vision = self._finalize_modality(self.vision)
+        action_domain_id = None
+        if self.action is not None:
+            if gen_data_clean.action_domain_id is not None:
+                action_domain_id = gen_data_clean.action_domain_id
+            else:
+                default_action_domain_id = torch.zeros(1, dtype=torch.long)  # [1]
+                action_domain_id = [default_action_domain_id] * len(self.action.token_shapes)
+        action = self._finalize_modality(
+            self.action,
+            domain_id=action_domain_id,
+            raw_action_dim=gen_data_clean.raw_action_dim,
+            residual_gate=gen_data_clean.action_residual_gate,
+            include_raw_action_dim=True,
+        )
+        sound = self._finalize_modality(self.sound)
+
+        # Finalize position IDs.
+        assert isinstance(self.position_ids, list)
+        assert all(isinstance(position_id, torch.Tensor) for position_id in self.position_ids)
+        if len(self.position_ids) > 0:
+            position_ids = torch.cat(self.position_ids, dim=1)  # [3,actual_seq_len]
+        else:
+            position_ids = torch.empty((3, 0), dtype=torch.long)  # [3,0]
+
+        return PackedSequence(
+            # Sequence structure
+            sequence_length=sequence_length,
+            sample_lens=sample_lens,
+            split_lens=split_lens,
+            attn_modes=attn_modes,
+            is_image_batch=gen_data_clean.is_image_batch,
+            uses_single_timestep=self.uses_single_timestep,
+            # Text modality (converted to tensors)
+            text_ids=torch.tensor(self.text_ids, dtype=torch.long),  # [N_text_tokens]
+            text_indexes=torch.tensor(self.text_indexes, dtype=torch.long),  # [N_text_tokens]
+            behavior_indexes=torch.tensor(self.behavior_indexes, dtype=torch.long),
+            online_memory_indexes=torch.tensor(self.online_memory_indexes, dtype=torch.long),
+            proprio_indexes=torch.tensor(self.proprio_indexes, dtype=torch.long),
+            position_ids=position_ids,  # [3,seq_len]
+            # Loss computation - Cross Entropy
+            label_ids=label_ids,
+            ce_loss_indexes=ce_loss_indexes,
+            ce_loss_weights=ce_loss_weights,
+            # Generation modalities
+            vision=vision,
+            action=action,
+            sound=sound,
+            # Temporal causal
+            null_action_supertokens=self.null_action_supertokens,
+            num_action_tokens_per_supertoken=self.num_action_tokens_per_supertoken,
+            # Multi-control transfer
+            vision_item_split_lens=list(self.vision_item_split_lens),
+            control_weights=gen_data_clean.control_weights,
+        )
+
+
+@dataclass
+class PackedSequence:
+    """Finalized model-facing packed sequence.
+
+    The object remains mutable for model-side payload replacement, clean replay,
+    and in-place device transfer. Build-only cursor and mRoPE state live on
+    ``PackedSequenceBuilder``.
+
+    Attributes:
+        sample_lens: Length of each sample in the packed sequence.
+        split_lens: Length of each attention split. Each sample contributes a causal
+            text split and a full generation split.
+        attn_modes: Attention mode for each split, such as ``"causal"`` or ``"full"``.
+        is_image_batch: Whether this batch contains images rather than videos.
+        uses_single_timestep: Whether all noised tokens share one input timestep scalar.
+        sequence_length: Total length of the packed sequence.
+        text_ids: Tensor of all text token IDs, including special tokens.
+        text_indexes: Tensor of global packed-sequence indexes for text tokens.
+        position_ids: Tensor of mRoPE position IDs for all packed tokens with shape
+            ``[3, sequence_length]``.
+        label_ids: Optional tensor of label IDs for text cross-entropy loss.
+        ce_loss_indexes: Optional tensor of global packed-sequence indexes for text
+            cross-entropy loss.
+        ce_loss_weights: Optional tensor of per-token weights for text cross-entropy loss.
+        null_action_supertokens: Whether temporal-causal supertoken 0 contains null
+            action tokens.
+        num_action_tokens_per_supertoken: Number of action tokens prefixing each
+            temporal-causal vision supertoken.
+        vision: Finalized vision modality data, or ``None`` if no vision is present.
+        action: Finalized action modality data, or ``None`` if no action is present.
+        sound: Finalized sound modality data, or ``None`` if no sound is present.
+        vision_item_split_lens: Per-sample per-vision-item token counts for multi-control
+            transfer.
+        control_weights: Per-sample per-control weights for multi-control weighted V-scaling.
+    """
+
+    # Sequence structure
+    sample_lens: list[int] = field(default_factory=list)
+    split_lens: list[int] = field(default_factory=list)
+    attn_modes: list[str] = field(default_factory=list)
+    is_image_batch: bool = False
+    uses_single_timestep: bool = False
+    sequence_length: int = 0
+
+    # Text modality
+    text_ids: torch.Tensor = field(default_factory=_empty_long_tensor)
+    text_indexes: torch.Tensor = field(default_factory=_empty_long_tensor)
+    behavior_indexes: torch.Tensor = field(default_factory=_empty_long_tensor)
+    online_memory_indexes: torch.Tensor = field(default_factory=_empty_long_tensor)
+    proprio_indexes: torch.Tensor = field(default_factory=_empty_long_tensor)
+    position_ids: torch.Tensor = field(default_factory=_empty_long_tensor)
+
+    # Loss computation - Cross Entropy (text)
+    label_ids: torch.Tensor | None = None
+    ce_loss_indexes: torch.Tensor | None = None
+    ce_loss_weights: torch.Tensor | None = None
+
+    # Temporal causal: whether supertoken 0's action slot contains null tokens.
+    # True for all training calls and AR frame 0; False for AR frame N>0 (real actions).
+    # Used by three_way_attention to zero out V for null action tokens (inline when attention_meta.null_action_supertokens=True).
+    null_action_supertokens: bool = False
+
+    # Temporal causal: number of action tokens prefixing each vision supertoken.
+    # Equals temporal_compression_factor when actions are packed inline; 0 when
+    # action_gen=False or for non-temporal-causal layouts. Single source of truth
+    # for downstream attention/KV-cache code (per-supertoken layout is
+    # num_action_tokens_per_supertoken + H_p * W_p).
+    num_action_tokens_per_supertoken: int = 0
+
+    # Generation modalities - NAMED FIELDS for type safety
+    vision: ModalityData | None = None
+    action: ModalityData | None = None
+    sound: ModalityData | None = None
+
+    # Multi-control transfer: per-sample list of per-vision-item token counts.
+    # For a multi-control transfer sample with N controls + 1 noisy target,
+    # vision_item_split_lens[i] = [L_ctrl0, L_ctrl1, ..., L_ctrlN-1, L_noisy].
+    # Used by cosmos3_vfm_network.py to derive gen-relative control/noisy ranges
+    # for multi_control_two_way_attention.
+    vision_item_split_lens: list[list[int]] = field(default_factory=list)
+
+    # Per-sample per-control weights for multi-control weighted V-scaling.
+    # Parallel to vision_item_split_lens[i][:-1] (excludes noisy item).
+    # None for non-transfer or standard single-control samples.
+    control_weights: list[list[float]] | None = None
+
+    def __post_init__(self) -> None:
+        self._sequence_pack_metadata: SequencePackMetadata | None = None
+        assert isinstance(self.text_ids, torch.Tensor), "PackedSequence.text_ids must be finalized"
+        assert isinstance(self.text_indexes, torch.Tensor), "PackedSequence.text_indexes must be finalized"
+        assert isinstance(self.behavior_indexes, torch.Tensor), "PackedSequence.behavior_indexes must be finalized"
+        assert isinstance(self.online_memory_indexes, torch.Tensor), "PackedSequence.online_memory_indexes must be finalized"
+        assert isinstance(self.proprio_indexes, torch.Tensor), "PackedSequence.proprio_indexes must be finalized"
+        assert isinstance(self.position_ids, torch.Tensor), "PackedSequence.position_ids must be finalized"
+        if self.label_ids is not None:
+            assert isinstance(self.label_ids, torch.Tensor), "PackedSequence.label_ids must be finalized"
+        if self.ce_loss_indexes is not None:
+            assert isinstance(self.ce_loss_indexes, torch.Tensor), "PackedSequence.ce_loss_indexes must be finalized"
+        if self.ce_loss_weights is not None:
+            assert isinstance(self.ce_loss_weights, torch.Tensor), "PackedSequence.ce_loss_weights must be finalized"
+        for modality in [self.vision, self.action, self.sound]:
+            assert modality is None or isinstance(modality, ModalityData), (
+                "PackedSequence modality fields must be finalized ModalityData"
+            )
+
+    def to_cuda(self) -> None:
+        """Move all tensor fields to CUDA in-place."""
+        self.text_ids = self.text_ids.cuda()
+        self.text_indexes = self.text_indexes.cuda()
+        self.behavior_indexes = self.behavior_indexes.cuda()
+        self.online_memory_indexes = self.online_memory_indexes.cuda()
+        self.proprio_indexes = self.proprio_indexes.cuda()
+        self.position_ids = self.position_ids.cuda()
+        if isinstance(self.label_ids, torch.Tensor):
+            self.label_ids = self.label_ids.cuda()
+        if isinstance(self.ce_loss_indexes, torch.Tensor):
+            self.ce_loss_indexes = self.ce_loss_indexes.cuda()
+        if isinstance(self.ce_loss_weights, torch.Tensor):
+            self.ce_loss_weights = self.ce_loss_weights.cuda()
+        if self.vision is not None:
+            self.vision.to_cuda()
+        if self.action is not None:
+            self.action.to_cuda()
+        if self.sound is not None:
+            self.sound.to_cuda()
+        self.prepare_sequence_pack_metadata()
+
+    def prepare_sequence_pack_metadata(self) -> None:
+        """Validate and prepare device-specific metadata for this layout."""
+        understanding_indexes = torch.cat(
+            (self.behavior_indexes, self.online_memory_indexes, self.proprio_indexes, self.text_indexes)
+        ).sort().values
+        self._sequence_pack_metadata = prepare_sequence_pack_metadata(
+            sample_lens=self.sample_lens,
+            split_lens=self.split_lens,
+            attn_modes=self.attn_modes,
+            packed_und_token_indexes=understanding_indexes,
+            device=self.text_indexes.device,
+        )
+
+    def get_sequence_pack_metadata(self) -> SequencePackMetadata | None:
+        """Return metadata prepared after the packed input reached its device."""
+        return self._sequence_pack_metadata
+
+
+@dataclass
+class SequencePlan:
+    """Plan describing which modalities are present in a sample.
+
+    This dataclass tracks the presence of different modalities (text, vision, action)
+    and their conditioning configurations for a dataset sample. Unlike SequencePlan
+    which holds the actual tensor data, this class provides a lightweight summary
+    of what modalities exist and how they should be conditioned.
+
+    Attributes:
+        has_text: Whether text/caption tokens are present for this sample.
+            Used for text-conditioned generation (e.g., text-to-image/video).
+        has_vision: Whether vision input (image or video latents) is present.
+            Defaults to False.
+        condition_frame_indexes_vision: Indexes of latent vision frames that are clean/conditioning.
+            [] means all frames are noised/supervised.
+            All frames specified means all frames are clean (no MSE supervision).
+            For multi-item samples (e.g. image editing where each sample has multiple
+            separately-encoded images), this applies to each vision item individually.
+            The number of items per sample is tracked by
+            ``GenerationDataClean.num_vision_items_per_sample``.
+        share_vision_temporal_positions: Whether all vision items in this sample share
+            the same temporal mRoPE grid.
+        has_action: Whether action input is present for robotics/embodied AI tasks.
+            Defaults to False.
+        condition_frame_indexes_action: Indexes of action steps that are clean/conditioning.
+            [] means all steps are noised/supervised.
+            All steps specified means all steps are clean (no MSE supervision).
+        action_start_frame_offset: Frame offset for aligning action[0] to vision frames.
+        has_sound: Whether sound/audio input is present.
+        condition_frame_indexes_sound: Indexes of sound frames that are clean/conditioning.
+            [] means all frames are noised/supervised.
+            All frames specified means all frames are clean (no MSE supervision).
+    """
+
+    # -- understanding (text conditioning) --
+    has_text: bool
+
+    # -- vision modality --
+    has_vision: bool = False
+    condition_frame_indexes_vision: list[int] = field(default_factory=list)
+    # If True, all vision items in this sample share the same temporal mRoPE grid
+    # (controlnet-style transfer: target frame i is spatio-temporally aligned with
+    # control frame i). Each item gets the same temporal_offset; spatial reset
+    # behavior is unchanged. Requires num_vision_items_per_sample > 1, equal latent_t,
+    # and equal fps across items. Default False preserves single-clip and
+    # image-editing semantics where items represent distinct time states.
+    share_vision_temporal_positions: bool = False
+
+    # -- action modality --
+    has_action: bool = False
+    condition_frame_indexes_action: list[int] = field(default_factory=list)
+    action_start_frame_offset: int = 1
+
+    # -- sound modality --
+    has_sound: bool = False
+    condition_frame_indexes_sound: list[int] = field(default_factory=list)
+
+    def as_dict(self) -> dict:
+        return {
+            "has_text": self.has_text,
+            "has_vision": self.has_vision,
+            "has_action": self.has_action,
+            "has_sound": self.has_sound,
+            "condition_frame_indexes_vision": self.condition_frame_indexes_vision,
+            "condition_frame_indexes_action": self.condition_frame_indexes_action,
+            "condition_frame_indexes_sound": self.condition_frame_indexes_sound,
+            "share_vision_temporal_positions": self.share_vision_temporal_positions,
+        }
+
+
+def build_sequence_plans_from_data_batch(
+    data_batch: dict,
+    input_video_key,
+    input_image_key: str,
+) -> list[SequencePlan]:
+    """Build or retrieve sequence plans from a data batch dictionary.
+
+    This function extracts sequence plans from the data batch if they exist,
+    otherwise creates default SequencePlan objects for each sample
+    in the batch.
+
+    Args:
+        data_batch: Dictionary containing the data batch from the dataloader.
+            Expected keys include 'video' or other tensors to determine batch size.
+            If 'sequence_plan' key exists, those plans are returned directly.
+        input_video_key: Data-batch key used to find video tensors when inferring batch size.
+        input_image_key: Data-batch key used to find image tensors when inferring batch size.
+
+    Returns:
+        List of SequencePlan objects, one per sample in the batch.
+    """
+    # For new modalities, please generate the sequence_plan in the dataset class!!!!
+
+    # If sequence_plan already exists in data_batch, return it
+    if "sequence_plan" in data_batch:
+        return data_batch["sequence_plan"]
+
+    assert "action" not in data_batch or data_batch["action"] is None, "Action data SHOULD have sequence_plans!"
+    assert "sound" not in data_batch or data_batch["sound"] is None, "Sound data SHOULD have sequence_plans!"
+
+    # Determine batch size from available tensors
+    batch_size = 0
+    for key in [input_video_key, input_image_key]:
+        if key in data_batch:
+            val = data_batch[key]
+            if isinstance(val, torch.Tensor):
+                batch_size = val.shape[0]
+                break
+            elif isinstance(val, list):
+                batch_size = len(val)
+                break
+
+    if batch_size == 0:
+        raise ValueError(
+            f"Cannot determine batch size from data_batch. Expected {input_video_key}, {input_image_key}, or similar key."
+        )
+
+    # Build default SequencePlan objects
+    return [
+        SequencePlan(
+            has_text=True,  # Has text prompt!
+            has_vision=True,
+            condition_frame_indexes_vision=[],  # No conditioning frames!
+        )
+        for _ in range(batch_size)
+    ]

releases/bridge-native-v3/code/cosmos_framework/inference/common/public_model_config.py

@@ -0,0 +1,361 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+from __future__ import annotations
+
+import copy
+import importlib.util
+from typing import Any
+
+_PUBLIC_TARGET_KEY = "_target"
+_PUBLIC_TYPE_KEY = "_type"
+_CLASS_NAME_KEY = "class_name"
+_CONFIG_NAME_KEY = "config_name"
+_TYPE_KEY = "_type"
+_TARGET_KEY = "_target_"
+
+_PUBLIC_VFM_URI_PREFIX = "cosmos3://vfm/"
+
+_TARGET_ALIASES = {
+    "projects.cosmos3.vfm.configs.base.defaults.vlm.create_qwen2_tokenizer_with_download": "create_qwen2_tokenizer_with_download",
+    "projects.cosmos3.vfm.configs.base.defaults.vlm.create_vlm_config": "create_vlm_config",
+    "projects.cosmos3.vfm.models.mot.unified_mot.Nemotron3DenseVLMoTConfig.from_json_file": "nemotron3_dense_vl_mot_config_from_json_file",
+    "projects.cosmos3.vfm.models.mot.unified_mot.Nemotron3DenseVLTextForCausalLM": "nemotron3_dense_vl_text_for_causal_lm",
+    "projects.cosmos3.vfm.models.mot.unified_mot.Qwen3VLMoTConfig.from_json_file": "qwen3_vl_mot_config_from_json_file",
+    "projects.cosmos3.vfm.models.mot.unified_mot.Qwen3VLTextForCausalLM": "qwen3_vl_text_for_causal_lm",
+    "projects.cosmos3.vfm.models.omni_mot_model.OmniMoTModel": "omni_mot_model",
+    "projects.cosmos3.vfm.processors.build_processor_lazy": "build_processor_lazy",
+    "projects.cosmos3.vfm.tokenizers.audio.avae.AVAEInterface": "avae_interface",
+    "projects.cosmos3.vfm.tokenizers.wan2pt2_vae_4x16x16.Wan2pt2VAEInterface": "wan2pt2_vae_interface",
+}
+
+# The Zeva training branch exposes config classes under the
+# cosmos_framework.configs.* package path (mirrors of projects.cosmos3.vfm.*).
+# Map them to the same public aliases so exported HF configs stay identical.
+_CF_TYPE_ALIASES = {
+    "cosmos_framework.configs.base.defaults.model_config.ZevaPolicyConfig": "zeva_policy_config",
+    "cosmos_framework.configs.base.defaults.model_config.BehaviorStage2Config": "behavior_stage2_config",
+    "cosmos_framework.configs.base.defaults.model_config.ProprioConditionConfig": "proprio_condition_config",
+    "cosmos_framework.configs.base.defaults.model_config.DiffusionExpertConfig": "diffusion_expert_config",
+    "cosmos_framework.configs.base.defaults.model_config.FixedStepSamplerConfig": "fixed_step_sampler_config",
+    "cosmos_framework.configs.base.defaults.model_config.LBLConfig": "lbl_config",
+    "cosmos_framework.configs.base.defaults.model_config.OmniMoTModelConfig": "omni_mot_model_config",
+    "cosmos_framework.configs.base.defaults.model_config.RectifiedFlowInferenceConfig": "rectified_flow_inference_config",
+    "cosmos_framework.configs.base.defaults.model_config.RectifiedFlowTrainingConfig": "rectified_flow_training_config",
+    "cosmos_framework.configs.base.defaults.activation_checkpointing.ActivationCheckpointingConfig": "activation_checkpointing_config",
+    "cosmos_framework.configs.base.defaults.compile.CompileConfig": "compile_config",
+    "cosmos_framework.configs.base.defaults.ema.EMAConfig": "ema_config",
+    "cosmos_framework.configs.base.defaults.parallelism.ParallelismConfig": "parallelism_config",
+    "cosmos_framework.configs.base.defaults.quantization.QuantizationConfig": "quantization_config",
+    "cosmos_framework.configs.base.defaults.vlm.PretrainedWeightsConfig": "pretrained_weights_config",
+    "cosmos_framework.configs.base.defaults.vlm.VLMConfig": "vlm_config",
+    "cosmos_framework.model.generator.tokenizers.wan2pt2_vae_4x16x16.Wan2pt2VAEInterface": "wan2pt2_vae_interface",
+    "cosmos_framework.model.generator.mot.unified_mot.Qwen3VLMoTConfig.from_json_file": "qwen3_vl_mot_config_from_json_file",
+    "cosmos_framework.model.generator.mot.unified_mot.Qwen3VLTextForCausalLM": "qwen3_vl_text_for_causal_lm",
+    "cosmos_framework.model.generator.omni_mot_model.OmniMoTModel": "omni_mot_model",
+    "cosmos_framework.configs.base.defaults.reasoner.create_qwen2_tokenizer_with_download": "create_qwen2_tokenizer_with_download",
+    "cosmos_framework.configs.base.defaults.reasoner.create_vlm_config": "create_vlm_config",
+}
+
+_TYPE_ALIASES = {
+    # Zeva branch (canonicalized under projects.cosmos3.vfm.configs.base.*)
+    "projects.cosmos3.vfm.configs.base.defaults.model_config.ZevaPolicyConfig": "zeva_policy_config",
+    "projects.cosmos3.vfm.configs.base.defaults.model_config.BehaviorStage2Config": "behavior_stage2_config",
+    "projects.cosmos3.vfm.configs.base.defaults.model_config.ProprioConditionConfig": "proprio_condition_config",
+    "projects.cosmos3.vfm.configs.base.defaults.activation_checkpointing.ActivationCheckpointingConfig": "activation_checkpointing_config",
+    "projects.cosmos3.vfm.configs.base.defaults.compile.CompileConfig": "compile_config",
+    "projects.cosmos3.vfm.configs.base.defaults.ema.EMAConfig": "ema_config",
+    "projects.cosmos3.vfm.configs.base.defaults.model_config.DiffusionExpertConfig": "diffusion_expert_config",
+    "projects.cosmos3.vfm.configs.base.defaults.model_config.FixedStepSamplerConfig": "fixed_step_sampler_config",
+    "projects.cosmos3.vfm.configs.base.defaults.model_config.LBLConfig": "lbl_config",
+    "projects.cosmos3.vfm.configs.base.defaults.model_config.OmniMoTModelConfig": "omni_mot_model_config",
+    "projects.cosmos3.vfm.configs.base.defaults.model_config.RectifiedFlowInferenceConfig": "rectified_flow_inference_config",
+    "projects.cosmos3.vfm.configs.base.defaults.model_config.RectifiedFlowTrainingConfig": "rectified_flow_training_config",
+    "projects.cosmos3.vfm.configs.base.defaults.parallelism.ParallelismConfig": "parallelism_config",
+    "projects.cosmos3.vfm.configs.base.defaults.quantization.QuantizationConfig": "quantization_config",
+    "projects.cosmos3.vfm.configs.base.defaults.vlm.PretrainedWeightsConfig": "pretrained_weights_config",
+    "projects.cosmos3.vfm.configs.base.defaults.vlm.VLMConfig": "vlm_config",
+}
+
+# Config objects can be serialized as either `_type` or `_target_` depending on
+# whether they came from structured config metadata or LazyCall construction.
+_TARGET_ALIASES.update(_TYPE_ALIASES)
+_TARGET_ALIASES.update(_CF_TYPE_ALIASES)
+_TYPE_ALIASES.update(_CF_TYPE_ALIASES)
+
+_ALIAS_TARGETS = {alias: target for target, alias in _TARGET_ALIASES.items()}
+_ALIAS_TYPES = {alias: target for target, alias in _TYPE_ALIASES.items()}
+
+
+def build_public_model_config(model_config: dict[str, Any]) -> dict[str, Any]:
+    """Build public HF model metadata from the internal LazyConfig model dict."""
+
+    return _to_public_model_config(model_config)
+
+
+def restore_model_config_from_public_model_config(model_config: dict[str, Any]) -> dict[str, Any]:
+    """Restore the internal LazyConfig model dict from public HF model metadata."""
+
+    return _from_public_model_config(model_config)
+
+
+def model_config_uses_public_aliases(model_config: Any) -> bool:
+    """Return whether a model config uses public aliases."""
+
+    return _has_public_alias(model_config)
+
+
+def load_model_config_from_hf_config(config_dict: dict[str, Any]) -> dict[str, Any]:
+    """Load internal model config from either legacy or public HF config."""
+
+    if "model" in config_dict:
+        model_config = config_dict["model"]
+        if model_config_uses_public_aliases(model_config):
+            return restore_model_config_from_public_model_config(model_config)
+        return model_config
+    raise KeyError("HF config must contain 'model'")
+
+
+def _has_key(obj: Any, key: str) -> bool:
+    if isinstance(obj, dict):
+        return key in obj or any(_has_key(value, key) for value in obj.values())
+    if isinstance(obj, list):
+        return any(_has_key(value, key) for value in obj)
+    return False
+
+
+def _has_public_alias(obj: Any) -> bool:
+    if isinstance(obj, dict):
+        for key, value in obj.items():
+            if key in {_PUBLIC_TARGET_KEY, _CLASS_NAME_KEY, _CONFIG_NAME_KEY}:
+                return True
+            if key == _PUBLIC_TYPE_KEY and _is_type_alias(value):
+                return True
+            if _has_public_alias(value):
+                return True
+    elif isinstance(obj, list):
+        return any(_has_public_alias(value) for value in obj)
+    return False
+
+
+def _to_public_model_config(obj: Any) -> Any:
+    if isinstance(obj, list):
+        return [_to_public_model_config(item) for item in obj]
+    if isinstance(obj, dict):
+        data = {}
+        for key, value in obj.items():
+            if key == _TARGET_KEY:
+                data[_PUBLIC_TARGET_KEY] = _target_to_alias(value)
+            elif key == _TYPE_KEY:
+                data[_PUBLIC_TYPE_KEY] = _type_to_alias(value)
+            else:
+                data[key] = _to_public_model_config(value)
+        return data
+    if isinstance(obj, str):
+        return _to_public_string(obj)
+    return copy.deepcopy(obj)
+
+
+def _from_public_model_config(obj: Any) -> Any:
+    if isinstance(obj, list):
+        return [_from_public_model_config(item) for item in obj]
+    if isinstance(obj, dict):
+        data = {}
+        for key, value in obj.items():
+            if key in {_PUBLIC_TARGET_KEY, _CLASS_NAME_KEY}:
+                data[_TARGET_KEY] = _alias_to_runtime_target(value)
+            elif key == _CONFIG_NAME_KEY or (key == _PUBLIC_TYPE_KEY and _is_type_alias(value)):
+                data[_TYPE_KEY] = _alias_to_runtime_type(value)
+            else:
+                data[key] = _from_public_model_config(value)
+        return data
+    if isinstance(obj, str):
+        return _from_public_string(obj)
+    return copy.deepcopy(obj)
+
+
+def _target_to_alias(target: Any) -> str:
+    if not isinstance(target, str):
+        raise TypeError(f"Expected target path string, got {type(target)}")
+    canonical = _canonicalize_module_path(target)
+    try:
+        return _TARGET_ALIASES[canonical]
+    except KeyError as exc:
+        raise ValueError(f"No public alias registered for target path: {target}") from exc
+
+
+def _type_to_alias(tp: Any) -> str:
+    if not isinstance(tp, str):
+        raise TypeError(f"Expected type path string, got {type(tp)}")
+    canonical = _canonicalize_module_path(tp)
+    try:
+        return _TYPE_ALIASES[canonical]
+    except KeyError as exc:
+        raise ValueError(f"No public alias registered for type path: {tp}") from exc
+
+
+def _alias_to_runtime_target(alias: Any) -> str:
+    if not isinstance(alias, str):
+        raise TypeError(f"Expected target alias string, got {type(alias)}")
+    try:
+        return _runtime_module_path(_ALIAS_TARGETS[alias])
+    except KeyError as exc:
+        raise ValueError(f"Unknown Cosmos target alias: {alias}") from exc
+
+
+def _alias_to_runtime_type(alias: Any) -> str:
+    if not isinstance(alias, str):
+        raise TypeError(f"Expected type alias string, got {type(alias)}")
+    try:
+        return _runtime_module_path(_ALIAS_TYPES[alias])
+    except KeyError as exc:
+        raise ValueError(f"Unknown Cosmos type alias: {alias}") from exc
+
+
+def _is_type_alias(value: Any) -> bool:
+    return isinstance(value, str) and value in _ALIAS_TYPES
+
+
+def _canonicalize_module_path(path: str) -> str:
+    replacements = (
+        # vlm → reasoner rename (upstream i4). Longer/more-specific rules first
+        # so the reasoner subtree isn't shadowed by the general vfm rules below.
+        ("cosmos_framework.configs.base.defaults.reasoner.", "projects.cosmos3.vfm.configs.base.defaults.vlm."),
+        ("cosmos_framework.configs.base.reasoner.", "projects.cosmos3.vfm.configs.base.vlm."),
+        ("cosmos_framework.model.generator.reasoner.", "projects.cosmos3.vfm.models.vlm."),
+        ("cosmos_framework.data.generator.reasoner.", "projects.cosmos3.vfm.datasets.vlm."),
+        ("cosmos_framework.data.generator.augmentors.reasoner.", "projects.cosmos3.vfm.datasets.augmentors.vlm."),
+        ("cosmos_framework.utils.generator.reasoner.", "projects.cosmos3.vfm.utils.vlm."),
+        ("cosmos3._src.vfm.", "projects.cosmos3.vfm."),
+        ("cosmos_framework.configs.base.", "projects.cosmos3.vfm.configs.base."),
+        ("cosmos_framework.model.generator.tokenizers.", "projects.cosmos3.vfm.tokenizers."),
+        ("cosmos_framework.model.generator.diffusion.", "projects.cosmos3.vfm.diffusion."),
+        ("cosmos_framework.model.generator.", "projects.cosmos3.vfm.models."),
+        ("cosmos_framework.data.generator.processors.", "projects.cosmos3.vfm.processors."),
+        ("cosmos.model.vfm.tokenizers.", "projects.cosmos3.vfm.tokenizers."),
+        ("cosmos.model.vfm.diffusion.", "projects.cosmos3.vfm.diffusion."),
+        ("cosmos.model.vfm.", "projects.cosmos3.vfm.models."),
+        ("cosmos.configs.base.", "projects.cosmos3.vfm.configs.base."),
+    )
+    for old, new in replacements:
+        if path.startswith(old):
+            return new + path[len(old) :]
+    return path
+
+
+def _runtime_module_path(canonical_path: str) -> str:
+    if _module_exists("cosmos_framework.model.generator"):
+        return _replace_vfm_module_prefix(canonical_path, package="cosmos_framework")
+    if _module_exists("cosmos.model.vfm"):
+        return _replace_vfm_module_prefix(canonical_path, package="cosmos")
+    if _module_exists("cosmos3._src.vfm"):
+        return canonical_path.replace("projects.cosmos3.vfm.", "cosmos3._src.vfm.", 1)
+    return canonical_path
+
+
+def _replace_vfm_module_prefix(canonical_path: str, *, package: str) -> str:
+    replacements = (
+        # vlm → reasoner (upstream rename). MUST precede the general vfm rules
+        # so the specific vlm→reasoner subtree isn't shadowed by them.
+        ("projects.cosmos3.vfm.configs.base.defaults.vlm.", f"{package}.configs.base.defaults.reasoner."),
+        ("projects.cosmos3.vfm.configs.base.vlm.", f"{package}.configs.base.reasoner."),
+        ("projects.cosmos3.vfm.models.vlm.", f"{package}.model.generator.reasoner."),
+        ("projects.cosmos3.vfm.datasets.vlm.", f"{package}.data.generator.reasoner."),
+        ("projects.cosmos3.vfm.datasets.augmentors.vlm.", f"{package}.data.generator.augmentors.reasoner."),
+        ("projects.cosmos3.vfm.utils.vlm.", f"{package}.utils.generator.reasoner."),
+        ("projects.cosmos3.vfm.configs.base.", f"{package}.configs.base."),
+        ("projects.cosmos3.vfm.models.", f"{package}.model.generator."),
+        ("projects.cosmos3.vfm.tokenizers.", f"{package}.model.generator.tokenizers."),
+        ("projects.cosmos3.vfm.diffusion.", f"{package}.model.generator.diffusion."),
+        ("projects.cosmos3.vfm.processors.", f"{package}.data.generator.processors."),
+        ("projects.cosmos3.vfm.datasets.", f"{package}.data.generator."),
+        ("projects.cosmos3.vfm.scripts.action.", f"{package}.data.generator.action_scripts."),
+        ("projects.cosmos3.vfm.utils.", f"{package}.utils.generator."),
+    )
+    for old, new in replacements:
+        if canonical_path.startswith(old):
+            return new + canonical_path[len(old) :]
+    return canonical_path
+
+
+def _to_public_string(value: str) -> str:
+    for prefix in ("projects/cosmos3/vfm/", "cosmos3/_src/vfm/"):
+        if value.startswith(prefix):
+            return _PUBLIC_VFM_URI_PREFIX + value[len(prefix) :]
+    for package in ("cosmos_framework", "cosmos"):
+        public_value = _public_string_from_runtime_file_prefix(value, package=package)
+        if public_value is not None:
+            return public_value
+    return value
+
+
+def _public_string_from_runtime_file_prefix(value: str, *, package: str) -> str | None:
+    replacements = (
+        # reasoner → vlm (inverse of the upstream rename). MUST precede the
+        # general vfm rules so the reasoner subtree isn't shadowed by them.
+        # Mirrors the module-path rules in ``_canonicalize_module_path``.
+        (f"{package}/configs/base/defaults/reasoner/", "configs/base/defaults/vlm/"),
+        (f"{package}/configs/base/reasoner/", "configs/base/vlm/"),
+        (f"{package}/model/generator/reasoner/", "models/vlm/"),
+        (f"{package}/data/generator/augmentors/reasoner/", "datasets/augmentors/vlm/"),
+        (f"{package}/data/generator/reasoner/", "datasets/vlm/"),
+        (f"{package}/utils/generator/reasoner/", "utils/vlm/"),
+        (f"{package}/configs/base/", "configs/base/"),
+        (f"{package}/model/generator/tokenizers/", "tokenizers/"),
+        (f"{package}/model/generator/diffusion/", "diffusion/"),
+        (f"{package}/model/generator/", "models/"),
+        (f"{package}/data/generator/processors/", "processors/"),
+        (f"{package}/data/generator/action_scripts/", "scripts/action/"),
+        (f"{package}/data/generator/", "datasets/"),
+        (f"{package}/utils/generator/", "utils/"),
+    )
+    for old, new in replacements:
+        if value.startswith(old):
+            return _PUBLIC_VFM_URI_PREFIX + new + value[len(old) :]
+    return None
+
+
+def _from_public_string(value: str) -> str:
+    if not value.startswith(_PUBLIC_VFM_URI_PREFIX):
+        return value
+    suffix = value[len(_PUBLIC_VFM_URI_PREFIX) :]
+    if _module_exists("cosmos_framework.model.generator"):
+        return _replace_vfm_file_prefix(suffix, package="cosmos_framework")
+    if _module_exists("cosmos.model.vfm"):
+        return _replace_vfm_file_prefix(suffix, package="cosmos")
+    if _module_exists("cosmos3._src.vfm"):
+        return f"cosmos3/_src/vfm/{suffix}"
+    return f"projects/cosmos3/vfm/{suffix}"
+
+
+def _replace_vfm_file_prefix(suffix: str, *, package: str) -> str:
+    replacements = (
+        # vlm → reasoner (upstream rename). MUST precede the general vfm rules
+        # so the specific vlm→reasoner subtree isn't shadowed by them. Mirrors
+        # the module-path rules in ``_replace_vfm_module_prefix``.
+        ("configs/base/defaults/vlm/", f"{package}/configs/base/defaults/reasoner/"),
+        ("configs/base/vlm/", f"{package}/configs/base/reasoner/"),
+        ("models/vlm/", f"{package}/model/generator/reasoner/"),
+        ("datasets/augmentors/vlm/", f"{package}/data/generator/augmentors/reasoner/"),
+        ("datasets/vlm/", f"{package}/data/generator/reasoner/"),
+        ("utils/vlm/", f"{package}/utils/generator/reasoner/"),
+        ("configs/base/", f"{package}/configs/base/"),
+        ("models/", f"{package}/model/generator/"),
+        ("tokenizers/", f"{package}/model/generator/tokenizers/"),
+        ("diffusion/", f"{package}/model/generator/diffusion/"),
+        ("processors/", f"{package}/data/generator/processors/"),
+        ("datasets/", f"{package}/data/generator/"),
+        ("scripts/action/", f"{package}/data/generator/action_scripts/"),
+        ("utils/", f"{package}/utils/generator/"),
+    )
+    for old, new in replacements:
+        if suffix.startswith(old):
+            return new + suffix[len(old) :]
+    return f"{package}/_vfm_unmapped/{suffix}"
+
+
+def _module_exists(module: str) -> bool:
+    try:
+        return importlib.util.find_spec(module) is not None
+    except ModuleNotFoundError:
+        return False

releases/bridge-native-v3/code/cosmos_framework/model/generator/mot/cosmos3_vfm_network.py

@@ -0,0 +1,1536 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+import math
+from typing import List, Tuple
+
+import torch
+from torch import nn
+from transformers.configuration_utils import PretrainedConfig
+from transformers.modeling_utils import PreTrainedModel
+
+from cosmos_framework.data.generator.sequence_packing import ModalityData, PackedSequence
+from cosmos_framework.data.generator.sequence_packing.natten import verify_natten_parameter_list
+from cosmos_framework.model.zeva import (
+    BriefInteractionTrace,
+    CausalPromptConfig,
+    CausalPromptEncoder,
+    CausalPromptPolicyAdapter,
+    PolicyInjectionConfig,
+    PolicyInjectionPrior,
+    inject_causal_prompt,
+)
+from cosmos_framework.model.generator.mot.attention import SplitInfo, build_packed_sequence
+from cosmos_framework.model.generator.mot.context_parallel_utils import (
+    get_context_parallel_last_hidden_state,
+    get_context_parallel_sharded_sequence,
+)
+from cosmos_framework.model.generator.mot.domain_aware_linear import DomainAwareLinear
+from cosmos_framework.model.generator.mot.modeling_utils import TimestepEmbedder, has_noisy_tokens
+from cosmos_framework.model.generator.utils.memory import MemoryState
+
+
+class Cosmos3VFMNetworkConfig(PretrainedConfig):
+    def __init__(
+        self,
+        vision_gen=True,
+        action_gen=False,
+        sound_gen=False,
+        vlm_config=None,
+        latent_patch_size=2,
+        latent_downsample_factor=8,
+        latent_channel_size=16,
+        max_latent_h=32,
+        max_latent_w=32,
+        max_latent_t=32,
+        enable_fps_modulation=False,
+        base_fps=24,
+        vit_max_num_patch_per_side=70,
+        connector_act="gelu_pytorch_tanh",
+        interpolate_pos=False,
+        timestep_shift=1.0,
+        timestep_scale=0.001,
+        predict_text_tokens=False,
+        joint_attn_implementation="two_way",
+        action_dim=32,
+        num_embodiment_domains=32,
+        temporal_compression_factor_vision=4,
+        temporal_compression_factor_action=1,
+        natten_parameter_list=None,
+        video_temporal_causal=False,
+        # Sound generation parameters
+        sound_dim: int | None = None,
+        temporal_compression_factor_sound=1,
+        sound_latent_fps: int = 25,
+        enable_input_bias: bool = True,
+        behavior_stage2_config: dict | None = None,
+        proprio_condition_config: dict | None = None,
+        residual_expert_config: dict | None = None,
+        **kwargs,
+    ):
+        self.vision_gen = vision_gen
+        self.sound_gen = sound_gen
+        self.vlm_config = vlm_config
+        self.latent_patch_size = latent_patch_size
+        self.latent_downsample_factor = latent_downsample_factor
+        self.latent_channel_size = latent_channel_size
+        self.max_latent_h = max_latent_h
+        self.max_latent_w = max_latent_w
+        self.max_latent_t = max_latent_t
+        self.enable_fps_modulation = enable_fps_modulation
+        self.base_fps = base_fps
+        self.vit_max_num_patch_per_side = vit_max_num_patch_per_side
+        self.connector_act = connector_act
+        self.interpolate_pos = interpolate_pos
+        self.timestep_shift = timestep_shift
+        self.timestep_scale = timestep_scale
+        self.predict_text_tokens = predict_text_tokens
+        self.joint_attn_implementation = joint_attn_implementation
+        self.temporal_compression_factor_vision = temporal_compression_factor_vision
+        self.natten_parameter_list = natten_parameter_list
+        self.video_temporal_causal = video_temporal_causal
+        self.enable_input_bias = enable_input_bias
+        self.behavior_stage2_config = behavior_stage2_config
+        self.proprio_condition_config = proprio_condition_config
+        self.residual_expert_config = residual_expert_config
+
+        # action related parameters
+        self.action_gen = action_gen  # whether to generate action tokens
+        self.action_dim = action_dim
+        self.num_embodiment_domains = num_embodiment_domains
+        self.temporal_compression_factor_action = temporal_compression_factor_action
+        if self.action_gen:
+            assert self.vision_gen, (
+                "Action generation requires visual generation! We do NOT support action only training!"
+            )
+
+        # sound related parameters
+        self.sound_dim = sound_dim
+        self.temporal_compression_factor_sound = temporal_compression_factor_sound
+        self.sound_latent_fps = sound_latent_fps
+        if self.sound_gen:
+            assert self.vision_gen, (
+                "Sound generation requires visual generation! We do NOT support sound only training!"
+            )
+
+        super().__init__(**kwargs)
+
+
+class Cosmos3VFMNetwork(PreTrainedModel):
+    config_class = Cosmos3VFMNetworkConfig
+    base_model_prefix = "cosmos3"
+
+    def __init__(self, language_model, config: Cosmos3VFMNetworkConfig):
+        super().__init__(config)
+        self.language_model = language_model
+
+        text_config = config.vlm_config.text_config if hasattr(config.vlm_config, "text_config") else config.vlm_config
+        self.hidden_size = text_config.hidden_size
+        self.num_heads = text_config.num_attention_heads
+        self.num_kv_heads = text_config.num_key_value_heads
+        self.head_dim = text_config.head_dim
+        self.num_hidden_layers = text_config.num_hidden_layers
+        self.attention_io_layout = "sequence_sharded"
+        self.predict_text_tokens = config.predict_text_tokens
+
+        if config.natten_parameter_list is not None and config.joint_attn_implementation != "three_way":
+            raise NotImplementedError(
+                f"Sparsity is only supported with 'three_way' attention, but got {config.joint_attn_implementation=}, "
+                "and 'natten_parameter_list' was not None."
+            )
+        self.natten_parameter_list = verify_natten_parameter_list(
+            config.natten_parameter_list, num_layers=self.num_hidden_layers
+        )
+
+        if config.video_temporal_causal and config.joint_attn_implementation != "three_way":
+            raise ValueError(
+                f"video_temporal_causal=True requires joint_attn_implementation='three_way', "
+                f"but got {config.joint_attn_implementation!r}."
+            )
+        self.video_temporal_causal = config.video_temporal_causal
+        self.pad_for_cuda_graphs = False
+
+        if config.vision_gen:
+            self.latent_patch_size = config.latent_patch_size
+            self.timestep_shift = config.timestep_shift
+            self.timestep_scale = config.timestep_scale
+            self.latent_downsample = config.latent_downsample_factor * config.latent_patch_size
+            self.max_latent_h = config.max_latent_h
+            self.max_latent_w = config.max_latent_w
+            self.max_latent_t = config.max_latent_t
+            self.latent_channel = config.latent_channel_size
+            self.patch_latent_dim = self.latent_patch_size**2 * self.latent_channel
+
+            _input_bias = config.enable_input_bias
+            self.time_embedder = TimestepEmbedder(self.hidden_size, bias=_input_bias)
+            self.vae2llm = nn.Linear(self.patch_latent_dim, self.hidden_size, bias=_input_bias)
+            self.llm2vae = nn.Linear(self.hidden_size, self.patch_latent_dim)
+
+        if config.action_gen:
+            self.action_dim = config.action_dim
+            self.num_embodiment_domains = config.num_embodiment_domains
+            self.action2llm = DomainAwareLinear(self.action_dim, self.hidden_size, self.num_embodiment_domains)
+            self.llm2action = DomainAwareLinear(self.hidden_size, self.action_dim, self.num_embodiment_domains)
+
+            self.action_modality_embed = nn.Parameter(torch.zeros(self.hidden_size))
+            residual_cfg = config.residual_expert_config or {}
+            self.action_residual_head: nn.Linear | None = None
+            if bool(residual_cfg.get("enabled", False)):
+                self.action_residual_head = nn.Linear(self.hidden_size, self.action_dim, bias=True)
+        else:
+            self.action_residual_head = None
+
+        behavior_cfg = config.behavior_stage2_config
+        if behavior_cfg is not None and behavior_cfg.get("enabled", False):
+            if not config.action_gen:
+                raise ValueError("Zeva policy injection requires action_gen=True")
+            prior_cfg = PolicyInjectionConfig(
+                global_dim=int(behavior_cfg["global_dim"]),
+                phase_dim=int(behavior_cfg["phase_dim"]),
+                effect_dim=int(behavior_cfg.get("effect_dim", 128)),
+                effect_history_length=int(behavior_cfg.get("effect_history_length", 4)),
+                action_dim=int(behavior_cfg["action_dim"]),
+                horizon=int(behavior_cfg["horizon"]),
+                num_anchors=int(behavior_cfg["num_anchors"]),
+                hidden_dim=int(behavior_cfg["hidden_dim"]),
+                num_heads=int(behavior_cfg["num_heads"]),
+            )
+            # These registered attribute names are retained solely because they
+            # are part of the released DCP state-dict schema.
+            self.behavior_pbd: PolicyInjectionPrior | None = PolicyInjectionPrior(prior_cfg)
+            self.behavior_adapter: CausalPromptPolicyAdapter | None = CausalPromptPolicyAdapter(
+                action_dim=prior_cfg.action_dim, hidden_dim=self.hidden_size
+            )
+            # Zeva's second conditioning path: the episode-level task context
+            # is a learned understanding-prefix token visible to every
+            # action query through Cosmos two-way attention.
+            self.behavior_global_projector: nn.Linear | None = nn.Linear(
+                prior_cfg.global_dim, self.hidden_size, bias=True
+            )
+            pim_enabled = bool(behavior_cfg.get("pim_memory_enabled", False))
+            pim_dim = int(behavior_cfg.get("pim_context_dim", 256))
+            self.behavior_pim_encoder: CausalPromptEncoder | None = (
+                CausalPromptEncoder(
+                    CausalPromptConfig(
+                        global_dim=prior_cfg.global_dim,
+                        phase_dim=prior_cfg.phase_dim,
+                        effect_dim=prior_cfg.effect_dim,
+                        brief_length=prior_cfg.effect_history_length,
+                        persistent_length=int(behavior_cfg.get("pim_persistent_length", 4)),
+                        hidden_dim=pim_dim,
+                        num_heads=int(behavior_cfg.get("num_heads", 4)),
+                    )
+                )
+                if pim_enabled
+                else None
+            )
+            self.behavior_pim_projector: nn.Linear | None = (
+                nn.Linear(pim_dim, self.hidden_size, bias=True) if pim_enabled else None
+            )
+            self.behavior_pim_gate: nn.Parameter | None = (
+                nn.Parameter(torch.empty(1)) if pim_enabled else None
+            )
+            self.behavior_pim_gate_init = float(behavior_cfg.get("pim_gate_init", 0.0))
+            self.behavior_pim_force_bypass = bool(behavior_cfg.get("pim_force_bypass", False))
+            online_enabled = bool(behavior_cfg.get("online_memory_enabled", False))
+            online_dim = int(behavior_cfg.get("online_context_dim", 256))
+            self.behavior_online_projector: nn.Linear | None = (
+                nn.Linear(online_dim, self.hidden_size, bias=True) if online_enabled else None
+            )
+            self.behavior_prior_dropout_rate = float(behavior_cfg["prior_dropout_rate"])
+            self.behavior_prior_inference_guidance_scale = float(
+                behavior_cfg.get("prior_inference_guidance_scale", 0.5)
+            )
+            self.behavior_prior_leading_condition_steps = int(
+                behavior_cfg.get("leading_condition_steps", 1)
+            )
+        else:
+            self.behavior_pbd = None
+            self.behavior_adapter = None
+            self.behavior_global_projector = None
+            self.behavior_pim_encoder = None
+            self.behavior_pim_projector = None
+            self.behavior_pim_gate = None
+            self.behavior_pim_gate_init = 0.0
+            self.behavior_pim_force_bypass = False
+            self.behavior_online_projector = None
+            self.behavior_prior_dropout_rate = 0.0
+            self.behavior_prior_inference_guidance_scale = 1.0
+            self.behavior_prior_leading_condition_steps = 1
+
+        proprio_cfg = config.proprio_condition_config
+        if proprio_cfg is not None and proprio_cfg.get("enabled", False):
+            if not config.action_gen:
+                raise ValueError("Proprio conditioning requires action_gen=True")
+            self.proprio_input_dim = int(proprio_cfg["input_dim"])
+            self.proprio_prefix_tokens = int(proprio_cfg.get("prefix_tokens", 1))
+            if self.proprio_input_dim < 1 or self.proprio_prefix_tokens < 1:
+                raise ValueError("Proprio input_dim and prefix_tokens must be positive")
+            self.proprio_projector: nn.Linear | None = nn.Linear(
+                self.proprio_input_dim,
+                self.proprio_prefix_tokens * self.hidden_size,
+                bias=True,
+            )
+        else:
+            self.proprio_input_dim = 0
+            self.proprio_prefix_tokens = 0
+            self.proprio_projector = None
+
+        if config.sound_gen:
+            self.sound_dim = config.sound_dim
+            self.sound2llm = nn.Linear(config.sound_dim, self.hidden_size, bias=config.enable_input_bias)
+            self.llm2sound = nn.Linear(self.hidden_size, config.sound_dim)
+            self.sound_modality_embed = nn.Parameter(torch.zeros(self.hidden_size))
+
+        self.config = config
+        self.parallel_dims = None
+
+    @property
+    def causal_prompt_encoder(self) -> CausalPromptEncoder | None:
+        """Paper-facing alias; storage keeps the legacy checkpoint key."""
+        return self.behavior_pim_encoder
+
+    @property
+    def causal_prompt_projector(self) -> nn.Linear | None:
+        """Project the fused causal prompt into the frozen policy prefix."""
+        return self.behavior_pim_projector
+
+    @property
+    def causal_prompt_gate(self) -> nn.Parameter | None:
+        """Zero-initialized policy-injection gate for the causal prompt."""
+        return self.behavior_pim_gate
+
+    def init_weights(self, buffer_device: torch.device | None):
+        if self.config.vision_gen or self.config.action_gen or self.config.sound_gen:
+            self.time_embedder._init_weights(buffer_device=buffer_device)
+
+        if self.config.vision_gen:
+            std = 1.0 / math.sqrt(self.patch_latent_dim)
+            torch.nn.init.trunc_normal_(self.vae2llm.weight, std=std, a=-3 * std, b=3 * std)
+            if self.config.enable_input_bias:
+                torch.nn.init.zeros_(self.vae2llm.bias)
+
+            std = 1.0 / math.sqrt(self.hidden_size)
+            torch.nn.init.trunc_normal_(self.llm2vae.weight, std=std, a=-3 * std, b=3 * std)
+            torch.nn.init.zeros_(self.llm2vae.bias)
+
+        if self.config.action_gen:
+            # DomainAwareLinear uses embeddings for weights, so we initialize them differently
+            # action2llm: input_size=action_dim, output_size=hidden_size
+            std = 1.0 / math.sqrt(self.action_dim)
+            torch.nn.init.trunc_normal_(self.action2llm.fc.weight, std=std, a=-3 * std, b=3 * std)
+            torch.nn.init.zeros_(self.action2llm.bias.weight)
+
+            # llm2action: input_size=hidden_size, output_size=action_dim
+            std = 1.0 / math.sqrt(self.hidden_size)
+            torch.nn.init.trunc_normal_(self.llm2action.fc.weight, std=std, a=-3 * std, b=3 * std)
+            torch.nn.init.zeros_(self.llm2action.bias.weight)
+
+            std = 1.0 / math.sqrt(self.hidden_size)
+            torch.nn.init.trunc_normal_(self.action_modality_embed, std=std, a=-3 * std, b=3 * std)
+            if self.action_residual_head is not None:
+                # Zero initialization makes residual enabling an exact base
+                # bypass before the first optimizer update.
+                torch.nn.init.zeros_(self.action_residual_head.weight)
+                torch.nn.init.zeros_(self.action_residual_head.bias)
+
+        if (
+            self.behavior_pbd is not None
+            and self.behavior_adapter is not None
+            and self.behavior_global_projector is not None
+        ):
+            self.behavior_pbd.init_weights()
+            self.behavior_adapter.init_weights()
+            nn.init.xavier_uniform_(self.behavior_global_projector.weight)
+            nn.init.zeros_(self.behavior_global_projector.bias)
+            if (
+                self.behavior_pim_encoder is not None
+                and self.behavior_pim_projector is not None
+                and self.behavior_pim_gate is not None
+            ):
+                self.behavior_pim_encoder.reset_parameters()
+                nn.init.xavier_uniform_(self.behavior_pim_projector.weight)
+                nn.init.zeros_(self.behavior_pim_projector.bias)
+                nn.init.constant_(self.behavior_pim_gate, self.behavior_pim_gate_init)
+            if self.behavior_online_projector is not None:
+                # Zero initialization makes enabling memory a strict no-op
+                # until the support-memory adaptation stage trains this path.
+                nn.init.zeros_(self.behavior_online_projector.weight)
+                nn.init.zeros_(self.behavior_online_projector.bias)
+        if self.proprio_projector is not None:
+            nn.init.xavier_uniform_(self.proprio_projector.weight)
+            nn.init.zeros_(self.proprio_projector.bias)
+
+        if self.config.sound_gen:
+            # sound2llm: input_size=sound_dim, output_size=hidden_size
+            std = 1.0 / math.sqrt(self.sound_dim)
+            torch.nn.init.trunc_normal_(self.sound2llm.weight, std=std, a=-3 * std, b=3 * std)
+            if self.config.enable_input_bias:
+                torch.nn.init.zeros_(self.sound2llm.bias)
+
+            # llm2sound: input_size=hidden_size, output_size=sound_dim
+            std = 1.0 / math.sqrt(self.hidden_size)
+            torch.nn.init.trunc_normal_(self.llm2sound.weight, std=std, a=-3 * std, b=3 * std)
+            torch.nn.init.zeros_(self.llm2sound.bias)
+
+            std = 1.0 / math.sqrt(self.hidden_size)
+            torch.nn.init.trunc_normal_(self.sound_modality_embed, std=std, a=-3 * std, b=3 * std)
+
+        self.language_model.init_weights(buffer_device=buffer_device)
+
+    def generate_reasoner_text(
+        self,
+        input_ids: torch.Tensor,
+        max_new_tokens: int,
+        *,
+        pixel_values: torch.Tensor | None = None,
+        image_grid_thw: torch.Tensor | None = None,
+        pixel_values_videos: torch.Tensor | None = None,
+        video_grid_thw: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        eos_token_id: int | list[int] | None = None,
+        pad_token_id: int | None = None,
+        do_sample: bool = False,
+        temperature: float = 1.0,
+        top_k: int | None = None,
+        top_p: float | None = None,
+        repetition_penalty: float = 1.0,
+        presence_penalty: float = 0.0,
+        seed: int | None = None,
+        return_only_new_tokens: bool = False,
+    ) -> torch.Tensor:
+        """Autoregressively generate text tokens using only the reasoner tower.
+
+        Thin pass-through to ``self.language_model.generate_reasoner_text``
+        (see ``unified_mot._impl_generate_reasoner_text`` for full argument
+        documentation).  Handles both text-only and image-conditioned (I2V)
+        prompts through this single entry point: pass
+        ``pixel_values`` + ``image_grid_thw`` (and optionally
+        ``attention_mask``) for image-conditioned prefill via the Qwen3-VL
+        visual encoder, or omit them for text-only prefill.  Video
+        conditioning is also supported via ``pixel_values_videos`` +
+        ``video_grid_thw``; the image and video pairs are mutually exclusive.
+        Uses the und-pathway weights (those WITHOUT the ``_moe_gen`` suffix)
+        plus ``embed_tokens`` / ``norm`` / ``lm_head``; the generation pathway
+        and all VFM-level multimodal embedders / heads (``vae2llm``,
+        ``llm2vae``, ``sound2llm``, etc.) are bypassed.
+
+        ``repetition_penalty`` / ``presence_penalty`` are pass-through
+        sampling controls applied inside
+        :func:`unified_mot._impl_generate_reasoner_text` as logit
+        transformations *before* the ``do_sample`` argmax / multinomial
+        branch (so they shift the greedy argmax too).  Identity defaults
+        (``1.0`` / ``0.0``) keep the un-penalized fast path
+        bit-identical.
+
+        ``seed`` is a pass-through sampling-RNG knob: when provided,
+        :func:`unified_mot._impl_generate_reasoner_text` allocates a
+        device-local ``torch.Generator``, seeds it once with
+        ``manual_seed(seed)``, and threads it into every
+        ``torch.multinomial`` draw — making the decoded sequence a
+        deterministic function of the seed, the prompt, and the
+        penalty masks.  ``None`` (default) consumes the device's
+        default RNG and is bit-identical to the pre-seed call surface.
+        Has no effect under greedy decoding (the argmax branch never
+        reads the generator).
+        """
+
+        return self.language_model.generate_reasoner_text(
+            input_ids=input_ids,
+            max_new_tokens=max_new_tokens,
+            pixel_values=pixel_values,
+            image_grid_thw=image_grid_thw,
+            pixel_values_videos=pixel_values_videos,
+            video_grid_thw=video_grid_thw,
+            attention_mask=attention_mask,
+            eos_token_id=eos_token_id,
+            pad_token_id=pad_token_id,
+            do_sample=do_sample,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+            repetition_penalty=repetition_penalty,
+            presence_penalty=presence_penalty,
+            seed=seed,
+            return_only_new_tokens=return_only_new_tokens,
+        )
+
+    def patchify_and_pack_latents(
+        self, tokens_vision: torch.Tensor, token_shapes_vision: List[Tuple[int, int, int]]
+    ) -> tuple[torch.Tensor, List[Tuple[int, int, int]]]:
+        p = self.latent_patch_size
+        # Patchify and pack the latents
+        packed_latent = []
+        original_latent_shapes = []  # Store original shapes for unpadding later
+
+        # C, T, H, W
+        for latent, (t, h, w) in zip(tokens_vision, token_shapes_vision):
+            latent = latent.squeeze(0)  # [C,T,H,W]
+
+            # Get original latent dimensions
+            _, t_actual, h_actual, w_actual = latent.shape
+            original_latent_shapes.append((t_actual, h_actual, w_actual))
+
+            # Compute padded dimensions (must be divisible by p)
+            h_padded = ((h_actual + p - 1) // p) * p
+            w_padded = ((w_actual + p - 1) // p) * p
+
+            # Zero-pad if dimensions are not divisible by p
+            if h_padded != h_actual or w_padded != w_actual:
+                padded = torch.zeros(
+                    (self.latent_channel, t_actual, h_padded, w_padded),
+                    device=latent.device,
+                    dtype=latent.dtype,
+                )  # [C,T,H_padded,W_padded]
+                padded[:, :, :h_actual, :w_actual] = latent
+                latent = padded  # [C,T,H_padded,W_padded]
+
+            # Compute number of patches after padding
+            h_patches = h_padded // p
+            w_patches = w_padded // p
+
+            # Patchify
+            latent = latent.reshape(
+                self.latent_channel, t_actual, h_patches, p, w_patches, p
+            )  # [C,T,h_patches,p,w_patches,p]
+            latent = torch.einsum("cthpwq->thwpqc", latent).reshape(
+                -1, p * p * self.latent_channel
+            )  # [T*h_patches*w_patches,patch_latent_dim]
+            packed_latent.append(latent)
+
+        # We assumed latents we get to the network is already noised
+        packed_latent = torch.cat(packed_latent, dim=0)  # [total_vision_patches,patch_latent_dim]
+        return packed_latent, original_latent_shapes
+
+    def unpatchify_and_unpack_latents(
+        self,
+        packed_mse_preds: torch.Tensor,
+        token_shapes_vision: List[Tuple[int, int, int]],
+        noisy_frame_indexes_vision: list[torch.Tensor],
+        original_latent_shapes: List[Tuple[int, int, int]] | None = None,
+    ) -> list[torch.Tensor]:
+        p = self.latent_patch_size
+        unpatchified_latents = []
+
+        # Split packed_mse_preds back into individual latents based on token_shapes_vision
+        start_idx = 0
+        for i, (t_c, h_c, w_c) in enumerate(token_shapes_vision):
+            # Get original shape for unpadding (if provided)
+            if original_latent_shapes is not None:
+                t_orig, h_orig, w_orig = original_latent_shapes[i]
+                # Compute padded dimensions used during patchify
+                h_padded = ((h_orig + p - 1) // p) * p
+                w_padded = ((w_orig + p - 1) // p) * p
+                h_patches = h_padded // p
+                w_patches = w_padded // p
+            else:
+                # Fallback: use token shapes directly (assumes no padding was needed)
+                t_orig, h_orig, w_orig = t_c, h_c * p, w_c * p
+                h_patches, w_patches = h_c, w_c
+
+            # noisy_frame_indexes_vision is a list of tensors, each with shape (T,),
+            # where the values are the noisy frame indices.
+            noisy_frame_indexes = noisy_frame_indexes_vision[i]
+            t_n = len(noisy_frame_indexes)
+
+            # Initialize with the original shape (after unpadding), zeros for clean frames
+            output_tensor = torch.zeros(
+                (self.latent_channel, t_c, h_orig, w_orig),
+                device=packed_mse_preds.device,
+                dtype=packed_mse_preds.dtype,
+            )  # [C,T,H_orig,W_orig]
+            num_patches = t_n * h_patches * w_patches
+            if num_patches > 0:
+                end_idx = start_idx + num_patches
+                # Extract patches for this latent
+                latent_patches = packed_mse_preds[start_idx:end_idx]  # [num_patches,patch_latent_dim]
+                # Reshape back to [t_n, h_patches, w_patches, p, p, channels]
+                latent_patches = latent_patches.reshape(
+                    t_n, h_patches, w_patches, p, p, self.latent_channel
+                )  # [T_n,h_patches,w_patches,p,p,C]
+                # Invert the einsum operation: "thwpqc->cthpwq"
+                latent = torch.einsum("thwpqc->cthpwq", latent_patches)  # [C,T_n,h_patches,p,w_patches,p]
+                # Reshape back to [channels, t_n, h_padded, w_padded]
+                latent = latent.reshape(
+                    self.latent_channel, t_n, h_patches * p, w_patches * p
+                )  # [C,T_n,H_padded,W_padded]
+
+                # Crop to original dimensions (unpad the zeros)
+                latent = latent[:, :, :h_orig, :w_orig]  # [C,T_n,H_orig,W_orig]
+
+                # Fill only the noisy frame positions using the actual mask indices
+                output_tensor[:, noisy_frame_indexes] = latent
+
+                start_idx = end_idx
+
+            unpatchified_latents.append(output_tensor.unsqueeze(0))  # [1,C,T,H,W]
+
+        # Return list of unpatchified latents (supports variable shapes)
+        return unpatchified_latents
+
+    def pack_action(
+        self,
+        tokens_action: list[torch.Tensor],
+        token_shapes_action: list[tuple[int, ...]],
+        domain_id_action: list[torch.Tensor],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Pack variable-length action tokens into a 1D sequence for transformer input.
+
+        Args:
+            tokens_action: List of action tensors, each [T_i, action_dim] (T_i may vary).
+            token_shapes_action: List of (T_i,) tuples per sample.
+            domain_id_action: List of domain ID tensors, each of shape [1].
+
+        Returns:
+            Tuple of (packed_tokens, per_token_domain_id):
+                packed_tokens: [total_action_tokens, action_dim]
+                per_token_domain_id: [total_action_tokens]
+        """
+        packed: list[torch.Tensor] = []
+        domain_ids: list[torch.Tensor] = []
+        for tokens, shape, d_id in zip(tokens_action, token_shapes_action, domain_id_action):
+            T = shape[0]
+            packed.append(tokens[:T])
+            domain_ids.append(d_id.expand(T))
+        return torch.cat(packed, dim=0), torch.cat(domain_ids, dim=0)
+
+    def unpack_action(
+        self,
+        packed_action_preds: torch.Tensor,
+        token_shapes_action: list[tuple[int, ...]],
+        noisy_frame_indexes_action: list[torch.Tensor],
+    ) -> list[torch.Tensor]:
+        """Unpack action predictions back into per-sample action tensors.
+
+        Args:
+            packed_action_preds: Packed action predictions of shape (total_noisy_tokens, action_dim)
+            token_shapes_action: Per-sample token shapes, each (T_i,) tuple.
+            noisy_frame_indexes_action: List of tensors, each with shape (Tn_i,), where the values
+                are the noisy frame indices for sample i.
+
+        Returns:
+            List of per-sample tensors, each of shape (T_i, action_dim), with predictions
+            placed at noisy positions. Clean positions are left as zeros.
+        """
+        unpacked: list[torch.Tensor] = []
+        start_idx = 0
+        for shape, noisy_frame_indexes in zip(token_shapes_action, noisy_frame_indexes_action):
+            T = shape[0]
+            output = torch.zeros(
+                (T, self.action_dim),
+                device=packed_action_preds.device,
+                dtype=packed_action_preds.dtype,
+            )
+            t_n = len(noisy_frame_indexes)
+            if t_n > 0:
+                end_idx = start_idx + t_n
+                output[noisy_frame_indexes] = packed_action_preds[start_idx:end_idx]
+                start_idx = end_idx
+            unpacked.append(output)
+        return unpacked
+
+    def pack_sound_latents(
+        self,
+        tokens_sound: list[torch.Tensor],
+        token_shapes_sound: list[tuple[int, int, int]],
+    ) -> torch.Tensor:
+        """Pack sound latents into a 1D sequence for transformer input.
+
+        Args:
+            tokens_sound: List of sound latent tensors, each [C, T]
+            token_shapes_sound: List of (T, 1, 1) tuples per sample
+
+        Returns:
+            Packed tensor of shape [total_sound_tokens, C]
+        """
+        packed = []
+        for sound, shape in zip(tokens_sound, token_shapes_sound):
+            T = shape[0]
+            # sound: [C, T] → take first T frames → [C, T]
+            # Then permute to [T, C] for packing
+            sound_tokens = sound[:, :T].permute(1, 0)  # [T,C]
+            packed.append(sound_tokens)
+        return torch.cat(packed, dim=0)  # [total_sound_tokens,C]
+
+    def unpack_sound_latents(
+        self,
+        packed_sound_preds: torch.Tensor,
+        token_shapes_sound: list[tuple[int, int, int]],
+        noisy_frame_indexes_sound: list[torch.Tensor],
+    ) -> list[torch.Tensor]:
+        """Unpack sound predictions back into per-sample sound latents.
+
+        Args:
+            packed_sound_preds: Packed sound predictions of shape (total_noisy_tokens, sound_dim)
+            token_shapes_sound: List of (T, 1, 1) tuples per sample
+            noisy_frame_indexes_action: List of tensors, each with shape (T_i,), where the values
+                are the noisy frame indices. T_i <= max_T.
+
+        Returns:
+            List of per-sample tensors, each [C, T], with predictions placed at noisy positions.
+            Clean positions are left as zeros.
+        """
+        unpacked = []
+        start_idx = 0
+        for shape, noisy_frame_indexes in zip(token_shapes_sound, noisy_frame_indexes_sound):
+            T = shape[0]
+            # Initialize output with zeros for clean positions
+            output = torch.zeros(
+                (self.sound_dim, T),
+                device=packed_sound_preds.device,
+                dtype=packed_sound_preds.dtype,
+            )
+
+            t_n = len(noisy_frame_indexes)
+
+            if t_n > 0:
+                end_idx = start_idx + t_n
+                # packed_sound_preds: [total_noisy_tokens, C] → transpose and fill at noisy positions
+                output[:, noisy_frame_indexes] = packed_sound_preds[
+                    start_idx:end_idx
+                ].T  # packed_sound_preds[...]: [T_n,C] → .T: [C,T_n]
+                start_idx = end_idx
+
+            unpacked.append(output)
+        return unpacked
+
+    def _encode_text(
+        self,
+        packed_seq: PackedSequence,
+    ) -> tuple[torch.Tensor, torch.dtype]:
+        """Embed text tokens and initialize packed_sequence.
+
+        Args:
+            packed_seq: PackedSequence containing text_ids and text_indexes.
+
+        Returns:
+            tuple of (packed_sequence, target_dtype) where packed_sequence has text embeddings filled in.
+        """
+        packed_text_embedding = self.language_model.model.embed_tokens(packed_seq.text_ids)  # [N_text,hidden_size]
+        packed_sequence = packed_text_embedding.new_zeros(
+            size=(packed_seq.sequence_length, self.hidden_size)
+        )  # [N_total,hidden_size]
+        packed_sequence[packed_seq.text_indexes] = (
+            packed_text_embedding  # [N_text,hidden_size] scattered into [N_total,hidden_size]
+        )
+        return packed_sequence, packed_text_embedding.dtype
+
+    def _encode_behavior_prefix(
+        self,
+        packed_seq: PackedSequence,
+        packed_sequence: torch.Tensor,
+        target_dtype: torch.dtype,
+    ) -> None:
+        """Fill the learned global-behavior causal slots reserved by the packer."""
+        if self.behavior_global_projector is None:
+            if packed_seq.behavior_indexes.numel():
+                raise ValueError("Packed behavior prefix slots require behavior_stage2.enabled=True")
+            return
+        if not packed_seq.behavior_indexes.numel():
+            # Clean readout / vision-only calls intentionally contain no action
+            # sample and therefore reserve no behavior prefix.
+            return
+        if not hasattr(packed_seq, "behavior_global"):
+            raise ValueError("Stage-2 behavior prefix requires behavior_global on PackedSequence")
+        global_features = packed_seq.behavior_global
+        if global_features.ndim != 2 or global_features.shape[0] != packed_seq.behavior_indexes.numel():
+            raise ValueError(
+                "behavior_global must contain one vector per reserved behavior prefix, "
+                f"got {tuple(global_features.shape)} for {packed_seq.behavior_indexes.numel()} slots"
+            )
+        prefix = self.behavior_global_projector(global_features).to(target_dtype)
+        if self.behavior_pim_encoder is not None and not self.behavior_pim_force_bypass:
+            required = (
+                "behavior_phase",
+                "behavior_effect",
+                "behavior_effect_valid",
+                "behavior_pim_phase",
+                "behavior_pim_effect",
+                "behavior_pim_valid",
+            )
+            if any(not hasattr(packed_seq, name) for name in required):
+                raise ValueError("PIM-enabled behavior prefix requires phase, BIT, and PIM tensors")
+            assert self.behavior_pim_projector is not None
+            assert self.behavior_pim_gate is not None
+            brief_interaction_trace = BriefInteractionTrace(
+                effects=packed_seq.behavior_effect,
+                valid=packed_seq.behavior_effect_valid,
+            )
+            causal_prompt_context = self.behavior_pim_encoder(
+                global_features,
+                packed_seq.behavior_phase,
+                brief_interaction_trace.effects,
+                brief_interaction_trace.valid,
+                packed_seq.behavior_pim_phase,
+                packed_seq.behavior_pim_effect,
+                packed_seq.behavior_pim_valid,
+            )
+            # Do not synthesize a residual from BOS slots when no persistent
+            # evidence exists.  The scalar tanh gate is exactly zero at
+            # initialization, preserving the released prefix bit-for-bit while
+            # letting the scalar gate learn first; encoder/projector gradients
+            # become active once that gate leaves zero.
+            causal_prompt = self.behavior_pim_projector(causal_prompt_context)
+            prefix = inject_causal_prompt(
+                prefix, causal_prompt, self.behavior_pim_gate, packed_seq.behavior_pim_valid
+            )
+            # Keep the legacy field for existing diagnostics and expose the
+            # paper term for new integrations.
+            packed_seq.behavior_pim_context = causal_prompt_context
+            packed_seq.causal_prompt = causal_prompt_context
+        packed_sequence[packed_seq.behavior_indexes] = prefix
+
+    def _encode_online_memory_prefix(
+        self,
+        packed_seq: PackedSequence,
+        packed_sequence: torch.Tensor,
+        target_dtype: torch.dtype,
+    ) -> None:
+        """Write one causal prefix from same-task completed transitions."""
+        indexes = packed_seq.online_memory_indexes
+        if not indexes.numel():
+            return
+        if self.behavior_online_projector is None:
+            raise ValueError("Online-memory prefix slots require behavior_stage2.online_memory_enabled=True")
+        if not hasattr(packed_seq, "behavior_online_context"):
+            raise ValueError("Online-memory prefix requires behavior_online_context on PackedSequence")
+        context = packed_seq.behavior_online_context
+        if context.ndim != 2 or context.shape[0] != indexes.numel():
+            raise ValueError(
+                "behavior_online_context must contain one vector per online prefix slot, "
+                f"got {tuple(context.shape)} for {indexes.numel()} slots"
+            )
+        packed_sequence[indexes] = self.behavior_online_projector(context).to(target_dtype)
+
+    def _encode_proprio_prefix(
+        self,
+        packed_seq: PackedSequence,
+        packed_sequence: torch.Tensor,
+        target_dtype: torch.dtype,
+    ) -> None:
+        """Project the robot state into reserved causal prefix slots."""
+        if self.proprio_projector is None:
+            if packed_seq.proprio_indexes.numel():
+                raise ValueError("Packed proprio prefix slots require proprio_condition.enabled=True")
+            return
+        if not packed_seq.proprio_indexes.numel():
+            return
+        if not hasattr(packed_seq, "proprio_condition"):
+            raise ValueError("Proprio prefix requires proprio_condition on PackedSequence")
+        proprio = packed_seq.proprio_condition
+        if proprio.ndim != 2 or proprio.shape[-1] != self.proprio_input_dim:
+            raise ValueError(
+                f"Expected proprio_condition [B,{self.proprio_input_dim}], got {tuple(proprio.shape)}"
+            )
+        prefix = self.proprio_projector(proprio).reshape(-1, self.hidden_size).to(target_dtype)
+        if prefix.shape[0] != packed_seq.proprio_indexes.numel():
+            raise ValueError("Proprio batch/prefix count does not match reserved slots")
+        packed_sequence[packed_seq.proprio_indexes] = prefix
+
+    def _embed_packed_timesteps(self, timesteps: torch.Tensor, packed_seq: PackedSequence) -> torch.Tensor:
+        """Embed noised-token timesteps, reusing work when packing proves they share one scalar."""
+        if packed_seq.uses_single_timestep and timesteps.numel() > 1:
+            timestep = timesteps[:1]  # [1]
+            with torch.autocast("cuda", enabled=True, dtype=torch.float32):
+                timestep_embed = self.time_embedder(timestep)  # [1,hidden_size]
+            # Materialize: expand() aliases storage; in-place ops on a float32 no-op .to() would corrupt all rows.
+            return timestep_embed.expand(timesteps.shape[0], -1).contiguous()  # [N_noisy_frames,hidden_size]
+
+        # Timesteps are computed in FP32 for numerical stability.
+        with torch.autocast("cuda", enabled=True, dtype=torch.float32):
+            return self.time_embedder(timesteps)  # [N_noisy_frames,hidden_size]
+
+    def _encode_vision(
+        self,
+        packed_seq: PackedSequence,
+        packed_sequence: torch.Tensor,
+        target_dtype: torch.dtype,
+    ) -> List[Tuple[int, int, int]] | None:
+        """Project vision tokens and fill into packed_sequence.
+
+        Args:
+            packed_seq: PackedSequence containing vision tokens and metadata.
+            packed_sequence: The packed sequence tensor to fill vision embeddings into (modified in-place).
+            target_dtype: Target dtype for embeddings (typically from text embedding).
+
+        Returns:
+            Original latent shapes before padding (for unpadding during decode), or None if no vision tokens.
+        """
+        if packed_seq.vision is None or packed_seq.vision.tokens is None:
+            # No vision tokens in this batch
+            return None
+
+        vision = packed_seq.vision
+        assert vision.tokens is not None  # Type narrowing (checked above but reassignment loses it)
+        assert vision.token_shapes is not None
+        assert isinstance(vision.sequence_indexes, torch.Tensor)
+        assert isinstance(vision.timesteps, torch.Tensor)
+        torch._assert(
+            vision.timesteps.dtype in (torch.long, torch.float32),
+            f"Timestep must be long/float32, got {vision.timesteps.dtype}",
+        )
+
+        assert isinstance(vision.mse_loss_indexes, torch.Tensor)
+
+        packed_tokens_vision, original_latent_shapes = self.patchify_and_pack_latents(
+            vision.tokens, vision.token_shapes
+        )  # packed_tokens_vision: [total_vision_patches,patch_latent_dim]
+        packed_tokens_vision = self.vae2llm(packed_tokens_vision.to(target_dtype))  # [total_vision_patches,hidden_size]
+
+        has_noisy_vision = vision.mse_loss_indexes.numel() > 0
+
+        if has_noisy_vision:
+            timesteps_vision = vision.timesteps.to(dtype=torch.float32) * self.timestep_scale  # [N_noisy_frames_vision]
+
+            packed_timestep_embeds_vision = self._embed_packed_timesteps(
+                timesteps_vision, packed_seq
+            )  # [N_noisy_frames_vision,hidden_size]
+            packed_timestep_embeds_vision = packed_timestep_embeds_vision.to(
+                target_dtype
+            )  # [N_noisy_frames_vision,hidden_size]
+
+            packed_tokens_vision = _apply_timestep_embeds_to_noisy_tokens(
+                packed_tokens=packed_tokens_vision,
+                packed_timestep_embeds=packed_timestep_embeds_vision,
+                noisy_frame_indexes=vision.noisy_frame_indexes,
+                token_shapes=vision.token_shapes,
+            )  # [total_vision_patches,hidden_size]
+
+        packed_sequence[vision.sequence_indexes] = (
+            packed_tokens_vision  # [total_vision_patches,hidden_size] scattered into [N_total,hidden_size]
+        )
+        return original_latent_shapes
+
+    def _decode_vision(
+        self,
+        packed_seq: PackedSequence,
+        last_hidden_state: torch.Tensor,
+        output_dict: dict,
+        original_latent_shapes: List[Tuple[int, int, int]] | None = None,
+    ) -> None:
+        """Decode vision tokens from hidden states and update output_dict.
+
+        Args:
+            packed_seq: PackedSequence containing mse_loss_indexes_vision and token_shapes_vision.
+            last_hidden_state: Hidden states from the transformer.
+            output_dict: Output dictionary to update with mse_preds (modified in-place).
+            original_latent_shapes: Original latent shapes before padding (for unpadding).
+        """
+        vision = packed_seq.vision
+        # Check if no vision or no noisy vision tokens
+        has_noisy_vision = (
+            vision is not None
+            and vision.tokens is not None
+            and isinstance(vision.mse_loss_indexes, torch.Tensor)
+            and vision.mse_loss_indexes.numel() > 0
+        )
+        if not has_noisy_vision:
+            # No noisy vision tokens present. The model is predicting actions
+            # given clean vision tokens. We need to execute a dummy forward to maintain
+            # computation graph consistency across ranks (FSDP should torch all weights).
+            preds_vision = torch.zeros(
+                [1, self.patch_latent_dim], device=last_hidden_state.device, dtype=last_hidden_state.dtype
+            )  # [1,patch_latent_dim]
+            preds_vision = self.vae2llm(preds_vision)  # [1,hidden_size]
+            preds_vision = self.llm2vae(preds_vision)  # [1,patch_latent_dim]
+            # Return a list of per-sample zero tensors with correct shapes (e.g. (C, T, H, W)),
+            # so downstream code (_get_velocity, _compute_flow_matching_loss) that iterates over preds_vision
+            # gets properly-shaped tensors. Without this, the dummy tensor (1, patch_latent_dim)
+            # would cause a size mismatch when concatenating vision+action velocities.
+            # When vision is None (no vision in batch), fall back to [preds_vision] purely for
+            # gradient graph consistency — it won't be iterated over.
+            if vision is not None and vision.tokens is not None:
+                preds_vision_list = [torch.zeros_like(tok) for tok in vision.tokens]
+                # Inject dummy forward's computation graph so vae2llm/llm2vae params
+                # stay in the autograd graph (zeros_like creates detached tensors).
+                preds_vision_list[0] = preds_vision_list[0] + 0.0 * preds_vision.sum()
+            else:
+                preds_vision_list = [preds_vision]
+            output_dict.update(preds_vision=preds_vision_list)
+        else:
+            assert vision is not None  # Type narrowing
+            assert isinstance(vision.mse_loss_indexes, torch.Tensor)
+            assert vision.noisy_frame_indexes is not None
+            preds_vision = self.llm2vae(
+                last_hidden_state[vision.mse_loss_indexes]
+            )  # [total_noisy_vision_patches,patch_latent_dim]
+            preds_vision = self.unpatchify_and_unpack_latents(
+                preds_vision,
+                token_shapes_vision=vision.token_shapes,
+                noisy_frame_indexes_vision=vision.noisy_frame_indexes,
+                original_latent_shapes=original_latent_shapes,
+            )
+            output_dict.update(preds_vision=preds_vision)
+
+    def _encode_action(
+        self,
+        packed_seq: PackedSequence,
+        packed_sequence: torch.Tensor,
+        target_dtype: torch.dtype,
+    ) -> None:
+        """Encode action tokens and fill into packed_sequence."""
+        if packed_seq.action is None or packed_seq.action.tokens is None:
+            # No action tokens in this batch
+            return
+
+        action: ModalityData = packed_seq.action
+        assert action.token_shapes is not None
+        assert isinstance(action.sequence_indexes, torch.Tensor)
+        assert isinstance(action.timesteps, torch.Tensor)
+        assert isinstance(action.mse_loss_indexes, torch.Tensor)
+
+        # Pack variable-length action tokens into a 1D sequence (same pattern as pack_sound_latents)
+        packed_tokens_action, per_token_domain_id = self.pack_action(
+            action.tokens, action.token_shapes, action.domain_id
+        )
+        packed_tokens_action = self.action2llm(packed_tokens_action, per_token_domain_id)
+
+        packed_tokens_action = packed_tokens_action + self.action_modality_embed.view(
+            1, -1
+        )  # [B_action*T_action,hidden_size]
+
+        has_noisy_actions = has_noisy_tokens(action)
+        if has_noisy_actions:
+            timesteps_action = action.timesteps * self.timestep_scale  # [N_noisy_frames_action]
+            packed_timestep_embeds_action = self._embed_packed_timesteps(
+                timesteps_action, packed_seq
+            )  # [N_noisy_frames_action,hidden_size]
+            packed_timestep_embeds_action = packed_timestep_embeds_action.to(
+                target_dtype
+            )  # [N_noisy_frames_action,hidden_size]
+
+            packed_tokens_action = _apply_timestep_embeds_to_noisy_tokens(
+                packed_tokens=packed_tokens_action,
+                packed_timestep_embeds=packed_timestep_embeds_action,
+                noisy_frame_indexes=action.noisy_frame_indexes,
+                token_shapes=action.token_shapes,
+            )  # [B_action*T_action,hidden_size]
+
+        if self.behavior_pbd is not None and self.behavior_adapter is not None:
+            if not all(hasattr(packed_seq, name) for name in ("behavior_global", "behavior_phase", "behavior_effect", "behavior_effect_valid")):
+                raise ValueError("Zeva requires task context, phase and BIT features")
+            action_lengths = {int(tokens.shape[0]) for tokens in action.tokens}
+            if len(action_lengths) != 1:
+                raise ValueError("Zeva policy injection requires a fixed action horizon within each batch")
+            prior_mean, prior_std = self.behavior_pbd(
+                packed_seq.behavior_global,
+                packed_seq.behavior_phase,
+                packed_seq.behavior_effect,
+                packed_seq.behavior_effect_valid,
+            )
+            residual = self.behavior_adapter(
+                prior_mean,
+                action_length=action_lengths.pop(),
+                leading_condition_steps=self.behavior_prior_leading_condition_steps,
+            )
+            if self.training and self.behavior_prior_dropout_rate:
+                keep = torch.rand((residual.shape[0], 1, 1), device=residual.device) >= self.behavior_prior_dropout_rate
+                residual = residual * keep.to(residual.dtype)
+            elif not self.training:
+                residual = residual * self.behavior_prior_inference_guidance_scale
+            packed_tokens_action = packed_tokens_action + residual.flatten(0, 1).to(packed_tokens_action.dtype)
+            packed_seq.behavior_prior_mean = prior_mean
+            packed_seq.behavior_prior_std = prior_std
+
+        packed_sequence[action.sequence_indexes] = (
+            packed_tokens_action  # [B_action*T_action,hidden_size] scattered into [N_total,hidden_size]
+        )
+
+    def _decode_action(
+        self,
+        packed_seq: PackedSequence,
+        last_hidden_state: torch.Tensor,
+        output_dict: dict,
+    ) -> None:
+        """Decode action tokens from hidden states and update output_dict."""
+        action = packed_seq.action
+        # Shared predicate with OmniMoTModel's has_noisy_actions gating: actions
+        # are decodable targets only when the packer marked action tokens noisy.
+        has_noisy_action = has_noisy_tokens(action)
+        if not has_noisy_action:
+            # dummy forward to maintain computation graph consistency across ranks
+            preds_action = torch.zeros(
+                [1, self.action_dim], device=last_hidden_state.device, dtype=last_hidden_state.dtype
+            )  # [1,action_dim]
+            dummy_domain_id = torch.zeros([1], device=last_hidden_state.device, dtype=torch.long)  # [1]
+            preds_action = self.action2llm(preds_action, dummy_domain_id) + self.action_modality_embed.view(
+                1, -1
+            )  # [1,hidden_size]
+            preds_action = self.llm2action(preds_action, dummy_domain_id)  # [1,action_dim]
+            # Return a list of per-sample zero tensors with correct shapes (e.g. (T, action_dim)),
+            # so downstream code (_get_velocity, _compute_flow_matching_loss) that iterates over preds_action
+            # gets properly-shaped tensors. Without this, the dummy tensor (1, action_dim)
+            # would cause a size mismatch when concatenating vision+action velocities.
+            if action is not None and action.tokens is not None:
+                preds_action_list = [torch.zeros_like(tok) for tok in action.tokens]
+                # Inject dummy forward's computation graph so DomainAwareLinear params
+                # stay in the autograd graph (zeros_like creates detached tensors).
+                preds_action_list[0] = preds_action_list[0] + 0.0 * preds_action.sum()
+            # When action is None (no action in batch), fall back to [preds_action] purely for
+            # gradient graph consistency — it won't be iterated over.
+            else:
+                preds_action_list = [preds_action]
+            output_dict.update(preds_action=preds_action_list)
+        else:
+            assert action is not None  # Type narrowing
+            assert isinstance(action.mse_loss_indexes, torch.Tensor)
+            assert action.condition_mask is not None
+            assert len(action.domain_id) > 0
+
+            action_hidden_states = last_hidden_state[action.mse_loss_indexes]  # [total_noisy_action_tokens,hidden_size]
+
+            # Build per-token domain IDs for the noisy tokens (same expansion logic as pack_action)
+            domain_ids: list[torch.Tensor] = []
+            for nfi, d_id in zip(action.noisy_frame_indexes, action.domain_id):
+                domain_ids.append(d_id.expand(len(nfi)))
+            per_token_domain_id = torch.cat(domain_ids, dim=0)
+
+            preds_action = self.llm2action(
+                action_hidden_states, per_token_domain_id
+            )  # [total_noisy_action_tokens,action_dim]
+            # The frozen base output is always computed first.  R3's
+            # zero-initialized residual expert is always on at inference: the
+            # training-only event metadata must never become a phase/task
+            # router.  The legacy explicit gate remains available only when
+            # ``always_on`` is false, preserving the old experiment's exact
+            # bypass semantics without confusing it with R3.
+            residual_cfg = self.config.residual_expert_config or {}
+            residual_always_on = bool(residual_cfg.get("always_on", False))
+            if self.action_residual_head is not None and (residual_always_on or action.residual_gate is not None):
+                residual_delta = self.action_residual_head(action_hidden_states)
+                arm_limit = float(residual_cfg.get("arm_limit", 0.05))
+                gripper_limit = float(residual_cfg.get("gripper_limit", 0.5))
+                gripper_index = int(residual_cfg.get("gripper_index", 7 if self.action_dim == 8 else 9))
+                if arm_limit < 0 or gripper_limit < 0:
+                    raise ValueError("Residual arm/gripper limits must be non-negative")
+                offset = 0
+                bounded_parts: list[torch.Tensor] = []
+                for sample_index, nfi in enumerate(action.noisy_frame_indexes):
+                    count = int(nfi.numel())
+                    delta_i = residual_delta[offset : offset + count]
+                    if residual_always_on:
+                        gate_i = torch.ones((count, 1), device=delta_i.device, dtype=delta_i.dtype)
+                    else:
+                        gate_i = action.residual_gate[sample_index].to(
+                            device=delta_i.device, dtype=delta_i.dtype
+                        )
+                        gate_i = gate_i.index_select(0, nfi.to(device=gate_i.device)).reshape(-1, 1)
+                    raw_dim = action.raw_action_dim[sample_index] if action.raw_action_dim is not None else None
+                    delta_i = bounded_action_residual(
+                        delta_i,
+                        gate_i,
+                        arm_limit=arm_limit,
+                        gripper_limit=gripper_limit,
+                        gripper_index=gripper_index,
+                        raw_action_dim=raw_dim,
+                    )
+                    bounded_parts.append(delta_i)
+                    offset += count
+                if bounded_parts:
+                    preds_action = preds_action + torch.cat(bounded_parts, dim=0)
+            preds_action = self.unpack_action(preds_action, action.token_shapes, action.noisy_frame_indexes)
+            output_dict.update(preds_action=preds_action)
+
+    def _encode_sound(
+        self,
+        packed_seq: PackedSequence,
+        packed_sequence: torch.Tensor,
+        target_dtype: torch.dtype,
+    ) -> None:
+        """Encode sound tokens and fill into packed_sequence.
+
+        Args:
+            packed_seq: PackedSequence containing sound tokens and metadata.
+            packed_sequence: The packed sequence tensor to fill sound embeddings into (modified in-place).
+            target_dtype: Target dtype for embeddings (typically from text embedding).
+        """
+        if packed_seq.sound is None or packed_seq.sound.tokens is None:
+            # No sound tokens in this batch
+            return
+
+        sound = packed_seq.sound
+        assert sound.token_shapes is not None
+        assert isinstance(sound.sequence_indexes, torch.Tensor)
+        assert isinstance(sound.timesteps, torch.Tensor)
+        assert isinstance(sound.mse_loss_indexes, torch.Tensor)
+
+        # Pack sound latents: list of [C, T] tensors → [total_tokens, C]
+        packed_tokens_sound = self.pack_sound_latents(
+            sound.tokens, sound.token_shapes
+        )  # [total_sound_tokens,sound_dim]
+        packed_tokens_sound = packed_tokens_sound.to(target_dtype)  # [total_sound_tokens,sound_dim]
+
+        # Project sound tokens + modality embedding. Position info comes from
+        # mRoPE position IDs in the attention layers.
+        packed_tokens_sound = (
+            self.sound2llm(packed_tokens_sound) + self.sound_modality_embed
+        )  # [total_sound_tokens,hidden_size]
+
+        has_noisy_sound = sound.mse_loss_indexes.numel() > 0
+        if has_noisy_sound:
+            timesteps_sound = sound.timesteps * self.timestep_scale  # [N_noisy_frames_sound]
+            packed_timestep_embeds_sound = self._embed_packed_timesteps(
+                timesteps_sound, packed_seq
+            )  # [N_noisy_frames_sound,hidden_size]
+            packed_timestep_embeds_sound = packed_timestep_embeds_sound.to(
+                target_dtype
+            )  # [N_noisy_frames_sound,hidden_size]
+
+            packed_tokens_sound = _apply_timestep_embeds_to_noisy_tokens(
+                packed_tokens=packed_tokens_sound,
+                packed_timestep_embeds=packed_timestep_embeds_sound,
+                noisy_frame_indexes=sound.noisy_frame_indexes,
+                token_shapes=sound.token_shapes,
+            )  # [total_sound_tokens,hidden_size]
+
+        packed_sequence[sound.sequence_indexes] = (
+            packed_tokens_sound  # [total_sound_tokens,hidden_size] scattered into [N_total,hidden_size]
+        )
+
+    def _decode_sound(
+        self,
+        packed_seq: PackedSequence,
+        last_hidden_state: torch.Tensor,
+        output_dict: dict,
+    ) -> None:
+        """Decode sound tokens from hidden states and update output_dict.
+
+        Args:
+            packed_seq: PackedSequence containing sound modality data.
+            last_hidden_state: Hidden states from the transformer.
+            output_dict: Output dictionary to update with preds_sound (modified in-place).
+        """
+        sound = packed_seq.sound
+        # Check if no sound or no noisy sound tokens
+        has_noisy_sound = (
+            sound is not None
+            and sound.tokens is not None
+            and isinstance(sound.mse_loss_indexes, torch.Tensor)
+            and sound.mse_loss_indexes.numel() > 0
+        )
+        if not has_noisy_sound:
+            # dummy forward to maintain computation graph consistency across ranks
+            preds_sound = torch.zeros(
+                [1, self.sound_dim], device=last_hidden_state.device, dtype=last_hidden_state.dtype
+            )  # [1,sound_dim]
+            preds_sound = self.sound2llm(preds_sound) + self.sound_modality_embed  # [1,hidden_size]
+            preds_sound = self.llm2sound(preds_sound)  # [1,sound_dim]
+            if sound is not None and sound.tokens is not None:
+                preds_sound_list = [torch.zeros_like(tok) for tok in sound.tokens]
+                preds_sound_list[0] = preds_sound_list[0] + 0.0 * preds_sound.sum()
+            else:
+                preds_sound_list = [preds_sound]
+            output_dict.update(preds_sound=preds_sound_list)
+        else:
+            assert sound is not None  # Type narrowing
+            assert isinstance(sound.mse_loss_indexes, torch.Tensor)
+            assert sound.condition_mask is not None
+            preds_sound = self.llm2sound(
+                last_hidden_state[sound.mse_loss_indexes]
+            )  # [total_noisy_sound_tokens,sound_dim]
+            preds_sound = self.unpack_sound_latents(
+                preds_sound, sound.token_shapes, sound.noisy_frame_indexes
+            )  # list of [C,T] per sample
+            output_dict.update(preds_sound=preds_sound)
+
+    def forward(
+        self,
+        packed_seq: PackedSequence,
+        memory: MemoryState | None = None,
+        video_temporal_causal: bool | None = None,
+    ) -> dict:
+        """
+        Forward pass for Cosmos3VFMNetwork.
+
+        Args:
+            packed_seq: PackedSequence containing all packed tensors and metadata.
+                See PackedSequence dataclass for field details.
+            memory: Optional MemoryState for persistent KV-cache memory
+                (AR inference or rolling-KV-cache training).  Built by
+                ``OmniMoTModel.build_memory_state()``.
+            video_temporal_causal: Per-call attention-mode override; ``None``
+                (default) uses the config-selected ``self.video_temporal_causal``.
+
+        Returns:
+            dict with keys:
+                - "preds_vision": list[Tensor[C,T,H,W]], one per sample.
+                - "preds_action": Velocity predictions for action tokens (if action_gen).
+                - "preds_sound": Velocity predictions for sound tokens (if sound_gen).
+                - "last_hidden_state": Last hidden state from the transformer.
+                - "lbl_metadata_*": Load balancing metadata.
+                - "ce_preds": Cross-entropy predictions (if predict_text_tokens is True).
+        """
+        # Note: During inference with @torch.no_grad(), model may be in training mode
+        # This is intentional for proper batch norm / dropout behavior
+        # assert self.training, "Cosmos3VFMNetwork only supports training mode"
+
+        packed_sequence, target_dtype = self._encode_text(packed_seq)  # packed_sequence: [N_total,hidden_size]
+        self._encode_behavior_prefix(packed_seq, packed_sequence, target_dtype)
+        self._encode_online_memory_prefix(packed_seq, packed_sequence, target_dtype)
+        self._encode_proprio_prefix(packed_seq, packed_sequence, target_dtype)
+
+        # encode vision tokens
+        original_latent_shapes: List[Tuple[int, int, int]] | None = None
+        if self.config.vision_gen:
+            original_latent_shapes = self._encode_vision(packed_seq, packed_sequence, target_dtype)
+
+        # encode action tokens
+        if self.config.action_gen:
+            self._encode_action(packed_seq, packed_sequence, target_dtype)
+
+        # encode sound tokens
+        if self.config.sound_gen:
+            self._encode_sound(packed_seq, packed_sequence, target_dtype)
+
+        assert packed_seq.attn_modes is not None
+        assert packed_seq.split_lens is not None
+        use_video_temporal_causal = (
+            self.video_temporal_causal if video_temporal_causal is None else video_temporal_causal
+        )
+
+        # Get all generation sequence indexes for MoE routing
+        # IMPORTANT: Include ALL latent tokens (video + action + sound), not just generation targets.
+        # Condition tokens still need to be routed to diffusion experts; they are excluded from
+        # LOSS computation, not from routing.
+        all_gen_indexes = []
+        if packed_seq.vision is not None:
+            assert packed_seq.vision.token_shapes is not None
+            assert isinstance(packed_seq.vision.sequence_indexes, torch.Tensor)
+            all_gen_indexes.append(packed_seq.vision.sequence_indexes)
+        if packed_seq.action is not None and isinstance(packed_seq.action.sequence_indexes, torch.Tensor):
+            all_gen_indexes.append(packed_seq.action.sequence_indexes)
+        if packed_seq.sound is not None and isinstance(packed_seq.sound.sequence_indexes, torch.Tensor):
+            all_gen_indexes.append(packed_seq.sound.sequence_indexes)
+        vision_sequence_indexes = torch.cat(all_gen_indexes, dim=0) if all_gen_indexes else None  # [N_gen_tokens]
+
+        # When temporal causal is enabled the buffer is [action_t0, vision_t0, action_t1, vision_t1, ...].
+        # After torch.cat([vision_indexes, action_indexes]) the interleaved order is lost; sorting restores it.
+        if self.video_temporal_causal:
+            assert packed_seq.sound is None, "Sound generation is not supported with video_temporal_causal=True."
+            if vision_sequence_indexes is not None:
+                vision_sequence_indexes = vision_sequence_indexes.sort().values  # [N_gen_tokens]
+
+        vision_token_shapes = packed_seq.vision.token_shapes if packed_seq.vision else None
+
+        # The packer is the single source of truth for the supertoken layout.
+        # ``num_action_tokens_per_supertoken`` is stamped onto ``packed_seq`` by
+        # ``pack_supertokens_temporal_causal`` (= tcf when actions are packed
+        # inline, 0 otherwise) and read unchanged by the attention builder, the
+        # NATTEN metadata generator, and the rolling KV-cache state — keeping
+        # all downstream supertoken geometry automatically in sync with the pack.
+        num_action_tokens_per_supertoken = packed_seq.num_action_tokens_per_supertoken
+
+        replicated_attention_io_cp = (
+            self.attention_io_layout == "replicated"
+            and self.parallel_dims is not None
+            and self.parallel_dims.cp_enabled
+        )
+        # ``sequence_sharded`` attention I/O shards the token sequence, so
+        # packing must pad sequence lengths to the CP size and the input/output
+        # sequence helpers need the CP mesh.  ``replicated`` attention I/O keeps
+        # current-frame sequences replicated and uses the CP mesh later inside
+        # attention to slice local heads, so the effective sequence-sharding
+        # world size is 1 here.
+        sequence_shard_parallel_dims = None if replicated_attention_io_cp else self.parallel_dims
+        sequence_shard_world_size = (
+            1 if replicated_attention_io_cp else (self.parallel_dims.cp_size if self.parallel_dims else 1)
+        )
+        prepared_sequence_pack_metadata = packed_seq.get_sequence_pack_metadata()
+
+        packed_understanding_indexes = torch.cat(
+            (
+                packed_seq.behavior_indexes,
+                packed_seq.online_memory_indexes,
+                packed_seq.proprio_indexes,
+                packed_seq.text_indexes,
+            )
+        ).sort().values
+        input_pack, attention_meta, natten_metadata_list = build_packed_sequence(
+            self.config.joint_attn_implementation,
+            packed_sequence=packed_sequence,
+            attn_modes=packed_seq.attn_modes,
+            split_lens=packed_seq.split_lens,
+            sample_lens=packed_seq.sample_lens,
+            packed_und_token_indexes=packed_understanding_indexes,
+            packed_gen_token_indexes=vision_sequence_indexes,
+            num_heads=self.num_heads,
+            is_image_batch=packed_seq.is_image_batch,
+            head_dim=self.head_dim,
+            num_layers=self.num_hidden_layers,
+            token_shapes=packed_seq.vision.token_shapes,
+            natten_parameter_list=self.natten_parameter_list,
+            cp_world_size=sequence_shard_world_size,
+            video_temporal_causal=use_video_temporal_causal,
+            skip_natten_metadata=memory is not None and not memory.requires_natten_metadata(),
+            vision_token_shapes=vision_token_shapes,
+            action_token_shapes=packed_seq.action.token_shapes if packed_seq.action else None,
+            num_action_tokens_per_supertoken=num_action_tokens_per_supertoken,
+            null_action_supertokens=packed_seq.null_action_supertokens,
+            pad_for_cuda_graphs=self.pad_for_cuda_graphs,
+            prepared_metadata=prepared_sequence_pack_metadata,
+        )
+
+        # ── Multi-control transfer: annotate SplitInfo with per-item ranges ──────
+        # Activated only when packed_seq carries control_weights, i.e. the caller
+        # has set up a multi-control batch via build_transfer_batch.
+        #
+        # multi_control_two_way_attention runs N independent maskless SDPA passes,
+        # one per control.  For each pass i, KV = [text | ctrl_i | noisy].
+        # The final noisy output is the weighted sum of the N pass outputs:
+        #   noisy_out = w_1 * noisy_out_1 + ... + w_N * noisy_out_N
+        # All SDPA calls are maskless → Flash Attention always active.
+        # N=1, w=1.0 → identical to two_way_attention.
+        #
+        # CP compatibility: control_stream_token_ranges are gen-relative global
+        # offsets computed here, before CP sharding.  Ulysses CP restores the full
+        # sequence on every rank (via all-to-all) before calling dispatch_attention,
+        # so the global ranges are valid indices inside multi_control_two_way_attention.
+        if (
+            isinstance(attention_meta, SplitInfo)
+            and packed_seq.control_weights is not None
+            and packed_seq.vision_item_split_lens
+        ):
+            # For multi-control, each sample must have N controls + 1 noisy item
+            # (items 0..N-2 are controls, item N-1 is the noisy target).
+            # Only batch_size=1 is supported; assert to catch misuse early.
+            assert len(packed_seq.vision_item_split_lens) == 1, (
+                f"Multi-control transfer requires batch_size=1, got {len(packed_seq.vision_item_split_lens)} samples."
+            )
+            item_lens = packed_seq.vision_item_split_lens[0]  # [L_ctrl0, L_ctrl1, ..., L_noisy]
+            weights = packed_seq.control_weights[0]  # [w_ctrl0, w_ctrl1, ...]
+            assert len(item_lens) > 1, (
+                f"Multi-control requires at least 1 control + 1 noisy item; got vision_item_split_lens={item_lens}."
+            )
+            assert len(weights) == len(item_lens) - 1, (
+                f"control_weights length ({len(weights)}) must equal number of control items ({len(item_lens) - 1})."
+            )
+            ctrl_ranges: list[tuple[int, int]] = []
+            cursor = 0
+            for lens in item_lens[:-1]:  # all but last = control streams
+                ctrl_ranges.append((cursor, cursor + lens))
+                cursor += lens
+            noisy_range = (cursor, cursor + item_lens[-1])
+            n_gen = int(vision_sequence_indexes.shape[0]) if vision_sequence_indexes is not None else 0
+            assert noisy_range[1] == n_gen, (
+                f"vision_item_split_lens sums to {noisy_range[1]} gen tokens but packed tensor has "
+                f"{n_gen}; packing inconsistency detected."
+            )
+            attention_meta.control_stream_token_ranges = ctrl_ranges
+            attention_meta.noisy_token_range = noisy_range
+            attention_meta.control_weights = weights
+
+        input_pack, packed_position_ids = get_context_parallel_sharded_sequence(
+            attn_implementation=self.config.joint_attn_implementation,
+            input_pack=input_pack,
+            position_ids=packed_seq.position_ids,
+            parallel_dims=sequence_shard_parallel_dims,
+        )
+
+        packed_outputs, lbl_metadata = self.language_model(
+            input_pack,
+            attention_mask=attention_meta,
+            position_ids=packed_position_ids,
+            natten_metadata_list=natten_metadata_list,
+            memory=memory,
+        )
+        last_hidden_state = get_context_parallel_last_hidden_state(
+            packed_outputs=packed_outputs,
+            parallel_dims=sequence_shard_parallel_dims,
+        )  # [N_total,hidden_size]
+        output_dict = dict()
+
+        # decode vision tokens
+        if self.config.vision_gen:
+            self._decode_vision(packed_seq, last_hidden_state, output_dict, original_latent_shapes)
+
+        # decode action tokens
+        if self.config.action_gen:
+            self._decode_action(packed_seq, last_hidden_state, output_dict)
+        if self.behavior_pbd is not None and hasattr(packed_seq, "behavior_prior_mean"):
+            output_dict["behavior_prior_mean"] = packed_seq.behavior_prior_mean
+            output_dict["behavior_prior_std"] = packed_seq.behavior_prior_std
+
+        # decode sound tokens
+        if self.config.sound_gen:
+            self._decode_sound(packed_seq, last_hidden_state, output_dict)
+
+        output_dict.update(last_hidden_state=last_hidden_state)
+        for lbl_metadata_key, lbl_metadata_value in lbl_metadata.items():
+            output_dict.update({f"lbl_metadata_{lbl_metadata_key}": lbl_metadata_value})
+        if self.predict_text_tokens:
+            packed_ce_preds = self.language_model.lm_head(
+                last_hidden_state[packed_seq.ce_loss_indexes]
+            )  # [N_ce_tokens,vocab_size]
+            output_dict["ce_preds"] = packed_ce_preds
+
+        return output_dict
+
+
+def _apply_timestep_embeds_to_noisy_tokens(
+    packed_tokens: torch.Tensor,
+    packed_timestep_embeds: torch.Tensor,
+    noisy_frame_indexes: List[torch.Tensor],
+    token_shapes: list[tuple[int, ...]],
+) -> torch.Tensor:
+    """Apply timestep embeddings to noisy tokens.
+    Tn is the number of noisy frames for a given sample.
+    Tc is the number of clean frames for a given sample.
+    T is the total number of frames for a given sample.
+    T = Tn + Tc
+
+    Args:
+        packed_tokens: The packed tokens to apply timestep embeddings to.
+        packed_timestep_embeds: The packed timestep embeddings to apply.
+        noisy_frame_indexes: The frame indices to apply timestep embeddings to
+            (list of tensors, each with shape (Tn,)).
+        token_shapes: The token shapes for each sample. Each entry is a tuple
+            shaped like ``(T, ...)`` where trailing dimensions represent the spatial grid.
+
+    Returns:
+        The packed tokens with timestep embeddings applied to the noisy tokens.
+    """
+
+    # Handle variable token shapes by processing each sample's noisy_frame_indexes individually.
+    # The noisy indices are first expanded to cover the entire spatial grid of each frame.
+    #
+    # For video frames, the spatial grid is (H, W).
+    # For action frames, the spatial grid is ().
+    # For sound frames, the spatial grid is (1, 1).
+    #
+    # The noisy indices are then flattened into a single tensor overall. When flattening,
+    # we must ensure that the noisy indices from each sample are unique by adding the
+    # cumulative sum of the token shapes of previous samples to the noisy indices for
+    # a given sample.
+    start_noisy_index = 0
+    flattened_noisy_frame_indexes = []
+
+    for noisy_indexes_i, token_shape_i in zip(noisy_frame_indexes, token_shapes):
+        assert noisy_indexes_i.numel() <= token_shape_i[0]
+        spatial_numel_i = math.prod(token_shape_i[1:])
+        spatial_indexes_i = torch.arange(spatial_numel_i, device=packed_tokens.device)  # [spatial_numel_i]
+        noisy_indexes_i = (
+            (noisy_indexes_i * spatial_numel_i).unsqueeze(-1).expand(-1, spatial_numel_i)
+        )  # [Tn_i,spatial_numel_i]
+        noisy_indexes_i = noisy_indexes_i.clone() + spatial_indexes_i + start_noisy_index  # [Tn_i,spatial_numel_i]
+        flattened_noisy_frame_indexes.append(noisy_indexes_i.flatten())  # [Tn_i*spatial_numel_i]
+        start_noisy_index += math.prod(token_shape_i)
+
+    flattened_noisy_frame_indexes = torch.cat(flattened_noisy_frame_indexes, dim=0)  # [total_noisy_patches]
+
+    assert packed_tokens.dim() == 2
+    assert packed_timestep_embeds.dim() == 2
+    assert packed_timestep_embeds.shape[1] == packed_tokens.shape[1]
+    assert packed_timestep_embeds.shape[0] <= packed_tokens.shape[0]
+    assert flattened_noisy_frame_indexes.dim() == 1
+    assert flattened_noisy_frame_indexes.shape[0] == packed_timestep_embeds.shape[0]
+
+    flattened_noisy_frame_indexes = flattened_noisy_frame_indexes.unsqueeze(-1).expand(
+        -1,
+        packed_tokens.shape[1],
+    )  # [total_noisy_patches,hidden_size]
+
+    return packed_tokens.scatter_add(
+        dim=0,
+        index=flattened_noisy_frame_indexes,
+        src=packed_timestep_embeds,
+    )  # [total_tokens,hidden_size]

releases/bridge-native-v3/code/cosmos_framework/model/generator/utils/data_and_condition.py

@@ -0,0 +1,182 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""
+Unified data and condition interface where we save the tokenized states and/or
+noised latent states for diffusion/flow-matching training.
+Used for the VFM generation model.
+"""
+
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass(slots=True)
+class GenerationDataClean:
+    """
+    Container for tokenized states and conditioning info (clean states)
+    for the multi-modal (vision, sound, action) MoT training.
+    Used for the VFM generation model.
+    """
+
+    batch_size: int
+    # Vision (list of per-sample tensors)
+    is_image_batch: bool
+    raw_state_vision: list[torch.Tensor] | None = None  # raw state in pixel space
+    x0_tokens_vision: list[torch.Tensor] | None = None  # tokenized latent state
+    fps_vision: torch.Tensor | None = None
+    temporal_positions_vision: list[torch.Tensor] | None = None  # one [T] tensor per vision latent item
+
+    # Image editing: number of vision items per sample.
+    # When set, x0_tokens_vision is a flat list of individually-encoded image latents
+    # (e.g. [src1, tgt1, src2, tgt2, ...]) and this field records how many items belong
+    # to each sample (e.g. [2, 2, ...]).  None for standard T2I/T2V (one item per sample).
+    num_vision_items_per_sample: list[int] | None = None
+
+    # Audio (Sound)
+    raw_state_sound: torch.Tensor | None = None
+    x0_tokens_sound: torch.Tensor | None = None
+    fps_sound: torch.Tensor | None = None
+
+    # Action (dense list of per-sample tensors, only action-having samples)
+    raw_state_action: list[torch.Tensor] | None = None
+    x0_tokens_action: list[torch.Tensor] | None = None
+    fps_action: torch.Tensor | None = None
+    action_domain_id: list[torch.Tensor] | None = None  # per-sample domain IDs, None when no action samples
+    raw_action_dim: list[torch.Tensor] | None = None  # raw action dimension, used adding masks to loss calculation
+    # Dense per-action-sample gate for the residual expert.  Values are in
+    # [0,1] and aligned with x0_tokens_action; absent means exact base bypass.
+    action_residual_gate: list[torch.Tensor] | None = None
+
+    # Optional per-sample low-dimensional state used as causal prefix tokens.
+    proprio_condition: list[torch.Tensor] | None = None
+
+    # Multi-control transfer: per-sample list of per-control weights.
+    # Shape: [num_samples], each element is a list of floats (one per control stream).
+    # None for non-transfer or single-control samples.
+    control_weights: list[list[float]] | None = None
+
+
+@dataclass(slots=True)
+class GenerationDataNoised:
+    """Container for states after noise addition, along with other
+    helper attributes for the flow-matching (gt velocity and noise)
+    for the multi-modal (vision, sound, action) MoT training.
+    Used for the VFM generation model.
+    """
+
+    batch_size: int
+    # Vision
+    epsilon_vision: torch.Tensor  # unit gaussian noise tensor
+    xt_tokens_vision: torch.Tensor  # tokens added with noise level t per flow-matching formulation
+    vt_target_vision: torch.Tensor  # gt rectified flow field
+    sigmas_vision: torch.Tensor | None = None  # SNR to add to the vision tokens
+
+    # Audio (Sound)
+    epsilon_sound: torch.Tensor | None = None
+    xt_tokens_sound: torch.Tensor | None = None
+    vt_target_sound: torch.Tensor | None = None
+    sigmas_sound: torch.Tensor | None = None
+
+    # Action
+    epsilon_action: torch.Tensor | None = None
+    xt_tokens_action: torch.Tensor | None = None
+    vt_target_action: torch.Tensor | None = None
+    sigmas_action: torch.Tensor | None = None
+    raw_action_dim: list[torch.Tensor] | None = None  # raw action dimension, used adding masks to loss calculation
+
+
+def unwrap_and_densify(raw: list | torch.Tensor | None, to_kwargs: dict) -> list[torch.Tensor] | None:
+    """Unwrap nested single-element lists and filter ``None`` entries.
+
+    The joint dataloader can produce data as nested single-element lists
+    (e.g. ``[[t1], [None], [t2]]``).  This helper flattens the nesting,
+    drops ``None`` entries, and moves the remaining tensors to the target
+    device/dtype.
+
+    Args:
+        raw: The raw value from ``data_batch``.  May be ``None``, a bare
+            tensor, or a (possibly nested) list of tensors / ``None`` s.
+            Each tensor in the list has shape ``(...)``.
+        to_kwargs: Keyword arguments forwarded to ``torch.Tensor.to``
+            (e.g. ``{"device": "cuda"}`` or ``{"device": "cuda", "dtype": torch.bfloat16}``).
+
+    Returns:
+        A dense list of device tensors each with shape ``(...)``, or ``None``
+        if the input is ``None`` or every entry is ``None``.
+
+    Examples:
+        >>> unwrap_and_densify([[t1], [None], [t2]], {"device": "cuda"})
+        [t1.cuda(), t2.cuda()]
+        >>> unwrap_and_densify(None, {"device": "cuda"})
+        None
+    """
+    if raw is None:
+        return None
+    if not isinstance(raw, list):
+        return [raw.to(**to_kwargs)]  # list of 1 tensor: [(...)]
+    # Unwrap single-element inner lists: [[t], [None]] -> [t, None]
+    if len(raw) > 0 and isinstance(raw[0], list):
+        raw = [item[0] if isinstance(item, list) else item for item in raw]
+    # Filter None entries and move to device
+    dense = [x.to(**to_kwargs) for x in raw if x is not None]  # list of B tensors: [(...), ...]
+    return dense if dense else None
+
+
+def _expand_per_sample_to_per_vision_item(
+    tensor: torch.Tensor,  # [B,...]
+    num_vision_items_per_sample: list[int] | None,
+) -> torch.Tensor:  # [N_vision_items,...]
+    """Expand a per-sample tensor to a per-vision-item tensor.
+
+    For image editing, each sample may contribute multiple vision items
+    (e.g. source + target).  This helper repeats each sample's value for
+    all of its vision items so that downstream indexing by vision-item
+    position works correctly.
+
+    Args:
+        tensor: Per-sample tensor of shape ``(N, ...)``.
+        num_vision_items_per_sample: Number of vision items per sample,
+            e.g. ``[2, 2, ...]``.  If ``None``, the tensor is returned as-is
+            (standard single-item-per-sample case).
+
+    Returns:
+        Tensor of shape ``(sum(num_vision_items_per_sample), ...)``, or the
+        original tensor when ``num_vision_items_per_sample`` is ``None``.
+    """
+    if num_vision_items_per_sample is None:
+        return tensor  # [B,...]
+    expanded = []
+    for sample_idx, num_items in enumerate(num_vision_items_per_sample):
+        for _ in range(
+            num_items
+        ):  # torch.stack(tensor[idx].repeat(num_vision_items_per_sample[idx]) for idx in range(len(num_vision_items_per_sample)))
+            expanded.append(tensor[sample_idx])  # [...]
+    return torch.stack(expanded)  # [N_vision_items,...]
+
+
+def build_dense_sound_schedule(
+    sequence_plans: list,
+    x0_tokens_sound: list[torch.Tensor] | None,
+    timesteps: torch.Tensor,  # [B,...]
+    sigmas: torch.Tensor,  # [B,...]
+) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+    """Reindex per-sample schedules to match the dense sound tensor list.
+
+    Sound tensors are dense over samples with ``has_sound=True``, while input
+    timesteps/sigmas are indexed by original batch position. This helper maps
+    dense sound entry ``i`` back to its source sample's schedule row.
+    """
+    sound_sample_indices = [i for i, plan in enumerate(sequence_plans) if getattr(plan, "has_sound", False)]
+    num_sound_tensors = 0 if x0_tokens_sound is None else len(x0_tokens_sound)
+    assert len(sound_sample_indices) == num_sound_tensors, (
+        "Sound tensor count must match sequence plans with has_sound=True. "
+        f"Got {num_sound_tensors} sound tensor(s) for {len(sound_sample_indices)} sound plan(s)."
+    )
+
+    if not sound_sample_indices:
+        return None, None
+
+    idx_sound = torch.tensor(sound_sample_indices, dtype=torch.long, device=timesteps.device)  # [n_sound]
+    return timesteps[idx_sound], sigmas[idx_sound]  # [n_sound,...], [n_sound,...]

releases/bridge-native-v3/code/cosmos_framework/model/zeva/__init__.py

@@ -0,0 +1,52 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 Z-Trans CORPORATION. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""Zeva causal-interaction modules required by the released inference path."""
+
+from .brief_interaction_trace import BriefInteractionTrace
+from .causal_prompt import CausalPromptConfig, CausalPromptEncoder, inject_causal_prompt
+from .causal_transition_encoder import CausalTransitionEncoder, CausalTransitionEncoderConfig
+from .checkpoint_io import normalize_cte_state_dict
+from .cte_losses import CTELossConfig, causal_transition_encoder_loss
+from .persistent_interaction_memory import (
+    PIMMemoryEntry,
+    PersistentInteractionMemory,
+    PersistentInteractionMemoryConfig,
+)
+from .phase_conditioned_retrieval import PhaseConditionedPIMRetrieval
+from .policy_injection import (
+    CausalPromptPolicyAdapter,
+    PolicyInjectionConfig,
+    PolicyInjectionPrior,
+    gaussian_prior_nll,
+)
+from .static_task_context_retrieval import (
+    StaticTaskContextRetrievalConfig,
+    StaticTaskContextRetrievalHead,
+    bidirectional_supervised_contrastive_loss,
+    retrieve_static_task_context,
+)
+
+__all__ = [
+    "CausalTransitionEncoder",
+    "CausalTransitionEncoderConfig",
+    "CTELossConfig",
+    "causal_transition_encoder_loss",
+    "BriefInteractionTrace",
+    "PolicyInjectionConfig",
+    "PolicyInjectionPrior",
+    "CausalPromptPolicyAdapter",
+    "gaussian_prior_nll",
+    "StaticTaskContextRetrievalConfig",
+    "StaticTaskContextRetrievalHead",
+    "bidirectional_supervised_contrastive_loss",
+    "retrieve_static_task_context",
+    "PIMMemoryEntry",
+    "PersistentInteractionMemory",
+    "PersistentInteractionMemoryConfig",
+    "PhaseConditionedPIMRetrieval",
+    "CausalPromptConfig",
+    "CausalPromptEncoder",
+    "inject_causal_prompt",
+    "normalize_cte_state_dict",
+]

releases/bridge-native-v3/code/cosmos_framework/model/zeva/attempt_artifacts.py

@@ -0,0 +1,77 @@
+"""Validation and manifests for fixed-seed PIM attempt artifacts."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from pathlib import Path
+from typing import Any, Iterable
+
+
+ATTEMPT_RECORD_VERSION = "zeva_pim_attempt_v1"
+CHUNK_RECORD_VERSION = "zeva_pim_chunk_v1"
+
+
+def sha256_file(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        for block in iter(lambda: handle.read(1024 * 1024), b""):
+            digest.update(block)
+    return digest.hexdigest()
+
+
+def validate_attempt_record(record: dict[str, Any]) -> None:
+    required = {
+        "schema_version",
+        "session_id",
+        "task",
+        "attempt_id",
+        "environment_seed",
+        "success",
+        "total_steps",
+        "chunk_records",
+        "artifact_npz",
+        "artifact_sha256",
+    }
+    missing = required.difference(record)
+    if missing:
+        raise ValueError(f"attempt record is missing fields: {sorted(missing)}")
+    if record["schema_version"] != ATTEMPT_RECORD_VERSION:
+        raise ValueError(f"unsupported attempt schema: {record['schema_version']!r}")
+    if not isinstance(record["session_id"], str) or not record["session_id"]:
+        raise ValueError("session_id must be a non-empty string")
+    if not isinstance(record["task"], str) or not record["task"]:
+        raise ValueError("task must be a non-empty string")
+    if int(record["attempt_id"]) < 0 or int(record["environment_seed"]) < 0:
+        raise ValueError("attempt_id and environment_seed must be non-negative")
+    if not isinstance(record["success"], bool):
+        raise ValueError("success must be boolean")
+    if int(record["total_steps"]) < 0:
+        raise ValueError("total_steps must be non-negative")
+    chunks = record["chunk_records"]
+    if not isinstance(chunks, list):
+        raise ValueError("chunk_records must be a list")
+    for index, chunk in enumerate(chunks):
+        if chunk.get("schema_version") != CHUNK_RECORD_VERSION:
+            raise ValueError(f"unsupported chunk schema at index {index}")
+        if int(chunk.get("replan_index", -1)) != index:
+            raise ValueError("chunk replan indices must be contiguous from zero")
+        start, end = int(chunk.get("start_step", -1)), int(chunk.get("end_step", -1))
+        if start < 0 or end < start or int(chunk.get("executed_count", -1)) != end - start:
+            raise ValueError(f"invalid chunk interval at index {index}")
+    artifact = str(record["artifact_npz"])
+    checksum = str(record["artifact_sha256"])
+    if not artifact.endswith(".npz") or len(checksum) != 64:
+        raise ValueError("invalid attempt artifact reference")
+
+
+def write_manifest(path: Path, records: Iterable[dict[str, Any]]) -> None:
+    rows = list(records)
+    for row in rows:
+        validate_attempt_record(row)
+    payload = {
+        "schema_version": ATTEMPT_RECORD_VERSION,
+        "attempts": len(rows),
+        "records": rows,
+    }
+    path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")

releases/bridge-native-v3/code/cosmos_framework/model/zeva/attempt_artifacts_test.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+from cosmos_framework.model.zeva.attempt_artifacts import (
+    ATTEMPT_RECORD_VERSION,
+    CHUNK_RECORD_VERSION,
+    validate_attempt_record,
+)
+
+
+def test_attempt_artifact_schema() -> None:
+    record = {
+        "schema_version": ATTEMPT_RECORD_VERSION,
+        "session_id": "kettle:197",
+        "task": "TurnOnElectricKettle",
+        "attempt_id": 0,
+        "environment_seed": 197,
+        "success": False,
+        "total_steps": 16,
+        "chunk_records": [
+            {
+                "schema_version": CHUNK_RECORD_VERSION,
+                "replan_index": 0,
+                "start_step": 0,
+                "end_step": 16,
+                "executed_count": 16,
+            }
+        ],
+        "artifact_npz": "attempt_000_seed_197.npz",
+        "artifact_sha256": "0" * 64,
+    }
+    validate_attempt_record(record)

releases/bridge-native-v3/code/cosmos_framework/model/zeva/attempt_protocol.py

@@ -0,0 +1,78 @@
+"""Frozen repeated-attempt protocol for RoboCasa Atomic-5 evaluation.
+
+This module intentionally has no simulator or model dependencies so clients,
+servers, tests, and result auditors can share one exact protocol definition.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+PROTOCOL_VERSION = "zeva_robocasa_atomic5_v1"
+PREDICTED_ACTION_HORIZON = 32
+EXECUTED_ACTION_HORIZON = 16
+EFFECT_ACTION_HORIZON = 16
+CTE_TRANSITION_ACTION_HORIZON = 4
+EFFECT_CTE_TRANSITIONS = 4
+MAX_CONTROLS_PER_ATTEMPT = 300
+REQUIRED_WARMUP_REQUESTS = 1
+
+
+def inference_seed(base_seed: int, attempt_id: int, replan_index: int) -> int:
+    """Return a reproducible per-request seed shared by paired conditions."""
+    if base_seed < 0 or attempt_id < 0 or replan_index < 0:
+        raise ValueError("base_seed, attempt_id, and replan_index must be non-negative")
+    # A RoboCasa attempt has at most ceil(300/16)=19 replans.  Reserving 1,000
+    # values per attempt makes the mapping readable and collision-free.
+    if replan_index >= 1_000:
+        raise ValueError("replan_index must be < 1000")
+    value = int(base_seed) + 1_000 * int(attempt_id) + int(replan_index)
+    if value >= 2**32:
+        raise ValueError("derived inference seed exceeds uint32")
+    return value
+
+
+@dataclass(frozen=True)
+class AttemptSessionKey:
+    """Identity of memory that may survive an environment reset."""
+
+    session_id: str
+    task_cluster: str
+    environment_seed: int
+
+    def __post_init__(self) -> None:
+        if not self.session_id:
+            raise ValueError("session_id must be non-empty")
+        if not self.task_cluster:
+            raise ValueError("task_cluster must be non-empty")
+        if self.environment_seed < 0:
+            raise ValueError("environment_seed must be non-negative")
+
+    def assert_compatible(self, other: "AttemptSessionKey") -> None:
+        if self != other:
+            raise ValueError(
+                "attempt-session mismatch: memory may persist only for the exact "
+                f"(session_id, task, seed); existing={self}, requested={other}"
+            )
+
+
+def default_session_id(task_cluster: str, environment_seed: int) -> str:
+    if not task_cluster:
+        raise ValueError("task_cluster must be non-empty")
+    if environment_seed < 0:
+        raise ValueError("environment_seed must be non-negative")
+    return f"{PROTOCOL_VERSION}:{task_cluster}:seed={int(environment_seed)}"
+
+
+def contract_manifest() -> dict[str, int | str]:
+    return {
+        "version": PROTOCOL_VERSION,
+        "predicted_action_horizon": PREDICTED_ACTION_HORIZON,
+        "executed_action_horizon": EXECUTED_ACTION_HORIZON,
+        "effect_action_horizon": EFFECT_ACTION_HORIZON,
+        "cte_transition_action_horizon": CTE_TRANSITION_ACTION_HORIZON,
+        "effect_cte_transitions": EFFECT_CTE_TRANSITIONS,
+        "max_controls_per_attempt": MAX_CONTROLS_PER_ATTEMPT,
+        "discarded_warmup_requests": REQUIRED_WARMUP_REQUESTS,
+    }

releases/bridge-native-v3/code/cosmos_framework/model/zeva/attempt_protocol_test.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+import pytest
+
+from cosmos_framework.model.zeva.attempt_protocol import (
+    EFFECT_ACTION_HORIZON,
+    EFFECT_CTE_TRANSITIONS,
+    EXECUTED_ACTION_HORIZON,
+    PREDICTED_ACTION_HORIZON,
+    REQUIRED_WARMUP_REQUESTS,
+    CTE_TRANSITION_ACTION_HORIZON,
+    AttemptSessionKey,
+    default_session_id,
+    inference_seed,
+)
+
+
+def test_zeva_temporal_contract_is_exact() -> None:
+    assert PREDICTED_ACTION_HORIZON == 32
+    assert EXECUTED_ACTION_HORIZON == 16
+    assert EFFECT_ACTION_HORIZON == 16
+    assert CTE_TRANSITION_ACTION_HORIZON == 4
+    assert EFFECT_CTE_TRANSITIONS == 4
+    assert CTE_TRANSITION_ACTION_HORIZON * EFFECT_CTE_TRANSITIONS == EFFECT_ACTION_HORIZON
+    assert REQUIRED_WARMUP_REQUESTS == 1
+
+
+def test_attempt_memory_survives_only_exact_task_seed_session() -> None:
+    session = AttemptSessionKey(default_session_id("OpenStandMixerHead", 195), "OpenStandMixerHead", 195)
+    session.assert_compatible(
+        AttemptSessionKey(default_session_id("OpenStandMixerHead", 195), "OpenStandMixerHead", 195)
+    )
+    with pytest.raises(ValueError, match="attempt-session mismatch"):
+        session.assert_compatible(
+            AttemptSessionKey(default_session_id("OpenStandMixerHead", 195), "OpenStandMixerHead", 196)
+        )
+    with pytest.raises(ValueError, match="attempt-session mismatch"):
+        session.assert_compatible(AttemptSessionKey(session.session_id, "TurnOnMicrowave", 195))
+
+
+def test_paired_inference_seed_schedule_is_deterministic() -> None:
+    seeds_a = [inference_seed(20260824, attempt, replan) for attempt in range(3) for replan in range(19)]
+    seeds_b = [inference_seed(20260824, attempt, replan) for attempt in range(3) for replan in range(19)]
+    assert seeds_a == seeds_b
+    assert len(seeds_a) == len(set(seeds_a))

releases/bridge-native-v3/code/cosmos_framework/model/zeva/brief_interaction_trace.py

@@ -0,0 +1,24 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 Z-Trans CORPORATION.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""Brief Interaction Trace (BIT) for the current attempt."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from torch import Tensor
+
+
+@dataclass(frozen=True)
+class BriefInteractionTrace:
+    """Ordered causal effects completed during the current attempt."""
+
+    effects: Tensor
+    valid: Tensor
+
+    def __post_init__(self) -> None:
+        if self.effects.ndim != 3:
+            raise ValueError("BIT effects must be [B,L,D]")
+        if self.valid.shape != self.effects.shape[:2]:
+            raise ValueError("BIT valid mask must be [B,L]")

releases/bridge-native-v3/code/cosmos_framework/model/zeva/causal_prompt.py

@@ -0,0 +1,12 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 Z-Trans CORPORATION.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""Causal Prompt construction and policy injection (paper Eq. 8)."""
+
+from .persistent_interaction_memory import (
+    CausalPromptConfig,
+    CausalPromptEncoder,
+    inject_causal_prompt,
+)
+
+__all__ = ["CausalPromptConfig", "CausalPromptEncoder", "inject_causal_prompt"]

releases/bridge-native-v3/code/cosmos_framework/model/zeva/causal_transition_encoder.py

@@ -0,0 +1,374 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 Z-Trans CORPORATION. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""Causal Transition Encoder (CTE) for robot trajectories.
+
+This is intentionally independent from the Cosmos action denoiser.  It learns a
+per-time-step causal interaction state from the same multi-view observation and
+8-D joint-position command format consumed by Cosmos Action.
+"""
+
+from __future__ import annotations
+
+from copy import deepcopy
+from dataclasses import asdict, dataclass
+
+import torch
+from torch import Tensor, nn
+from torch.nn import functional as F
+from torch.utils.checkpoint import checkpoint
+
+
+@dataclass
+class CausalTransitionEncoderConfig:
+    action_dim: int = 8
+    hidden_dim: int = 256
+    retrieval_dim: int = 128
+    phase_dim: int = 128
+    effect_dim: int = 128
+    transition_steps: int = 4
+    # Phase stays at the VAE's 4x cadence.  Effect uses four such transitions
+    # (16 raw controls) so it is supervised by an observable visual outcome.
+    effect_window_transitions: int = 4
+    effect_target_grid: tuple[int, int] = (4, 6)
+    num_layers: int = 4
+    num_heads: int = 8
+    image_channels: int = 3
+    ema_decay: float = 0.996
+    # GRU is the stable/default RoboCasa lineage. Mamba remains opt-in so that
+    # installing mamba_ssm cannot silently change the checkpoint architecture.
+    use_mamba: bool = False
+    vision_chunk_size: int = 256
+    vision_gradient_checkpointing: bool = True
+    vision_checkpoint_threshold: int = 1024
+
+    def to_dict(self) -> dict[str, object]:
+        return asdict(self)
+
+
+class _VisionStem(nn.Module):
+    """Small trainable visual stem; inputs are DROID's composite multi-view frames."""
+
+    def __init__(
+        self,
+        channels: int,
+        hidden_dim: int,
+        chunk_size: int,
+        gradient_checkpointing: bool,
+        checkpoint_threshold: int,
+    ) -> None:
+        super().__init__()
+        width = max(hidden_dim // 4, 32)
+        self.net = nn.Sequential(
+            nn.Conv2d(channels, width, kernel_size=7, stride=4, padding=3),
+            nn.GELU(),
+            nn.Conv2d(width, width * 2, kernel_size=3, stride=2, padding=1),
+            nn.GELU(),
+            nn.Conv2d(width * 2, hidden_dim, kernel_size=3, stride=2, padding=1),
+            nn.GELU(),
+            nn.AdaptiveAvgPool2d(1),
+        )
+        self.proj = nn.Linear(hidden_dim, hidden_dim)
+        self.chunk_size = chunk_size
+        self.gradient_checkpointing = gradient_checkpointing
+        self.checkpoint_threshold = checkpoint_threshold
+
+    def _encode_flat(self, flat: Tensor) -> Tensor:
+        return self.proj(self.net(flat).flatten(1))
+
+    def forward(self, frames: Tensor) -> Tensor:
+        # [B,T,C,H,W] -> [B,T,D]
+        batch, steps = frames.shape[:2]
+        flat = frames.flatten(0, 1)
+        # A complete DROID segment can exceed 2,000 frames.  Checkpoint each
+        # visual micro-batch so the CTE remains full-trajectory causal while
+        # its convolutional activations do not scale with segment length.
+        use_checkpoint = (
+            self.training
+            and self.gradient_checkpointing
+            and torch.is_grad_enabled()
+            and flat.shape[0] > self.checkpoint_threshold
+        )
+        if not use_checkpoint:
+            return self._encode_flat(flat).unflatten(0, (batch, steps))
+        chunks = []
+        for frame_chunk in flat.split(self.chunk_size, dim=0):
+            if use_checkpoint:
+                chunks.append(checkpoint(self._encode_flat, frame_chunk, use_reentrant=False))
+            else:
+                chunks.append(self._encode_flat(frame_chunk))
+        return torch.cat(chunks, dim=0).unflatten(0, (batch, steps))
+
+
+class _FrozenVAEDeltaTarget(nn.Module):
+    """Fixed spatial projection of cached Cosmos-VAE latents.
+
+    Unlike the EMA vision stem, this target cannot drift to make an auxiliary
+    loss easy.  Retaining a 4x6 pooled grid makes camera/object displacement
+    visible instead of reducing every frame to one spatial mean.
+    """
+
+    def __init__(self, channels: int, hidden_dim: int, grid: tuple[int, int]) -> None:
+        super().__init__()
+        self.grid = grid
+        input_dim = channels * grid[0] * grid[1]
+        generator = torch.Generator().manual_seed(20260815)
+        projection = torch.randn(hidden_dim, input_dim, generator=generator) / input_dim**0.5
+        self.register_buffer("projection", projection, persistent=True)
+
+    @torch.no_grad()
+    def forward(self, frames: Tensor) -> Tensor:
+        batch, steps = frames.shape[:2]
+        flat = frames.flatten(0, 1).float()
+        pooled = F.adaptive_avg_pool2d(flat, self.grid).flatten(1)
+        pooled = F.layer_norm(pooled, (pooled.shape[-1],))
+        return F.linear(pooled, self.projection).unflatten(0, (batch, steps))
+
+
+class _CausalMixer(nn.Module):
+    """Causal temporal mixer with a portable GRU fallback when mamba_ssm is absent."""
+
+    def __init__(self, hidden_dim: int, use_mamba: bool) -> None:
+        super().__init__()
+        self.norm = nn.LayerNorm(hidden_dim)
+        self.kind = "gru"
+        if use_mamba:
+            try:
+                from mamba_ssm import Mamba  # type: ignore[import-not-found]
+
+                self.mixer: nn.Module = Mamba(d_model=hidden_dim, d_state=16, d_conv=4, expand=2)
+                self.kind = "mamba"
+            except ImportError:
+                self.mixer = nn.GRU(hidden_dim, hidden_dim, batch_first=True)
+        else:
+            self.mixer = nn.GRU(hidden_dim, hidden_dim, batch_first=True)
+
+    def forward(self, x: Tensor) -> Tensor:
+        y = self.norm(x)
+        if self.kind == "mamba":
+            y = self.mixer(y)
+        else:
+            y, _ = self.mixer(y)
+        return x + y
+
+
+class _CausalInteractionBlock(nn.Module):
+    """Causal visual/action/interaction streams with same-step cross attention."""
+
+    def __init__(self, cfg: CausalTransitionEncoderConfig) -> None:
+        super().__init__()
+        self.visual = _CausalMixer(cfg.hidden_dim, cfg.use_mamba)
+        self.action = _CausalMixer(cfg.hidden_dim, cfg.use_mamba)
+        self.interaction_state = _CausalMixer(cfg.hidden_dim, cfg.use_mamba)
+        self.cross = nn.MultiheadAttention(cfg.hidden_dim, cfg.num_heads, batch_first=True)
+        self.cross_norm = nn.LayerNorm(cfg.hidden_dim)
+        self.ffn = nn.Sequential(
+            nn.LayerNorm(cfg.hidden_dim),
+            nn.Linear(cfg.hidden_dim, 4 * cfg.hidden_dim),
+            nn.GELU(),
+            nn.Linear(4 * cfg.hidden_dim, cfg.hidden_dim),
+        )
+
+    def forward(self, visual: Tensor, action: Tensor, interaction_state: Tensor, valid: Tensor) -> tuple[Tensor, Tensor, Tensor]:
+        visual, action, interaction_state = (
+            self.visual(visual),
+            self.action(action),
+            self.interaction_state(interaction_state),
+        )
+        # The three tokens for each timestep attend only within that timestep.
+        tokens = torch.stack((visual, action, interaction_state), dim=2)
+        batch, steps, streams, dim = tokens.shape
+        tokens = tokens.reshape(batch * steps, streams, dim)
+        attended, _ = self.cross(tokens, tokens, tokens, need_weights=False)
+        tokens = tokens + attended
+        tokens = tokens + self.ffn(tokens)
+        tokens = tokens.reshape(batch, steps, streams, dim)
+        visual, action, interaction_state = tokens.unbind(dim=2)
+        mask = valid.unsqueeze(-1).to(visual.dtype)
+        return visual * mask, action * mask, interaction_state * mask
+
+
+class CausalTransitionEncoder(nn.Module):
+    """Causal transition encoder with phase and causal-effect projections."""
+
+    def __init__(self, cfg: CausalTransitionEncoderConfig | None = None) -> None:
+        super().__init__()
+        self.cfg = cfg or CausalTransitionEncoderConfig()
+        self.visual_encoder = _VisionStem(
+            self.cfg.image_channels,
+            self.cfg.hidden_dim,
+            self.cfg.vision_chunk_size,
+            self.cfg.vision_gradient_checkpointing,
+            self.cfg.vision_checkpoint_threshold,
+        )
+        self.target_visual_encoder = deepcopy(self.visual_encoder).requires_grad_(False)
+        # The interaction/action stream is intentionally right shifted. At state
+        # v_b it sees only the completed transition u_(b-1), never u_b.  The
+        # latter is supplied exclusively to the auxiliary effect heads.
+        transition_dim = self.cfg.transition_steps * self.cfg.action_dim
+        self.action_encoder = nn.Sequential(nn.LayerNorm(transition_dim), nn.Linear(transition_dim, self.cfg.hidden_dim))
+        self.bos_action = nn.Parameter(torch.zeros(1, 1, self.cfg.hidden_dim))
+        nn.init.normal_(self.bos_action, std=0.02)
+        self.interaction_state_token = nn.Parameter(torch.zeros(1, 1, self.cfg.hidden_dim))
+        nn.init.normal_(self.interaction_state_token, std=0.02)
+        self.blocks = nn.ModuleList([_CausalInteractionBlock(self.cfg) for _ in range(self.cfg.num_layers)])
+        self.final_norm = nn.LayerNorm(self.cfg.hidden_dim)
+        self.retrieval_head = nn.Linear(self.cfg.hidden_dim, self.cfg.retrieval_dim)
+        self.phase_head = nn.Linear(self.cfg.hidden_dim, self.cfg.phase_dim)
+        self.action_head = nn.Sequential(nn.LayerNorm(self.cfg.hidden_dim), nn.Linear(self.cfg.hidden_dim, transition_dim))
+        self.visual_head = nn.Sequential(nn.LayerNorm(self.cfg.hidden_dim), nn.Linear(self.cfg.hidden_dim, self.cfg.hidden_dim))
+        effect_window_dim = self.cfg.effect_window_transitions * transition_dim
+        self.effect_action_encoder = nn.Sequential(
+            nn.LayerNorm(effect_window_dim), nn.Linear(effect_window_dim, self.cfg.hidden_dim)
+        )
+        self.frozen_effect_target = _FrozenVAEDeltaTarget(
+            self.cfg.image_channels, self.cfg.hidden_dim, self.cfg.effect_target_grid
+        )
+        self.effect_pre_head = nn.Sequential(
+            nn.LayerNorm(2 * self.cfg.hidden_dim), nn.Linear(2 * self.cfg.hidden_dim, self.cfg.effect_dim)
+        )
+        # Observed effect is visual only.  Passing executed actions here lets
+        # post simply echo the action-conditioned pre code and caused v2's
+        # rank-one collapse.
+        self.effect_post_head = nn.Sequential(
+            nn.LayerNorm(2 * self.cfg.hidden_dim), nn.Linear(2 * self.cfg.hidden_dim, self.cfg.effect_dim)
+        )
+        self.effect_action_head = nn.Sequential(
+            nn.LayerNorm(self.cfg.effect_dim), nn.Linear(self.cfg.effect_dim, 4 * self.cfg.action_dim)
+        )
+        self.effect_outcome_head = nn.Sequential(
+            nn.LayerNorm(self.cfg.effect_dim), nn.Linear(self.cfg.effect_dim, self.cfg.hidden_dim)
+        )
+
+    @torch.no_grad()
+    def update_ema_target(self) -> None:
+        """Call once after every optimizer step, before building the next target."""
+        for target, online in zip(self.target_visual_encoder.parameters(), self.visual_encoder.parameters(), strict=True):
+            target.lerp_(online, 1.0 - self.cfg.ema_decay)
+
+    @torch.no_grad()
+    def encode_target_vision(self, frames: Tensor) -> Tensor:
+        return self.target_visual_encoder(frames)
+
+    def forward(
+        self,
+        frames: Tensor,
+        transition_actions: Tensor,
+        valid_mask: Tensor | None = None,
+        transition_valid: Tensor | None = None,
+    ) -> dict[str, Tensor]:
+        """Encode padded full trajectories.
+
+        Args:
+            frames: DROID composite observations ``[B,T,3,H,W]`` in [0, 1].
+            transition_actions: completed/raw four-action transitions
+                ``[B,T-1,4,A]``.  At state ``b`` the main action stream only
+                receives transition ``b-1``; transition ``b`` is used solely
+                to supervise the effect heads.
+            valid_mask: ``[B,T]`` true for real timesteps, false for padding.
+        """
+        if frames.ndim != 5 or transition_actions.ndim != 4:
+            raise ValueError("Expected frames [B,T,C,H,W] and transition_actions [B,T-1,4,A].")
+        if (
+            transition_actions.shape[0] != frames.shape[0]
+            or transition_actions.shape[1] != frames.shape[1] - 1
+            or transition_actions.shape[2:] != (self.cfg.transition_steps, self.cfg.action_dim)
+        ):
+            raise ValueError("Frame/transition time axes or action shape do not match CTE config.")
+        if valid_mask is None:
+            valid_mask = torch.ones(frames.shape[:2], dtype=torch.bool, device=frames.device)
+        if transition_valid is None:
+            transition_valid = torch.ones(transition_actions.shape[:-1], dtype=torch.bool, device=frames.device)
+        if transition_valid.shape != transition_actions.shape[:-1]:
+            raise ValueError("transition_valid must have shape [B,T-1,4].")
+        visual = self.visual_encoder(frames)
+        transition_flat = transition_actions.flatten(-2)
+        transition_embed = self.action_encoder(transition_flat)
+        action = self.bos_action.expand(frames.shape[0], frames.shape[1], -1).clone()
+        action[:, 1:] = transition_embed
+        action[:, 1:] *= transition_valid.all(dim=-1, keepdim=True).to(action.dtype)
+        interaction_state = self.interaction_state_token.expand_as(visual)
+        for block in self.blocks:
+            visual, action, interaction_state = block(visual, action, interaction_state, valid_mask)
+        z = self.final_norm(interaction_state) * valid_mask.unsqueeze(-1)
+        with torch.no_grad():
+            target_visual = self.target_visual_encoder(frames)
+        transition_complete = valid_mask[:, :-1] & valid_mask[:, 1:] & transition_valid.all(dim=-1)
+        effect_windows = (frames.shape[1] - 1) // self.cfg.effect_window_transitions
+        effect_steps = self.cfg.effect_window_transitions
+        if effect_windows:
+            effect_start = torch.arange(effect_windows, device=frames.device) * effect_steps
+            effect_end = effect_start + effect_steps
+            effect_actions = transition_actions[:, : effect_windows * effect_steps].reshape(
+                frames.shape[0], effect_windows, effect_steps, self.cfg.transition_steps, self.cfg.action_dim
+            )
+            effect_action_valid = transition_valid[:, : effect_windows * effect_steps].reshape(
+                frames.shape[0], effect_windows, effect_steps, self.cfg.transition_steps
+            )
+            effect_complete = (
+                valid_mask[:, effect_start]
+                & valid_mask[:, effect_end]
+                & effect_action_valid.all(dim=(-1, -2))
+            )
+            effect_embed = self.effect_action_encoder(effect_actions.flatten(-3))
+            frozen_visual = self.frozen_effect_target(frames)
+            effect_delta = frozen_visual[:, effect_end] - frozen_visual[:, effect_start]
+            effect_pre_raw = self.effect_pre_head(torch.cat((z[:, effect_start], effect_embed), dim=-1))
+            effect_post_raw = self.effect_post_head(torch.cat((frozen_visual[:, effect_start], effect_delta), dim=-1))
+        else:
+            effect_actions = transition_actions.new_zeros((frames.shape[0], 0, effect_steps, self.cfg.transition_steps, self.cfg.action_dim))
+            effect_complete = torch.zeros((frames.shape[0], 0), dtype=torch.bool, device=frames.device)
+            effect_delta = z[:, :0]
+            effect_pre_raw = z[:, :0]
+            effect_post_raw = z[:, :0]
+        # At reset there is a single observed frame and no completed
+        # transition/effect window.  LayerNorm cannot be applied to the empty
+        # ``[..., 0, hidden_dim]`` fallback above because the effect heads
+        # expect ``effect_dim``.  Preserve the causal no-effect state with
+        # correctly shaped empty outputs; callers then provide only BOS/pad
+        # effect tokens to the policy.
+        if effect_pre_raw.shape[1] == 0:
+            effect_pre = effect_pre_raw.new_zeros(
+                (frames.shape[0], 0, self.cfg.effect_dim)
+            )
+            effect_post = effect_pre_raw.new_zeros(
+                (frames.shape[0], 0, self.cfg.effect_dim)
+            )
+            effect_outcome_pre = effect_pre_raw.new_zeros(
+                (frames.shape[0], 0, self.cfg.hidden_dim)
+            )
+            effect_outcome_post = effect_pre_raw.new_zeros(
+                (frames.shape[0], 0, self.cfg.hidden_dim)
+            )
+            effect_action = effect_pre_raw.new_zeros(
+                (frames.shape[0], 0, 4 * self.cfg.action_dim)
+            )
+        else:
+            effect_pre = F.normalize(effect_pre_raw, dim=-1)
+            effect_post = F.normalize(effect_post_raw, dim=-1)
+            effect_outcome_pre = self.effect_outcome_head(effect_pre_raw)
+            effect_outcome_post = self.effect_outcome_head(effect_post_raw)
+            effect_action = self.effect_action_head(effect_pre_raw)
+        return {
+            "causal_interaction_state": z,
+            "retrieval": F.normalize(self.retrieval_head(z), dim=-1),
+            "phase": F.normalize(self.phase_head(z), dim=-1),
+            "next_action": self.action_head(z[:, :-1]).view(
+                frames.shape[0], frames.shape[1] - 1, self.cfg.transition_steps, self.cfg.action_dim
+            ),
+            "next_vision": self.visual_head(z),
+            # Frozen target-visual representation used as the PIM retrieval key.
+            "visual_key": F.normalize(target_visual[..., : self.cfg.phase_dim], dim=-1),
+            "effect_pre": effect_pre,
+            "effect_post": effect_post,
+            "effect_pre_raw": effect_pre_raw,
+            "effect_post_raw": effect_post_raw,
+            "effect_outcome_pre": effect_outcome_pre,
+            "effect_outcome_post": effect_outcome_post,
+            "effect_action": effect_action,
+            "effect_actions": effect_actions,
+            "effect_delta_target": effect_delta,
+            "effect_complete": effect_complete,
+            "target_visual": target_visual,
+            "transition_complete": transition_complete,
+        }

releases/bridge-native-v3/code/cosmos_framework/model/zeva/causal_transition_encoder_test.py

@@ -0,0 +1,28 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 Z-Trans CORPORATION. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+import torch
+
+from cosmos_framework.model.zeva import (
+    CausalTransitionEncoder,
+    CausalTransitionEncoderConfig,
+    causal_transition_encoder_loss,
+)
+
+
+def test_cte_forward_loss_and_ema() -> None:
+    model = CausalTransitionEncoder(CausalTransitionEncoderConfig(hidden_dim=32, retrieval_dim=16, phase_dim=16, num_layers=1, num_heads=4, use_mamba=False))
+    frames = torch.rand(3, 5, 3, 48, 48)
+    actions = torch.rand(3, 4, 4, 8)
+    valid = torch.tensor([[1, 1, 1, 1, 1], [1, 1, 1, 1, 0], [1, 1, 1, 0, 0]], dtype=torch.bool)
+    output = model(frames, actions, valid)
+    assert output["causal_interaction_state"].shape == (3, 5, 32)
+    assert output["retrieval"].shape == (3, 5, 16)
+    losses = causal_transition_encoder_loss(output, actions, valid, torch.tensor([0, 0, 1]))
+    assert torch.isfinite(losses["total"])
+    losses["total"].backward()
+    before = next(model.target_visual_encoder.parameters()).clone()
+    with torch.no_grad():
+        next(model.visual_encoder.parameters()).add_(0.01)
+    model.update_ema_target()
+    assert not torch.equal(before, next(model.target_visual_encoder.parameters()))

releases/bridge-native-v3/code/cosmos_framework/model/zeva/checkpoint_io.py

@@ -0,0 +1,22 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 Z-Trans CORPORATION.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""Checkpoint loading utilities for Zeva modules."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any
+
+
+def normalize_cte_state_dict(state_dict: Mapping[str, Any]) -> dict[str, Any]:
+    """Normalize serialized CTE parameter names without changing tensors."""
+
+    normalized: dict[str, Any] = {}
+    for key, value in state_dict.items():
+        new_key = key.replace("behavior_token", "interaction_state_token")
+        new_key = new_key.replace(".behavior.", ".interaction_state.")
+        if new_key in normalized:
+            raise ValueError(f"Duplicate CTE checkpoint key after normalization: {new_key}")
+        normalized[new_key] = value
+    return normalized

releases/bridge-native-v3/code/cosmos_framework/model/zeva/checkpoint_io_test.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+import torch
+
+from cosmos_framework.model.zeva import (
+    CausalTransitionEncoder,
+    CausalTransitionEncoderConfig,
+    normalize_cte_state_dict,
+)
+
+
+def test_cte_key_normalization_preserves_every_tensor() -> None:
+    config = CausalTransitionEncoderConfig(
+        hidden_dim=16,
+        retrieval_dim=8,
+        phase_dim=8,
+        effect_dim=8,
+        num_layers=1,
+        num_heads=4,
+        use_mamba=False,
+    )
+    source = CausalTransitionEncoder(config)
+    serialized = {}
+    for key, value in source.state_dict().items():
+        serialized_key = key.replace("interaction_state_token", "behavior_token")
+        serialized_key = serialized_key.replace(".interaction_state.", ".behavior.")
+        serialized[serialized_key] = value.clone()
+
+    restored = CausalTransitionEncoder(config)
+    restored.load_state_dict(normalize_cte_state_dict(serialized), strict=True)
+    for key, value in source.state_dict().items():
+        assert torch.equal(value, restored.state_dict()[key]), key

releases/bridge-native-v3/code/cosmos_framework/model/zeva/cte_losses.py

@@ -0,0 +1,177 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 Z-Trans CORPORATION. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""Learning objectives for the Causal Transition Encoder (CTE)."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+from torch import Tensor
+from torch.nn import functional as F
+
+
+@dataclass
+class CTELossConfig:
+    action_weight: float = 1.0
+    vision_weight: float = 1.0
+    task_weight: float = 0.2
+    phase_weight: float = 0.1
+    # Effect-v3: observed outcomes are contrastive and explicitly protected
+    # against rank collapse. Action reconstruction remains a weak auxiliary.
+    effect_weight: float = 0.25
+    effect_contrastive_weight: float = 1.0
+    effect_action_weight: float = 0.05
+    effect_align_weight: float = 0.1
+    effect_variance_weight: float = 1.0
+    effect_covariance_weight: float = 0.04
+    effect_temperature: float = 0.07
+    temperature: float = 0.1
+
+
+def _masked_mean(x: Tensor, mask: Tensor) -> Tensor:
+    weight = mask.unsqueeze(-1).to(x.dtype)
+    return (x * weight).sum(dim=1) / weight.sum(dim=1).clamp_min(1.0)
+
+
+def _task_identity_clustering_loss(
+    retrieval: Tensor, valid: Tensor, semantic_ids: Tensor, temperature: float
+) -> Tensor:
+    """Multi-positive supervised contrastive loss; samples without a positive are ignored."""
+    z = F.normalize(_masked_mean(retrieval, valid), dim=-1)
+    logits = z @ z.T / temperature
+    logits.fill_diagonal_(-torch.inf)
+    positives = semantic_ids[:, None].eq(semantic_ids[None, :])
+    positives.fill_diagonal_(False)
+    usable = positives.any(dim=1)
+    if not bool(usable.any()):
+        return z.new_zeros(())
+    log_prob = logits - torch.logsumexp(logits, dim=1, keepdim=True)
+    return -(log_prob.masked_fill(~positives, 0.0).sum(dim=1) / positives.sum(dim=1).clamp_min(1))[usable].mean()
+
+
+def _phase_progression_loss(phase: Tensor, valid: Tensor, margin: float = 0.01) -> Tensor:
+    """Encourage consecutive phase tokens to advance along one trajectory direction."""
+    losses: list[Tensor] = []
+    for sequence, mask in zip(phase, valid, strict=True):
+        z = sequence[mask]
+        if z.shape[0] < 3:
+            continue
+        direction = F.normalize(z[-1] - z[0], dim=-1)
+        increments = (z[1:] - z[:-1]) @ direction
+        losses.append(F.relu(margin - increments).mean())
+    return torch.stack(losses).mean() if losses else phase.new_zeros(())
+
+
+def summarize_effect_window(actions: Tensor) -> Tensor:
+    """Summarize an executed low-rate effect window ``[B,W,K,4,A]``.
+
+    Keeping first, net displacement, mean and standard deviation makes this a
+    light action regularizer, not a loss that can be solved by copying the
+    current four-control input into the effect code.
+    """
+    if actions.ndim != 5:
+        raise ValueError("Effect actions must have shape [B,W,K,4,A].")
+    sequence = actions.flatten(-3, -2)
+    return torch.cat(
+        (
+            sequence[..., 0, :],
+            sequence[..., -1, :] - sequence[..., 0, :],
+            sequence.mean(dim=-2),
+            sequence.std(dim=-2, correction=0),
+        ),
+        dim=-1,
+    )
+
+
+def _effect_nce(query: Tensor, target: Tensor, mask: Tensor, temperature: float) -> Tensor:
+    query, target = query[mask], target[mask]
+    if query.shape[0] < 2:
+        return query.new_zeros(())
+    logits = F.normalize(query, dim=-1) @ F.normalize(target, dim=-1).T / temperature
+    return F.cross_entropy(logits, torch.arange(logits.shape[0], device=logits.device))
+
+
+def _vicreg(code: Tensor, mask: Tensor) -> tuple[Tensor, Tensor]:
+    code = code[mask]
+    if code.shape[0] < 2:
+        zero = code.new_zeros(())
+        return zero, zero
+    centered = code - code.mean(dim=0, keepdim=True)
+    std = (centered.var(dim=0, correction=0) + 1e-4).sqrt()
+    variance = F.relu(1.0 - std).mean()
+    covariance = centered.T @ centered / max(code.shape[0] - 1, 1)
+    covariance.fill_diagonal_(0)
+    covariance_loss = covariance.square().sum() / code.shape[-1]
+    return variance, covariance_loss
+
+
+def causal_transition_encoder_loss(
+    outputs: dict[str, Tensor],
+    transition_actions: Tensor,
+    valid_mask: Tensor,
+    semantic_ids: Tensor,
+    cfg: CTELossConfig | None = None,
+) -> dict[str, Tensor]:
+    """CTE objective: causal effect, task identity, and phase progression."""
+    cfg = cfg or CTELossConfig()
+    transition_complete = outputs["transition_complete"]
+    if transition_complete.shape != transition_actions.shape[:2]:
+        raise ValueError("Effect outputs and transition-action cache are not temporally aligned.")
+    action_error = F.smooth_l1_loss(
+        outputs["next_action"], transition_actions, reduction="none"
+    ).mean(dim=(-1, -2))
+    action_loss = (action_error * transition_complete).sum() / transition_complete.sum().clamp_min(1)
+    target_vision = outputs["target_visual"]
+    vision_error = F.mse_loss(
+        outputs["next_vision"][:, :-1], target_vision[:, 1:], reduction="none"
+    ).mean(dim=-1)
+    vision_loss = (vision_error * transition_complete).sum() / transition_complete.sum().clamp_min(1)
+    task_loss = _task_identity_clustering_loss(outputs["retrieval"], valid_mask, semantic_ids, cfg.temperature)
+    phase_loss = _phase_progression_loss(outputs["phase"], valid_mask)
+    effect_mask = outputs["effect_complete"]
+    effect_target = outputs["effect_delta_target"]
+    effect_nce_pre = _effect_nce(outputs["effect_outcome_pre"], effect_target, effect_mask, cfg.effect_temperature)
+    effect_nce_post = _effect_nce(outputs["effect_outcome_post"], effect_target, effect_mask, cfg.effect_temperature)
+    effect_contrastive_loss = 0.5 * (effect_nce_pre + effect_nce_post)
+    # Diagnostic only: absolute deltas have a strong zero baseline, so this
+    # MSE is intentionally not optimized directly.
+    effect_visual_error = F.mse_loss(outputs["effect_outcome_pre"], effect_target, reduction="none").mean(dim=-1)
+    effect_visual_loss = (effect_visual_error * effect_mask).sum() / effect_mask.sum().clamp_min(1)
+    action_target = summarize_effect_window(outputs["effect_actions"])
+    effect_action_error = F.smooth_l1_loss(
+        outputs["effect_action"], action_target, reduction="none"
+    ).mean(dim=-1)
+    effect_action_loss = (effect_action_error * effect_mask).sum() / effect_mask.sum().clamp_min(1)
+    effect_align_error = 1.0 - (outputs["effect_pre"] * outputs["effect_post"]).sum(dim=-1)
+    effect_align_loss = (effect_align_error * effect_mask).sum() / effect_mask.sum().clamp_min(1)
+    effect_variance_loss, effect_covariance_loss = _vicreg(outputs["effect_post_raw"], effect_mask)
+    effect_loss = (
+        cfg.effect_contrastive_weight * effect_contrastive_loss
+        + cfg.effect_action_weight * effect_action_loss
+        + cfg.effect_align_weight * effect_align_loss
+        + cfg.effect_variance_weight * effect_variance_loss
+        + cfg.effect_covariance_weight * effect_covariance_loss
+    )
+    total = (
+        cfg.action_weight * action_loss
+        + cfg.vision_weight * vision_loss
+        + cfg.task_weight * task_loss
+        + cfg.phase_weight * phase_loss
+        + cfg.effect_weight * effect_loss
+    )
+    return {
+        "total": total,
+        "action": action_loss,
+        "vision": vision_loss,
+        "loss_task": task_loss,
+        "loss_phase": phase_loss,
+        "loss_effect": effect_loss,
+        "effect_contrastive": effect_contrastive_loss,
+        "effect_visual": effect_visual_loss,
+        "effect_action": effect_action_loss,
+        "effect_align": effect_align_loss,
+        "effect_variance": effect_variance_loss,
+        "effect_covariance": effect_covariance_loss,
+    }

releases/bridge-native-v3/code/cosmos_framework/model/zeva/experimental/__init__.py

@@ -0,0 +1 @@
+"""Experimental Zeva components that are not part of the paper architecture."""

releases/bridge-native-v3/code/cosmos_framework/model/zeva/experimental/robocasa_transition_memory.py

@@ -0,0 +1,49 @@
+"""RoboCasa Atomic-5 experimental transition-memory contract.
+
+This module keeps the RoboCasa contract separate from the ARX defaults in
+the generic transition-memory defaults. In particular, RoboCasa uses the
+frozen Atomic-5 CTE,
+7-dimensional arm actions and the left-agent/wrist 20-Hz temporal contract.
+"""
+
+from __future__ import annotations
+
+from .transition_memory import TransitionMemorySchema
+
+
+ROBOCASA_ATOMIC5_TASKS = (
+    "OpenStandMixerHead",
+    "TurnOnElectricKettle",
+    "CloseToasterOvenDoor",
+    "TurnOnMicrowave",
+    "CoffeeSetupMug",
+)
+
+
+def make_robocasa_atomic5_schema(
+    *,
+    cte_hash: str = "",
+    vae_temporal_hash: str = "",
+    capacity: int = 64,
+    top_k: int = 4,
+) -> TransitionMemorySchema:
+    """Return the Atomic-5 transition-memory schema.
+
+    ``cte_hash`` and ``vae_temporal_hash`` are part of the schema hash.  A
+    runtime/cache produced with a different CTE or temporal contract is
+    therefore rejected instead of silently mixing representations.
+    """
+
+    return TransitionMemorySchema(
+        version="robocasa_atomic5_transition_memory_v1",
+        task_contract="robocasa365_atomic5_5task_left_wrist_arm7_20hz_cosmos32_exec16",
+        phase_dim=128,
+        visual_key_dim=128,
+        effect_dim=128,
+        action_dim=7,
+        action_horizon=16,
+        capacity=int(capacity),
+        top_k=int(top_k),
+        cte_hash=str(cte_hash),
+        vae_temporal_hash=str(vae_temporal_hash),
+    )

releases/bridge-native-v3/code/cosmos_framework/model/zeva/experimental/transition_memory.py

@@ -0,0 +1,246 @@
+"""Experimental transition memory retained for non-paper ablations.
+
+The memory stores only completed 16-control transitions.  It is deliberately
+independent of the Cosmos model so that real-robot code can run it in shadow
+mode before enabling the extra conditioning prefix.
+"""
+
+from __future__ import annotations
+
+from collections import deque
+from dataclasses import dataclass, field
+import hashlib
+import json
+from typing import Iterable
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+
+@dataclass(frozen=True)
+class TransitionMemorySchema:
+    """Dimensions and temporal contract used by one memory session."""
+
+    version: str = "arx_transition_memory_v1"
+    task_contract: str = "arx_task7_model_a_5task"
+    phase_dim: int = 128
+    visual_key_dim: int = 128
+    effect_dim: int = 128
+    action_dim: int = 14
+    action_horizon: int = 16
+    capacity: int = 64
+    top_k: int = 4
+    cte_hash: str = ""
+    vae_temporal_hash: str = ""
+
+    @property
+    def hash(self) -> str:
+        payload = json.dumps(self.__dict__, sort_keys=True, separators=(",", ":"))
+        return hashlib.sha256(payload.encode("utf-8")).hexdigest()
+
+
+@dataclass
+class TransitionRecord:
+    """One causally completed replan transition."""
+
+    task_cluster: str
+    phase: torch.Tensor
+    visual_key: torch.Tensor
+    effect_post: torch.Tensor
+    executed_action: torch.Tensor
+    next_visual_key: torch.Tensor
+    next_phase: torch.Tensor
+    latent_index: int
+    schema_hash: str
+    valid: bool = True
+    metadata: dict[str, object] = field(default_factory=dict)
+
+    def cpu(self) -> "TransitionRecord":
+        """Detach tensors so the session never retains a training graph."""
+        return TransitionRecord(
+            task_cluster=self.task_cluster,
+            phase=self.phase.detach().float().cpu().contiguous(),
+            visual_key=self.visual_key.detach().float().cpu().contiguous(),
+            effect_post=self.effect_post.detach().float().cpu().contiguous(),
+            executed_action=self.executed_action.detach().float().cpu().contiguous(),
+            next_visual_key=self.next_visual_key.detach().float().cpu().contiguous(),
+            next_phase=self.next_phase.detach().float().cpu().contiguous(),
+            latent_index=int(self.latent_index),
+            schema_hash=self.schema_hash,
+            valid=bool(self.valid),
+            metadata=dict(self.metadata),
+        )
+
+
+class TransitionMemory:
+    """Bounded FIFO memory scoped to one task session."""
+
+    def __init__(self, schema: TransitionMemorySchema | None = None) -> None:
+        self.schema = schema or TransitionMemorySchema()
+        self._task_cluster: str | None = None
+        self._entries: deque[TransitionRecord] = deque(maxlen=self.schema.capacity)
+
+    @property
+    def task_cluster(self) -> str | None:
+        return self._task_cluster
+
+    def __len__(self) -> int:
+        return len(self._entries)
+
+    @property
+    def entries(self) -> tuple[TransitionRecord, ...]:
+        """Return a stable, read-only snapshot for audit/finalization."""
+        return tuple(self._entries)
+
+    def reset(self, task_cluster: str | None = None) -> None:
+        self._entries.clear()
+        self._task_cluster = task_cluster
+
+    def _validate(self, transition: TransitionRecord) -> None:
+        if transition.schema_hash != self.schema.hash:
+            raise ValueError(
+                "transition memory schema mismatch: "
+                f"got {transition.schema_hash}, expected {self.schema.hash}"
+            )
+        if transition.phase.shape != (self.schema.phase_dim,):
+            raise ValueError(f"phase must be [{self.schema.phase_dim}], got {tuple(transition.phase.shape)}")
+        if transition.visual_key.shape != (self.schema.visual_key_dim,):
+            raise ValueError("visual_key has an invalid shape")
+        if transition.effect_post.shape != (self.schema.effect_dim,):
+            raise ValueError("effect_post has an invalid shape")
+        if transition.executed_action.shape != (self.schema.action_horizon, self.schema.action_dim):
+            raise ValueError("executed_action must contain exactly one completed 16-step window")
+        if transition.next_visual_key.shape != (self.schema.visual_key_dim,):
+            raise ValueError("next_visual_key has an invalid shape")
+        if transition.next_phase.shape != (self.schema.phase_dim,):
+            raise ValueError("next_phase has an invalid shape")
+        if not transition.valid:
+            raise ValueError("invalid or incomplete transitions must not enter transition memory")
+
+    def append(self, transition: TransitionRecord) -> None:
+        """Append one completed transition, resetting on a task change."""
+        self._validate(transition)
+        if self._task_cluster is None:
+            self._task_cluster = transition.task_cluster
+        if transition.task_cluster != self._task_cluster:
+            raise ValueError(
+                f"task mismatch: session={self._task_cluster!r}, transition={transition.task_cluster!r}; reset first"
+            )
+        self._entries.append(transition.cpu())
+
+    def annotate_attempt_outcome(
+        self,
+        attempt_id: int,
+        *,
+        outcome: str,
+        termination_reason: str,
+        total_steps: int,
+        final_progress: float,
+    ) -> int:
+        """Attach a terminal label to every transition from one attempt."""
+        if outcome not in {"success", "failure"}:
+            raise ValueError(f"unsupported attempt outcome: {outcome!r}")
+        if termination_reason not in {"success", "env_done", "max_steps", "client_error", "safety"}:
+            raise ValueError(f"unsupported termination reason: {termination_reason!r}")
+        if attempt_id < 0 or total_steps < 0:
+            raise ValueError("attempt_id and total_steps must be non-negative")
+        if not 0.0 <= float(final_progress) <= 1.0:
+            raise ValueError("final_progress must be in [0,1]")
+        updated = 0
+        for entry in self._entries:
+            if int(entry.metadata.get("attempt_id", -1)) != attempt_id:
+                continue
+            entry.metadata.update(
+                {
+                    "terminal_outcome": outcome,
+                    "termination_reason": termination_reason,
+                    "total_steps": int(total_steps),
+                    "final_progress": float(final_progress),
+                    "progress_source": "success_only",
+                }
+            )
+            updated += 1
+        return updated
+
+    def query(
+        self,
+        phase: torch.Tensor,
+        visual_key: torch.Tensor,
+        *,
+        top_k: int | None = None,
+    ) -> tuple[list[TransitionRecord], torch.Tensor]:
+        """Return cosine-ranked history entries and scores."""
+        if phase.shape != (self.schema.phase_dim,) or visual_key.shape != (self.schema.visual_key_dim,):
+            raise ValueError("query phase/visual_key shape mismatch")
+        if not self._entries:
+            return [], torch.empty(0, dtype=torch.float32)
+        k = min(int(top_k or self.schema.top_k), len(self._entries))
+        q = F.normalize(torch.cat((phase.float(), visual_key.float())), dim=0)
+        keys = torch.stack([
+            F.normalize(torch.cat((entry.phase, entry.visual_key)), dim=0)
+            for entry in self._entries
+        ])
+        scores = keys @ q
+        values, indices = torch.topk(scores, k=k, largest=True, sorted=True)
+        entries = [list(self._entries)[int(index)] for index in indices]
+        return entries, values
+
+
+class TransitionMemoryEncoder(nn.Module):
+    """Encode top-k transitions into one 256-D causal context token."""
+
+    def __init__(self, schema: TransitionMemorySchema | None = None, hidden_dim: int = 256) -> None:
+        super().__init__()
+        self.schema = schema or TransitionMemorySchema()
+        self.hidden_dim = hidden_dim
+        value_dim = (
+            self.schema.phase_dim
+            + self.schema.visual_key_dim
+            + self.schema.effect_dim
+            + self.schema.action_horizon * self.schema.action_dim
+            + self.schema.visual_key_dim
+            + self.schema.phase_dim
+        )
+        query_dim = self.schema.phase_dim + self.schema.visual_key_dim
+        self.query_proj = nn.Sequential(nn.Linear(query_dim, hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, hidden_dim))
+        self.key_proj = nn.Sequential(nn.Linear(query_dim, hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, hidden_dim))
+        self.value_proj = nn.Sequential(nn.Linear(value_dim, hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, hidden_dim))
+        self.output = nn.Sequential(nn.LayerNorm(hidden_dim), nn.Linear(hidden_dim, hidden_dim), nn.SiLU())
+        self.bos_online = nn.Parameter(torch.zeros(hidden_dim))
+        nn.init.normal_(self.bos_online, std=0.02)
+
+    def forward(
+        self,
+        phase: torch.Tensor,
+        visual_key: torch.Tensor,
+        entries: Iterable[TransitionRecord] | None = None,
+    ) -> torch.Tensor:
+        """Return ``[B, hidden_dim]`` or ``[hidden_dim]`` for one query."""
+        if phase.ndim == 1:
+            phase = phase.unsqueeze(0)
+        if visual_key.ndim == 1:
+            visual_key = visual_key.unsqueeze(0)
+        if phase.shape[0] != 1 or visual_key.shape[0] != 1:
+            raise ValueError("TransitionMemoryEncoder currently accepts one query at a time")
+        device = phase.device
+        query = torch.cat((phase.float(), visual_key.float()), dim=-1)
+        q = self.query_proj(query)
+        items = list(entries or [])[: self.schema.top_k]
+        if not items:
+            return self.output(q + self.bos_online.to(q)).to(phase.dtype)
+        keys = torch.stack([torch.cat((item.phase, item.visual_key)) for item in items]).to(device=device)
+        values = torch.stack([
+            torch.cat((
+                item.phase,
+                item.visual_key,
+                item.effect_post,
+                item.executed_action.reshape(-1),
+                item.next_visual_key,
+                item.next_phase,
+            ))
+            for item in items
+        ]).to(device=device)
+        attention = torch.softmax(F.normalize(self.key_proj(keys), dim=-1) @ F.normalize(q, dim=-1).squeeze(0), dim=0)
+        context = (attention.unsqueeze(-1) * self.value_proj(values)).sum(dim=0, keepdim=True)
+        return self.output(q + context).to(phase.dtype)

releases/bridge-native-v3/code/cosmos_framework/model/zeva/experimental/transition_memory_policy_adapter.py

@@ -0,0 +1,116 @@
+"""Learned action residual projector for the Cosmos transition-memory prefix.
+
+This module is intentionally tiny and dependency-light so the same class can
+be copied into the remote Cosmos framework checkout.  It is not a selector:
+it receives only the server-produced action chunk, the current phase/visual
+features and the learned transition-memory context, and emits a bounded residual.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import Tensor, nn
+
+
+ACTION_PROJECTOR_FORMAT = "robocasa_atomic5_transition_memory_action_projector_v1"
+
+
+class TransitionMemoryPolicyAdapter(nn.Module):
+    """Predict a bounded 32x7 residual from the learned memory context."""
+
+    def __init__(
+        self,
+        *,
+        context_dim: int = 256,
+        phase_dim: int = 128,
+        visual_dim: int = 128,
+        action_horizon: int = 32,
+        action_dim: int = 7,
+        hidden_dim: int = 256,
+        residual_scale: float = 0.5,
+        effective_horizon: int = 16,
+        residual_scale_max: float | None = None,
+    ) -> None:
+        super().__init__()
+        if action_horizon != 32 or action_dim != 7:
+            raise ValueError("Atomic-5 action projector requires horizon=32 and action_dim=7")
+        if effective_horizon < 1 or effective_horizon > action_horizon:
+            raise ValueError("effective_horizon must be in [1, action_horizon]")
+        if residual_scale_max is None:
+            residual_scale_max = float(residual_scale)
+        if not torch.isfinite(torch.tensor(float(residual_scale_max))):
+            raise ValueError("residual_scale_max must be finite")
+        if float(residual_scale_max) <= 0.0:
+            raise ValueError("residual_scale_max must be positive")
+        if not torch.isfinite(torch.tensor(float(residual_scale))):
+            raise ValueError("residual_scale must be finite")
+        if float(residual_scale) < 0.0 or float(residual_scale) > float(residual_scale_max):
+            raise ValueError("residual_scale must lie in [0, residual_scale_max]")
+        self.context_dim = int(context_dim)
+        self.phase_dim = int(phase_dim)
+        self.visual_dim = int(visual_dim)
+        self.action_horizon = int(action_horizon)
+        self.action_dim = int(action_dim)
+        self.hidden_dim = int(hidden_dim)
+        self.residual_scale = float(residual_scale)
+        self.residual_scale_max = float(residual_scale_max)
+        self.effective_horizon = int(effective_horizon)
+        input_dim = context_dim + phase_dim + visual_dim + action_horizon * action_dim
+        self.body = nn.Sequential(
+            nn.Linear(input_dim, hidden_dim),
+            nn.SiLU(),
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.SiLU(),
+        )
+        self.residual_head = nn.Linear(hidden_dim, action_horizon * action_dim)
+
+    def forward(
+        self,
+        context: Tensor,
+        phase: Tensor,
+        visual_key: Tensor,
+        baseline_actions: Tensor,
+    ) -> tuple[Tensor, Tensor]:
+        if context.ndim == 1:
+            context = context.unsqueeze(0)
+        if phase.ndim == 1:
+            phase = phase.unsqueeze(0)
+        if visual_key.ndim == 1:
+            visual_key = visual_key.unsqueeze(0)
+        if baseline_actions.ndim != 3:
+            raise ValueError("baseline_actions must be [B,32,7]")
+        expected = (self.action_horizon, self.action_dim)
+        if tuple(baseline_actions.shape[1:]) != expected:
+            raise ValueError(f"baseline_actions must be [B,{expected[0]},{expected[1]}]")
+        if context.shape[0] != baseline_actions.shape[0] or phase.shape[0] != baseline_actions.shape[0] or visual_key.shape[0] != baseline_actions.shape[0]:
+            raise ValueError("projector batch dimensions do not match")
+        # The evaluator executes only the first ``effective_horizon`` controls
+        # of a 32-action prediction.  Do not let the unexecuted suffix affect
+        # the learned correction: it is often stochastic padding/noise.
+        baseline_for_features = baseline_actions.float()
+        if self.effective_horizon < self.action_horizon:
+            mask = torch.zeros_like(baseline_for_features)
+            mask[:, : self.effective_horizon] = 1.0
+            baseline_for_features = baseline_for_features * mask
+        features = torch.cat(
+            (context.float(), phase.float(), visual_key.float(), baseline_for_features.reshape(baseline_actions.shape[0], -1)),
+            dim=-1,
+        )
+        residual = torch.tanh(self.residual_head(self.body(features))).reshape(-1, *expected)
+        if self.effective_horizon < self.action_horizon:
+            mask = torch.zeros_like(residual)
+            mask[:, : self.effective_horizon] = 1.0
+            residual = residual * mask
+        residual = residual * self.residual_scale
+        return baseline_actions + residual, residual
+
+    def set_residual_scale(self, value: float) -> None:
+        """Set a calibrated bounded residual scale without changing weights."""
+        value = float(value)
+        if not torch.isfinite(torch.tensor(value)):
+            raise ValueError("residual scale must be finite")
+        if value < 0.0 or value > self.residual_scale_max:
+            raise ValueError(
+                f"residual scale {value} is outside [0, {self.residual_scale_max}]"
+            )
+        self.residual_scale = value

releases/bridge-native-v3/code/cosmos_framework/model/zeva/experimental/transition_memory_runtime.py

@@ -0,0 +1,100 @@
+"""Runtime glue for the experimental task-session transition memory."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+
+from .transition_memory import TransitionMemory, TransitionMemoryEncoder, TransitionMemorySchema, TransitionRecord
+
+
+@dataclass
+class TransitionMemoryReadout:
+    context: torch.Tensor
+    scores: torch.Tensor
+    num_entries: int
+
+
+class TransitionMemoryController:
+    """Owns task reset, causal writes, and context reads for one policy server."""
+
+    def __init__(
+        self,
+        schema: TransitionMemorySchema | None = None,
+        *,
+        encoder: TransitionMemoryEncoder | None = None,
+        enabled: bool = False,
+    ) -> None:
+        self.schema = schema or TransitionMemorySchema()
+        self.memory = TransitionMemory(self.schema)
+        self.encoder = encoder or TransitionMemoryEncoder(self.schema)
+        self.enabled = bool(enabled)
+        self._replan_open = False
+
+    def reset(self, task_cluster: str) -> None:
+        self.memory.reset(task_cluster)
+        self._replan_open = False
+
+    def begin_replan(self) -> None:
+        if self.memory.task_cluster is None:
+            raise RuntimeError("reset(task_cluster) is required before begin_replan")
+        if self._replan_open:
+            raise RuntimeError("previous replan is still open")
+        self._replan_open = True
+
+    def discard_open_replan(self) -> bool:
+        """Discard a terminal partial window without writing it to memory."""
+        was_open = self._replan_open
+        self._replan_open = False
+        return was_open
+
+    @torch.inference_mode()
+    def read(self, phase: torch.Tensor, visual_key: torch.Tensor) -> TransitionMemoryReadout:
+        """Read only transitions written before this replan."""
+        if not self._replan_open:
+            raise RuntimeError("begin_replan() must precede an transition-memory read")
+        entries, scores = self.memory.query(phase.detach().cpu(), visual_key.detach().cpu())
+        if self.enabled:
+            context = self.encoder(phase, visual_key, entries)
+        else:
+            context = phase.new_zeros((1, self.encoder.hidden_dim))
+        return TransitionMemoryReadout(context=context, scores=scores, num_entries=len(entries))
+
+    def complete_replan(
+        self,
+        *,
+        phase: torch.Tensor,
+        visual_key: torch.Tensor,
+        effect_post: torch.Tensor,
+        executed_action: torch.Tensor,
+        next_visual_key: torch.Tensor,
+        next_phase: torch.Tensor,
+        latent_index: int,
+        completed: bool,
+        metadata: dict[str, object] | None = None,
+    ) -> bool:
+        """Commit exactly one completed 16-step transition, or discard it."""
+        if not self._replan_open:
+            raise RuntimeError("begin_replan() must precede complete_replan()")
+        self._replan_open = False
+        if not completed:
+            return False
+        if executed_action.shape != (self.schema.action_horizon, self.schema.action_dim):
+            raise ValueError("only a completed 16-step action window may enter transition memory")
+        assert self.memory.task_cluster is not None
+        self.memory.append(
+            TransitionRecord(
+                task_cluster=self.memory.task_cluster,
+                phase=phase,
+                visual_key=visual_key,
+                effect_post=effect_post,
+                executed_action=executed_action,
+                next_visual_key=next_visual_key,
+                next_phase=next_phase,
+                latent_index=int(latent_index),
+                schema_hash=self.schema.hash,
+                metadata=dict(metadata or {}),
+            )
+        )
+        return True

releases/bridge-native-v3/code/cosmos_framework/model/zeva/experimental/transition_memory_test.py

@@ -0,0 +1,161 @@
+from __future__ import annotations
+
+import pytest
+import torch
+
+from cosmos_framework.model.zeva.experimental.transition_memory import (
+    TransitionMemory,
+    TransitionMemoryEncoder,
+    TransitionMemorySchema,
+    TransitionRecord,
+)
+from cosmos_framework.model.zeva.experimental.transition_memory_runtime import TransitionMemoryController
+from cosmos_framework.model.zeva.experimental.robocasa_transition_memory import make_robocasa_atomic5_schema
+
+
+def make_transition(schema: TransitionMemorySchema, task: str, i: int, *, valid: bool = True) -> TransitionRecord:
+    return TransitionRecord(
+        task_cluster=task,
+        phase=torch.full((schema.phase_dim,), float(i)),
+        visual_key=torch.full((schema.visual_key_dim,), float(i + 1)),
+        effect_post=torch.full((schema.effect_dim,), float(i + 2)),
+        executed_action=torch.full((schema.action_horizon, schema.action_dim), float(i + 3)),
+        next_visual_key=torch.full((schema.visual_key_dim,), float(i + 4)),
+        next_phase=torch.full((schema.phase_dim,), float(i + 5)),
+        latent_index=4 * i,
+        schema_hash=schema.hash,
+        valid=valid,
+    )
+
+
+def test_fifo_reset_and_task_isolation() -> None:
+    schema = TransitionMemorySchema(capacity=2)
+    memory = TransitionMemory(schema)
+    memory.append(make_transition(schema, "pick", 0))
+    memory.append(make_transition(schema, "pick", 1))
+    memory.append(make_transition(schema, "pick", 2))
+    assert len(memory) == 2
+    with pytest.raises(ValueError, match="task mismatch"):
+        memory.append(make_transition(schema, "pour", 3))
+    memory.reset("pour")
+    memory.append(make_transition(schema, "pour", 3))
+    assert len(memory) == 1
+    assert memory.task_cluster == "pour"
+
+
+def test_schema_and_incomplete_transition_fail_fast() -> None:
+    schema = TransitionMemorySchema()
+    memory = TransitionMemory(schema)
+    bad_hash = make_transition(schema, "pick", 0)
+    bad_hash.schema_hash = "wrong"
+    with pytest.raises(ValueError, match="schema mismatch"):
+        memory.append(bad_hash)
+    with pytest.raises(ValueError, match="invalid or incomplete"):
+        memory.append(make_transition(schema, "pick", 0, valid=False))
+
+
+def test_query_is_bounded_and_encoder_handles_empty_memory() -> None:
+    schema = TransitionMemorySchema(capacity=64, top_k=4)
+    memory = TransitionMemory(schema)
+    for i in range(10):
+        memory.append(make_transition(schema, "pick", i))
+    entries, scores = memory.query(torch.zeros(schema.phase_dim), torch.zeros(schema.visual_key_dim))
+    assert len(entries) == 4
+    assert scores.shape == (4,)
+    encoder = TransitionMemoryEncoder(schema)
+    context = encoder(torch.zeros(schema.phase_dim), torch.zeros(schema.visual_key_dim), entries)
+    empty_context = encoder(torch.zeros(schema.phase_dim), torch.zeros(schema.visual_key_dim), [])
+    assert context.shape == (1, 256)
+    assert empty_context.shape == (1, 256)
+    assert torch.isfinite(context).all()
+
+
+def test_controller_does_not_commit_incomplete_replan() -> None:
+    schema = TransitionMemorySchema()
+    controller = TransitionMemoryController(schema, enabled=True)
+    controller.reset("pick")
+    controller.begin_replan()
+    readout = controller.read(torch.zeros(schema.phase_dim), torch.zeros(schema.visual_key_dim))
+    assert readout.num_entries == 0
+    assert not controller.complete_replan(
+        phase=torch.zeros(schema.phase_dim),
+        visual_key=torch.zeros(schema.visual_key_dim),
+        effect_post=torch.zeros(schema.effect_dim),
+        executed_action=torch.zeros(schema.action_horizon, schema.action_dim),
+        next_visual_key=torch.zeros(schema.visual_key_dim),
+        next_phase=torch.zeros(schema.phase_dim),
+        latent_index=0,
+        completed=False,
+    )
+    assert len(controller.memory) == 0
+
+
+def test_robocasa_atomic5_contract_is_arm7_and_temporally_bounded() -> None:
+    schema = make_robocasa_atomic5_schema(cte_hash="cte", vae_temporal_hash="vae")
+    assert schema.action_dim == 7
+    assert schema.action_horizon == 16
+    assert schema.capacity == 64
+    assert schema.top_k == 4
+    # The provenance fields participate in the hash; a cache from another
+    # A different CTE/VAE contract cannot be accepted by the runtime.
+    assert schema.hash != make_robocasa_atomic5_schema(cte_hash="other", vae_temporal_hash="vae").hash
+
+
+def test_robocasa_controller_writes_only_after_completed_window() -> None:
+    schema = make_robocasa_atomic5_schema()
+    controller = TransitionMemoryController(schema, enabled=True)
+    controller.reset("OpenStandMixerHead")
+    def zeros(shape: tuple[int, ...]) -> torch.Tensor:
+        return torch.zeros(shape)
+    controller.begin_replan()
+    controller.read(zeros((128,)), zeros((128,)))
+    assert not controller.complete_replan(
+        phase=zeros((128,)), visual_key=zeros((128,)), effect_post=zeros((128,)),
+        executed_action=zeros((16, 7)), next_visual_key=zeros((128,)),
+        next_phase=zeros((128,)), latent_index=0, completed=False,
+    )
+    assert len(controller.memory) == 0
+    controller.begin_replan()
+    controller.read(zeros((128,)), zeros((128,)))
+    assert controller.complete_replan(
+        phase=zeros((128,)), visual_key=zeros((128,)), effect_post=zeros((128,)),
+        executed_action=zeros((16, 7)), next_visual_key=zeros((128,)),
+        next_phase=zeros((128,)), latent_index=0, completed=True,
+    )
+    assert len(controller.memory) == 1
+
+
+def test_attempt_outcome_annotation_is_scoped_and_auditable() -> None:
+    schema = make_robocasa_atomic5_schema()
+    memory = TransitionMemory(schema)
+    first = make_transition(schema, "OpenStandMixerHead", 0)
+    first.metadata = {"attempt_id": 0, "session_id": "mixer:195", "environment_seed": 195}
+    second = make_transition(schema, "OpenStandMixerHead", 1)
+    second.metadata = {"attempt_id": 1, "session_id": "mixer:195", "environment_seed": 195}
+    memory.append(first)
+    memory.append(second)
+    updated = memory.annotate_attempt_outcome(
+        0,
+        outcome="failure",
+        termination_reason="max_steps",
+        total_steps=300,
+        final_progress=0.0,
+    )
+    assert updated == 1
+    assert memory.entries[0].metadata["terminal_outcome"] == "failure"
+    assert "terminal_outcome" not in memory.entries[1].metadata
+
+
+def test_discard_open_replan_never_commits_partial_window() -> None:
+    schema = make_robocasa_atomic5_schema()
+    controller = TransitionMemoryController(schema, enabled=False)
+    controller.reset("OpenStandMixerHead")
+    controller.begin_replan()
+    assert controller.discard_open_replan()
+    assert len(controller.memory) == 0
+    assert not controller.discard_open_replan()
+
+
+    expected = {(95, frame) for frame in range(85)}
+    assert (95, 83) in expected
+    assert (93, 83) not in expected