Release 2026-08-31 (i4 592d5717)

.file_mapping.json

@@ -1,7 +1,7 @@
 {
-  "_source_commit": "7597afdd014f52b5f6756e02a238608ed18d8522-dirty",
-  "_dest_commit": "505041197b257b5bace2405415a3f3d5512fc94d",
-  "_generated_at": "2026-08-28T09:10:22Z",
+  "_source_commit": "592d5717b8e3c95ebc92dd9d82e9f2192fc90c1b-dirty",
+  "_dest_commit": "0e034bc98ffa3c3dfa19f037871f3a8bbc1c4d05",
+  "_generated_at": "2026-08-31T13:26:08Z",
   "files": {
     "imaginaire/__init__.py": "cosmos_framework/__init__.py",
     "imaginaire/attention/__init__.py": "cosmos_framework/model/attention/__init__.py",
@@ -108,7 +108,6 @@
     "imaginaire/trainer.py": "cosmos_framework/trainer/__init__.py",
     "imaginaire/utils/__init__.py": "cosmos_framework/utils/__init__.py",
     "imaginaire/utils/callback.py": "cosmos_framework/utils/callback.py",
-    "imaginaire/utils/callback_test.py": "cosmos_framework/utils/callback_test.py",
     "imaginaire/utils/checkpoint_db.py": "cosmos_framework/utils/checkpoint_db.py",
     "imaginaire/utils/checkpointer.py": "cosmos_framework/utils/checkpointer.py",
     "imaginaire/utils/cluster_env.py": "cosmos_framework/utils/cluster_env.py",
@@ -257,6 +256,7 @@
     "projects/cosmos3/cosmos3/configs/distillation_export.py": "cosmos_framework/inference/common/distillation_export.py",
     "projects/cosmos3/cosmos3/configs/distillation_export_test.py": "cosmos_framework/inference/common/distillation_export_test.py",
     "projects/cosmos3/cosmos3/cost_model/__init__.py": "cosmos_framework/utils/generator/cost_model/__init__.py",
+    "projects/cosmos3/cosmos3/cost_model/benchmark.py": "cosmos_framework/utils/generator/cost_model/benchmark.py",
     "projects/cosmos3/cosmos3/cost_model/budget.py": "cosmos_framework/utils/generator/cost_model/budget.py",
     "projects/cosmos3/cosmos3/cost_model/estimator.py": "cosmos_framework/utils/generator/cost_model/estimator.py",
     "projects/cosmos3/cosmos3/cost_model/spec.py": "cosmos_framework/utils/generator/cost_model/spec.py",
@@ -270,6 +270,7 @@
     "projects/cosmos3/cosmos3/datasets/action/utils/pose_utils_test.py": "cosmos_framework/data/generator/action/utils/pose_utils_test.py",
     "projects/cosmos3/cosmos3/datasets/action/utils/transforms.py": "cosmos_framework/data/generator/action/utils/transforms.py",
     "projects/cosmos3/cosmos3/datasets/action/utils/transforms_test.py": "cosmos_framework/data/generator/action/utils/transforms_test.py",
+    "projects/cosmos3/cosmos3/datasets/action/utils/unified_action_schema.py": "cosmos_framework/data/generator/action/utils/unified_action_schema.py",
     "projects/cosmos3/cosmos3/datasets/action/utils/viewpoint_utils.py": "cosmos_framework/data/generator/action/utils/viewpoint_utils.py",
     "projects/cosmos3/cosmos3/datasets/augmentor_provider.py": "cosmos_framework/data/generator/augmentor_provider.py",
     "projects/cosmos3/cosmos3/datasets/augmentors/__init__.py": "cosmos_framework/data/generator/augmentors/__init__.py",
@@ -372,6 +373,8 @@
     "projects/cosmos3/cosmos3/models/mot/und_k_norm_example_test.py": "cosmos_framework/model/generator/mot/und_k_norm_example_test.py",
     "projects/cosmos3/cosmos3/models/mot/unified_mot.py": "cosmos_framework/model/generator/mot/unified_mot.py",
     "projects/cosmos3/cosmos3/models/mot/unified_mot_test.py": "cosmos_framework/model/generator/mot/unified_mot_test.py",
+    "projects/cosmos3/cosmos3/models/mot/vae_load_balance.py": "cosmos_framework/model/generator/mot/vae_load_balance.py",
+    "projects/cosmos3/cosmos3/models/mot/vae_load_balance_test.py": "cosmos_framework/model/generator/mot/vae_load_balance_test.py",
     "projects/cosmos3/cosmos3/models/omni_mot_model.py": "cosmos_framework/model/generator/omni_mot_model.py",
     "projects/cosmos3/cosmos3/models/parallelize_vlm.py": "cosmos_framework/model/generator/parallelize_vlm.py",
     "projects/cosmos3/cosmos3/models/reasoner/__init__.py": "cosmos_framework/model/generator/reasoner/__init__.py",
@@ -427,6 +430,7 @@
     "projects/cosmos3/cosmos3/models/utils/safetensors_loader.py": "cosmos_framework/model/generator/utils/safetensors_loader.py",
     "projects/cosmos3/cosmos3/models/utils/safetensors_loader_indexed_test.py": "cosmos_framework/model/generator/utils/safetensors_loader_indexed_test.py",
     "projects/cosmos3/cosmos3/models/utils/safetensors_loader_test.py": "cosmos_framework/model/generator/utils/safetensors_loader_test.py",
+    "projects/cosmos3/cosmos3/models/vision_encoder.py": "cosmos_framework/model/generator/vision_encoder.py",
     "projects/cosmos3/cosmos3/models/vlm_model.py": "cosmos_framework/model/generator/vlm_model.py",
     "projects/cosmos3/cosmos3/processors/__init__.py": "cosmos_framework/data/generator/processors/__init__.py",
     "projects/cosmos3/cosmos3/processors/audio_utils.py": "cosmos_framework/data/generator/processors/audio_utils.py",

cosmos_framework/callbacks/every_n_draw_sample.py

@@ -283,7 +283,8 @@ def _decode_transfer_pixel_row(
 
     ``decode`` defaults to the model's main VAE; a joint camera + LiDAR sample passes
     ``model.decode_lidar`` for range clips because the two streams have separate VAEs.
-    LiDAR never uses the per-camera encode path, even when the camera items do.
+    LiDAR callers pass ``decode_per_view=False``: V0 and V1 are both 1x temporal on a
+    single range view, not 4x camera-major WAN clips.
     """
     decode = decode if decode is not None else model.decode
     if decode_per_view is None:

cosmos_framework/callbacks/training_stats.py

@@ -10,14 +10,20 @@
 from cosmos_framework.model._base import ImaginaireModel
 from cosmos_framework.utils import distributed
 from cosmos_framework.utils.callback import Callback
+from cosmos_framework.model.generator.algorithm.loss.flow_matching import (
+    ACTION_SLOT_SAMPLE_COUNT_KEY,
+    ACTION_SLOT_SAMPLE_LOSS_KEY,
+)
 from cosmos_framework.callbacks.wandb_log import _LossRecord
 from cosmos_framework.data.generator.action.utils.domain_utils import EMBODIMENT_TO_DOMAIN_ID
+from cosmos_framework.data.generator.action.utils.unified_action_schema import UNIFIED_ACTION_SLOT_GROUPS
 
 # Build inverse mapping: domain_id -> embodiment_type. First occurrence wins when multiple embodiment names share the
 # same domain id.
 DOMAIN_ID_TO_EMBODIMENT: dict[int, str] = {}
 for _k, _v in EMBODIMENT_TO_DOMAIN_ID.items():
     DOMAIN_ID_TO_EMBODIMENT.setdefault(_v, _k)
+NUM_ACTION_DOMAINS = max(DOMAIN_ID_TO_EMBODIMENT) + 1
 
 
 class TrainingStatsCallback(Callback):
@@ -32,6 +38,8 @@ def __init__(self, log_freq: int = 100):
         self._embodiment_total_count: int = 0
         self._per_embodiment_loss: dict[str, _LossRecord] = {}
         self._per_embodiment_sub_loss: dict[str, dict[str, _LossRecord]] = {}
+        self._action_slot_family_stats: dict[str, torch.Tensor] = {}
+        self._action_slot_rank_stats: torch.Tensor | None = None
 
     def _accumulate_mode_counts(self, data_batch: dict[str, torch.Tensor]) -> None:
         modes = data_batch.get("mode", None)
@@ -186,12 +194,136 @@ def _accumulate_per_embodiment_loss(
         if embodiment not in self._per_embodiment_sub_loss:
             self._per_embodiment_sub_loss[embodiment] = {}
         for key in output_batch:
+            if key.startswith("_"):
+                continue
             if "loss" in key and "per_instance" not in key:
                 if key not in self._per_embodiment_sub_loss[embodiment]:
                     self._per_embodiment_sub_loss[embodiment][key] = _LossRecord()
                 self._per_embodiment_sub_loss[embodiment][key].loss += output_batch[key].detach().float()
                 self._per_embodiment_sub_loss[embodiment][key].iter_count += 1
 
+    def _accumulate_rank_action_slot_stats(self, sample_loss: torch.Tensor, sample_count: torch.Tensor) -> None:
+        """Accumulate optimizer-aligned rank/microbatch means."""
+        # Match the optimizer reduction: every contributing rank-local
+        # microbatch mean has equal weight, independent of its sample count.
+        rank_sample_count = sample_count.sum(dim=0)
+        rank_has_slot = rank_sample_count.gt(0)
+        if self._action_slot_rank_stats is None:
+            self._action_slot_rank_stats = sample_loss.new_zeros(
+                (2, len(UNIFIED_ACTION_SLOT_GROUPS)), dtype=torch.float32
+            )
+        self._action_slot_rank_stats[0].add_(sample_loss.sum(dim=0) / rank_sample_count.clamp(min=1))
+        self._action_slot_rank_stats[1].add_(rank_has_slot.to(dtype=torch.float32))
+
+    def _accumulate_family_action_slot_stats(
+        self, sample_loss: torch.Tensor, sample_count: torch.Tensor, dataset_names: list[str]
+    ) -> None:
+        """Accumulate sample-weighted statistics for each action dataset."""
+        for dataset_name in sorted(set(dataset_names)):
+            row_indices = [index for index, name in enumerate(dataset_names) if name == dataset_name]
+            index = torch.tensor(row_indices, device=sample_loss.device)
+            stats = self._action_slot_family_stats.get(dataset_name)
+            if stats is None:
+                stats = sample_loss.new_zeros((2, len(UNIFIED_ACTION_SLOT_GROUPS)), dtype=torch.float32)
+                self._action_slot_family_stats[dataset_name] = stats
+            stats[0].add_(sample_loss.index_select(0, index).sum(dim=0))
+            stats[1].add_(sample_count.index_select(0, index).sum(dim=0))
+
+    def _accumulate_action_slot_stats(self, output_batch: dict[str, torch.Tensor]) -> None:
+        """Accumulate global and per-dataset action-slot diagnostics."""
+        sample_loss = output_batch.get(ACTION_SLOT_SAMPLE_LOSS_KEY)
+        sample_count = output_batch.get(ACTION_SLOT_SAMPLE_COUNT_KEY)
+        if (
+            not isinstance(sample_loss, torch.Tensor)
+            or not isinstance(sample_count, torch.Tensor)
+            or sample_loss.ndim != 2
+            or sample_loss.shape != sample_count.shape
+            or sample_loss.shape[1] != len(UNIFIED_ACTION_SLOT_GROUPS)
+        ):
+            return
+
+        sample_loss = sample_loss.detach().float()
+        sample_count = sample_count.detach().float()
+        self._accumulate_rank_action_slot_stats(sample_loss, sample_count)
+
+        dataset_names = output_batch.get("_action_family")
+        if isinstance(dataset_names, str):
+            dataset_names = [dataset_names]
+        if (
+            not isinstance(dataset_names, (list, tuple))
+            or len(dataset_names) != sample_loss.shape[0]
+            or any(not isinstance(name, str) or not name for name in dataset_names)
+        ):
+            return
+        self._accumulate_family_action_slot_stats(sample_loss, sample_count, list(dataset_names))
+
+    def _compute_action_slot_loss_stats(self, log_prefix: str) -> dict[str, float]:
+        """Aggregate rank-weighted global and sample-weighted per-dataset slot losses."""
+        local_dataset_stats = {
+            dataset_name: stats.cpu().tolist() for dataset_name, stats in self._action_slot_family_stats.items()
+        }
+        local_rank_stats = self._action_slot_rank_stats
+        if local_rank_stats is None:
+            local_rank_stats = torch.zeros(2, len(UNIFIED_ACTION_SLOT_GROUPS), dtype=torch.float32)
+        local_payload = (local_dataset_stats, local_rank_stats.cpu().tolist())
+        if dist.is_available() and dist.is_initialized():
+            gathered: list[tuple[dict[str, list[list[float]]], list[list[float]]] | None] = [
+                None for _ in range(dist.get_world_size())
+            ]
+            dist.all_gather_object(gathered, local_payload)
+        else:
+            gathered = [local_payload]
+
+        family_totals: dict[str, torch.Tensor] = {}
+        rank_totals = torch.zeros(2, len(UNIFIED_ACTION_SLOT_GROUPS), dtype=torch.float64)
+        for payload in gathered:
+            if payload is None:
+                continue
+            dataset_stats, rank_stats = payload
+            rank_totals.add_(torch.tensor(rank_stats, dtype=torch.float64))
+            for dataset_name, stats in dataset_stats.items():
+                family_totals.setdefault(
+                    dataset_name,
+                    torch.zeros(2, len(UNIFIED_ACTION_SLOT_GROUPS), dtype=torch.float64),
+                ).add_(torch.tensor(stats, dtype=torch.float64))
+
+        result: dict[str, float] = {}
+        for slot_index, (slot_name, _) in enumerate(UNIFIED_ACTION_SLOT_GROUPS):
+            count = rank_totals[1, slot_index].item()
+            if count > 0:
+                result[f"{log_prefix}_stats_loss_action_slot/{slot_name}"] = rank_totals[0, slot_index].item() / count
+        for dataset_name, stats in sorted(family_totals.items()):
+            sample_loss_sum = stats[0]
+            sample_count = stats[1]
+            for slot_index, (slot_name, _) in enumerate(UNIFIED_ACTION_SLOT_GROUPS):
+                count = sample_count[slot_index].item()
+                if count > 0:
+                    result[f"{log_prefix}_stats_action_slot_loss/{dataset_name}/{slot_name}"] = (
+                        sample_loss_sum[slot_index].item() / count
+                    )
+                    result[f"{log_prefix}_stats_action_slot_count/{dataset_name}/{slot_name}"] = count
+        self._action_slot_family_stats = {}
+        self._action_slot_rank_stats = None
+        return result
+
+    @torch.no_grad()
+    def on_training_step_batch_end(
+        self,
+        model: ImaginaireModel,
+        data_batch: dict[str, torch.Tensor],
+        output_batch: dict[str, torch.Tensor],
+        loss: torch.Tensor,
+        iteration: int = 0,
+    ) -> None:
+        parallel_dims = getattr(model, "parallel_dims", None)
+        if (
+            parallel_dims is not None
+            and getattr(parallel_dims, "cp_enabled", False)
+            and getattr(parallel_dims, "cp_rank", 0) != 0
+        ):
+            return
+        self._accumulate_action_slot_stats(output_batch)
+
     def _compute_per_embodiment_loss_stats(self, log_prefix: str) -> dict[str, float]:
         """Compute per-embodiment loss averages across all ranks.
 
@@ -274,6 +406,7 @@ def on_training_step_end(
         mode_total, mode_counts = self._gather_global_mode_counts()
         embodiment_total, embodiment_counts = self._gather_global_embodiment_counts()
         per_embodiment_loss_dict = self._compute_per_embodiment_loss_stats(log_prefix="train")
+        action_slot_loss_dict = self._compute_action_slot_loss_stats(log_prefix="train")
 
         if not distributed.is_rank0():
             return
@@ -291,5 +424,6 @@ def on_training_step_end(
             )
         )
         log_dict.update(per_embodiment_loss_dict)
+        log_dict.update(action_slot_loss_dict)
 
         wandb.log({k: float(v) for k, v in log_dict.items()}, step=iteration)

cosmos_framework/configs/base/defaults/flex_attention.py

@@ -14,7 +14,7 @@
 
 import attrs
 
-# Which noisy tokens a noisy token attends to, among those of its own sample.
+# Which RGB tokens an RGB token attends to, among those of its own sample.
 #
 # The noisy square is the largest quadrant of the multiview mask -- every other rule is already
 # confined to a single ``(frame, view)`` cell -- so this is the choice that sets what the mask
@@ -25,10 +25,12 @@
 #   the default. Cost is (FVS)^2.
 # * ``"same_view"``: ``1/V`` of them. Each camera only attends to its own noisy tokens. Cost
 #   is V*(FS)^2.
-# * ``"same_view_or_frame"``: Each camera attends to its own noisy tokens plus the same frame in
-#   every other camera. Cost is V*(FS)^2 + F*(VS)^2. Not allowed on a joint camera + LiDAR pack:
-#   the two streams do not share a frame index.
-AttentionScope = Literal["all_views", "same_view", "same_view_or_frame"]
+# * ``"decomposed"``: Each camera attends to its own noisy tokens plus the same frame in
+#   every other camera, which decomposes the square into a temporal half and a spatial one.
+#   Cost is V*(FS)^2 + F*(VS)^2. Rejected on a joint camera + LiDAR pack unless
+#   ``decomposed_temporal_window_seconds`` is set: the two streams do not share a frame
+#   index, but they do share real capture time, which the window compares instead.
+AttentionScope = Literal["all_views", "same_view", "decomposed"]
 
 # The scopes of ``AttentionScope`` at runtime, which the annotation itself is not.
 ATTENTION_SCOPES = get_args(AttentionScope)
@@ -49,12 +51,11 @@ class FlexAttentionMaskConfig:
     view.
     """
 
-    # Which same-kind (RGB) tokens of its sample a token attends to: conditioning tokens
-    # reach conditioning tokens, and noisy tokens reach both noisy and conditioning tokens
-    # alike. Cross-view attention is what lets the rig agree with itself, so the full square
-    # is the default; the narrower scopes buy attention that grows with the rig rather than
-    # with its square, per the comment above. Never widens a WSM (World Scenario Map) control
-    # token's reach, which is always its own view -- see
+    # Which RGB tokens of its sample an RGB token attends to, independent of whether the
+    # query or key is conditioning. Cross-view attention is what lets the rig agree with
+    # itself, so the full square is the default; the narrower scopes buy attention that grows
+    # with the rig rather than with its square, per the comment above. Never widens a WSM
+    # (World Scenario Map) control token's reach, which is always its own view -- see
     # flex_attention.build_multiview_flex_metadata's ``is_control_per_item``, which the
     # network derives per generation stream, and which a batch without a control stream
     # leaves empty.
@@ -63,6 +64,17 @@ class FlexAttentionMaskConfig:
         validator=attrs.validators.in_(ATTENTION_SCOPES),
     )
 
+    # Only read under attention_scope="decomposed". Replaces that scope's temporal half --
+    # "the query's own frame index" -- with "any key within this many seconds at or before the
+    # query's own capture time", i.e. 0 <= query_timestamp - key_timestamp <= this value. None
+    # (the default) keeps the frame-index form, which only agrees across sensors that share one
+    # clock; a joint camera + LiDAR pack needs a window instead; see
+    # flex_attention.build_multiview_flex_metadata and ._multiview_pair_predicate.
+    decomposed_temporal_window_seconds: float | None = attrs.field(
+        default=None,
+        validator=attrs.validators.optional(attrs.validators.ge(0)),
+    )
+
 
 @attrs.define(slots=False)
 class FlexAttentionConfig:

cosmos_framework/configs/base/defaults/parallelism.py

@@ -43,6 +43,17 @@ class ParallelismConfig:
     # Number of ranks for CFG parallelism.
     cfg_parallel_shard_degree: int = 1
 
+    # Size of the group ranks exchange raw video samples within before the VAE encode, to
+    # equalize each rank's predicted VAE-encode cost for the step (see
+    # models.mot.vae_load_balance and OmniMoTModel._prepare_training_data). 1 (the default)
+    # disables load balancing entirely -- same convention as context_parallel_shard_degree
+    # and cfg_parallel_shard_degree. Must evenly divide world_size. Unlike cp/cfgp this is
+    # NOT an attention-sharding overlay: it only ever runs once per cp-window (on the step
+    # that actually calls the VAE encoder), and has no interaction with how the packed
+    # sequence is sharded for compute. Size it to stay within one node (NVLink) to keep the
+    # raw-pixel exchange cheap -- nothing here enforces that, it is a placement choice.
+    vae_load_balance_group_size: int = 1
+
     # Inference-mode mesh toggle for ParallelDims.
     enable_inference_mode: bool = False
 

cosmos_framework/data/generator/action/utils/action_processing.py

@@ -321,6 +321,28 @@ def load_action_normalizer(
     )
 
 
+def make_camera_global_asinh_normalizer(
+    stats_path: str | Path = GLOBAL_ASINH_METRIC_GRIPPER_STATS_PATH,
+    *,
+    pose_convention: str = "backward_anchored",
+) -> ActionAsinhNormalization:
+    """Build the native 9D camera normalizer from the canonical global-asinh profile."""
+    if pose_convention not in {
+        "backward_anchored",
+        "backward_chunk_anchored_8f",
+        "backward_chunk_anchored_16f",
+    }:
+        raise ValueError(f"global_asinh normalization requires an anchored pose convention, got {pose_convention!r}")
+    stats = load_action_normalization_stats(
+        "global_asinh",
+        stats_path=stats_path,
+        expected_dim=59,
+    )
+    lo = stats["q01"][:9]
+    hi = stats["q99"][:9]
+    return ActionAsinhNormalization(offset=(hi + lo) / 2.0, scale=(hi - lo).clamp(min=1e-8) / 2.0)
+
+
 def make_pose_action_scale_normalizer(
     action_dim: int,
     *,