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

[harshitwandhare]

Addresses #174, the reorder you proposed as the cheaper first step.

What was slow

_compose_multi_view ran crop, resize and ColorJitter on cat([wrist, left, right]) at full size, then downscaled the two exterior views by 2x while tiling them. So the jitter paid for 3H x W pixels per frame to build a composite that only keeps 3H/2 x W of them. Half that work was thrown away by the next line.

What the change does

Splits the pipeline around the composition. Crop and resize stay per view and run first, because cropping the composite would cut across the tile boundaries. Only ColorJitter moves, onto the composite.

Both orderings draw one set of colour factors and apply it to every view, so the views still agree on lighting and per-sample colour diversity is unchanged. What changes is that the exterior views are jittered after downscaling rather than before. For a random augmentation that is a re-parameterization rather than a different distribution, which is the same reasoning you gave for the device swap in point (1).

Numbers

Measured after syncing onto main at #223. The branch has since merged main again at #226; that sync touched inference and MoT files only, none of them in this path, and the directory run below is from the current head. _compose_multi_view called directly on
a bare instance, 16 frames at 180x320 per view, float32, torch.set_num_threads(4),
20 calls per round after a warmup. The two variants alternate round by round, so machine
load drifts across both sides rather than favouring whichever ran first:

round 1: before=208.2  after=115.0
round 2: before=223.7  after=128.9
round 3: before=209.6  after=135.7
round 4: before=222.7  after=129.6
round 5: before=216.3  after=119.1

median before = 216.3 ms, median after = 128.9 ms, ratio = 1.68x
max(after) = 135.7 ms, min(before) = 208.2 ms

About 1.7x on the whole method. The sets do not overlap. This measures the full method
rather than an isolated ColorJitter call, so it carries the crop, resize, interpolate and
cat that did not change. That is why it lands short of the 2x the pixel count predicts.

An earlier run on this branch read 1.9x from three runs per side taken one side after the
other. Alternating the rounds gives 1.7x and I trust that number more, so it replaces the
first one here.

Output shape is identical either way, (16, 3, 270, 320).

The uint8 half is deliberately not here

You suggested casting to uint8 ahead of the jitter for another 4x off the memory traffic. Measured separately, that made the reordered path 1.35x slower, not faster: ColorJitter in torchvision 0.20.1 round-trips uint8 through float internally for the hue and saturation ops, so the conversion cost eats the bandwidth win. Numbers are in the issue thread. Happy to revisit if it behaves differently on GPU, but the workers being optimised here are CPU-bound.

Tests

New droid_lerobot_dataset_test.py, four L0 tests: the tiling itself, that augmentation preserves shape and does not bleed across tiles, that one colour draw covers all three views, and that the augmentor is not silently a no-op.

The three behavioural ones pin what the reorder must not change, so I ran them against main first, with the helper adapted to the pre-change attribute name, and they pass there too:

4 passed, 4 warnings in 32.91s

and on this branch:

4 passed, 4 warnings in 8.57s

Whole directory, to catch anything that read the old attribute:

$ python -m pytest cosmos_framework/data/generator/action/ -q
64 passed, 30 warnings in 44.65s

ruff check and ruff format --check clean on both files.

What I could not run

pre-commit run --files <the two files> passes every hook that applies except uv-lock-script, which fails on this machine because the hook shells out to /bin/bash. That looks like a Windows environment problem on my side rather than anything about this change, but I have not proved that, so flagging it rather than claiming the gate is green.

gpu-tests.yml needs a GPU runner and the cu128-train extras, so it has not run here.

Correcting something an earlier version of this description said. I claimed the new test file was not in the list that workflow runs and offered to wire it in. That was wrong. The unittest job runs the whole colocated suite in one invocation and its own comment says new tests are picked up automatically:

uv run --all-extras --group=cu128-train python -m pytest -v -s   cosmos_framework/ -o addopts=   --deselect cosmos_framework/model/generator/mot/attention_test.py::test_two_way_attention_flex_matches_dense_across_batch_shapes

The files named individually elsewhere in that workflow are the torchrun and deselect special cases, not the general path. Collection confirmed locally under the same command:

$ python -m pytest cosmos_framework/ -o addopts= --collect-only -q
cosmos_framework/data/generator/action/datasets/droid_lerobot_dataset_test.py::test_compose_multi_view_tiles_wrist_over_left_and_right
cosmos_framework/data/generator/action/datasets/droid_lerobot_dataset_test.py::test_compose_multi_view_keeps_layout_under_augmentation
cosmos_framework/data/generator/action/datasets/droid_lerobot_dataset_test.py::test_color_jitter_is_shared_across_the_three_views
cosmos_framework/data/generator/action/datasets/droid_lerobot_dataset_test.py::test_augmentation_actually_changes_the_composite

So nothing needs wiring. The four are CPU-only and run on that job alongside the rest of the colocated suite.