perf(action): jitter the DROID composite instead of the three full-size views

cosmos_framework/data/generator/action/datasets/droid_lerobot_dataset.py

@@ -183,7 +183,8 @@ def __init__(
             with open(self._filter_dict_path) as f:
                 self._filter_dict = json.load(f)
 
-        self._image_augmentor: T.Compose | None = None
+        self._geometric_augmentor: T.Compose | None = None
+        self._color_augmentor: T.ColorJitter | None = None
 
         # Eager source registration. i4 defers this to its own dataloader's
         # ActionUnifiedIterableDataset.assign_worker(); cosmos-framework instead
@@ -313,6 +314,18 @@ def _compose_multi_view(self, sample: dict[str, Any]) -> torch.Tensor:
         Left and right exterior cameras are downscaled by 2x so that they
         tile to the same width as the wrist view. The output height is 3H/2.
 
+        Augmentation is split around the composition. The random crop and
+        resize stay per view and run first, because cropping the composite
+        would cut across the tile boundaries. ColorJitter runs last, on the
+        composite, where the exterior views have already been downscaled: that
+        is 3H/2 x W pixels instead of the 3H x W of the three full-size views,
+        so it costs roughly half as much for the same augmentation. Both
+        orderings draw one set of colour factors and apply it to every view, so
+        the views still agree on lighting; what changes is that the exterior
+        views are jittered after downscaling rather than before, which for a
+        random augmentation is a re-parameterization rather than a different
+        distribution. See issue #174.
+
         Returns:
             Composited raw video tensor in ``(T,C,H_out,W)`` float format.
         """
@@ -321,17 +334,17 @@ def _compose_multi_view(self, sample: dict[str, Any]) -> torch.Tensor:
         right = sample[self._image_features["right"]]  # [T,C,H_r,W_r]
 
         if self._use_image_augmentation:
-            if self._image_augmentor is None:
+            if self._geometric_augmentor is None:
                 _, _, h, w = wrist.shape
-                self._image_augmentor = T.Compose(
+                self._geometric_augmentor = T.Compose(
                     [
                         T.RandomCrop((int(h * 0.95), int(w * 0.95))),
                         T.Resize((h, w), antialias=True),
-                        T.ColorJitter(brightness=0.3, contrast=0.4, saturation=0.5, hue=0.08),
                     ]
                 )
+                self._color_augmentor = T.ColorJitter(brightness=0.3, contrast=0.4, saturation=0.5, hue=0.08)
             n, m = wrist.shape[0], wrist.shape[0] + left.shape[0]
-            combined = self._image_augmentor(torch.cat([wrist, left, right], dim=0))
+            combined = self._geometric_augmentor(torch.cat([wrist, left, right], dim=0))
             wrist, left, right = combined[:n], combined[n:m], combined[m:]
 
         _, _, h_w, w_w = wrist.shape
@@ -342,6 +355,10 @@ def _compose_multi_view(self, sample: dict[str, Any]) -> torch.Tensor:
         bottom = torch.cat([left, right], dim=-1)  # [T,C,H/2,W]
 
         composite = torch.cat([wrist, bottom], dim=-2)  # [T,C,3H/2,W]
+
+        if self._use_image_augmentation:
+            composite = self._color_augmentor(composite)
+
         return composite  # [T,C,3H/2,W]
 
     def _build_action_spec(self) -> ActionSpec:

cosmos_framework/data/generator/action/datasets/droid_lerobot_dataset_test.py

@@ -0,0 +1,84 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from cosmos_framework.data.generator.action.datasets.droid_lerobot_dataset import DROIDLeRobotDataset
+
+_IMAGE_FEATURES = {"wrist": "wrist_key", "left": "left_key", "right": "right_key"}
+
+
+def _dataset(use_image_augmentation: bool) -> DROIDLeRobotDataset:
+    """Build a bare instance exercising only ``_compose_multi_view``.
+
+    ``__init__`` opens LeRobot datasets from disk, which this test does not need:
+    the method under test reads three tensors out of ``sample`` and the
+    augmentation attributes, and touches nothing else.
+    """
+    dataset = object.__new__(DROIDLeRobotDataset)
+    dataset._image_features = _IMAGE_FEATURES
+    dataset._use_image_augmentation = use_image_augmentation
+    dataset._geometric_augmentor = None
+    dataset._color_augmentor = None
+    return dataset
+
+
+def _sample(wrist: float, left: float, right: float, t: int = 2, h: int = 16, w: int = 32) -> dict:
+    return {
+        "wrist_key": torch.full((t, 3, h, w), wrist),
+        "left_key": torch.full((t, 3, h, w), left),
+        "right_key": torch.full((t, 3, h, w), right),
+    }
+
+
+@pytest.mark.L0
+def test_compose_multi_view_tiles_wrist_over_left_and_right() -> None:
+    """Wrist spans the top, left and right tile the bottom at half scale."""
+    composite = _dataset(use_image_augmentation=False)._compose_multi_view(_sample(0.1, 0.5, 0.9))
+
+    assert composite.shape == (2, 3, 24, 32)  # [T,C,3H/2,W]
+    torch.testing.assert_close(composite[:, :, :16, :], torch.full((2, 3, 16, 32), 0.1))
+    torch.testing.assert_close(composite[:, :, 16:, :16], torch.full((2, 3, 8, 16), 0.5))
+    torch.testing.assert_close(composite[:, :, 16:, 16:], torch.full((2, 3, 8, 16), 0.9))
+
+
+@pytest.mark.L0
+def test_compose_multi_view_keeps_layout_under_augmentation() -> None:
+    """Augmentation must not change the composite's shape or tiling."""
+    torch.manual_seed(0)
+    composite = _dataset(use_image_augmentation=True)._compose_multi_view(_sample(0.1, 0.5, 0.9))
+
+    assert composite.shape == (2, 3, 24, 32)
+    # Each region is still uniform, so the tiles did not bleed into each other.
+    for region in (composite[:, :, :16, :], composite[:, :, 16:, :16], composite[:, :, 16:, 16:]):
+        assert torch.allclose(region, region.flatten()[0].expand_as(region), atol=1e-5)
+
+
+@pytest.mark.L0
+def test_color_jitter_is_shared_across_the_three_views() -> None:
+    """One colour draw covers the whole composite.
+
+    Jittering the views independently would cut colour diversity across the
+    frame and make the three cameras disagree on lighting, so identical input
+    views have to stay identical after augmentation.
+    """
+    torch.manual_seed(0)
+    composite = _dataset(use_image_augmentation=True)._compose_multi_view(_sample(0.4, 0.4, 0.4))
+
+    wrist_px = composite[:, :, :16, :].flatten()[0]
+    torch.testing.assert_close(composite[:, :, 16:, :16], torch.full((2, 3, 8, 16), wrist_px))
+    torch.testing.assert_close(composite[:, :, 16:, 16:], torch.full((2, 3, 8, 16), wrist_px))
+
+
+@pytest.mark.L0
+def test_augmentation_actually_changes_the_composite() -> None:
+    """Guards against the augmentor silently becoming a no-op."""
+    sample = _sample(0.4, 0.6, 0.8)
+    plain = _dataset(use_image_augmentation=False)._compose_multi_view(sample)
+    torch.manual_seed(0)
+    augmented = _dataset(use_image_augmentation=True)._compose_multi_view(sample)
+
+    assert not torch.allclose(plain, augmented)