fix(train): ensure consistent color augmentation RNG across base and wrist cameras

Fixes #859. augmax.Chain splits RNG by transform count, so base cameras
(4 transforms) and wrist cameras (1 transform) received different sub-RNGs
for ColorJitter despite the same input seed. This caused inconsistent color
semantics between camera views.

Split augmentation into spatial (camera-specific) and color (shared) stages
with deterministic RNG derivation via jax.random.fold_in, ensuring all cameras
see the same ColorJitter parameters within each sample.

Author: wadeKeith

src/openpi/models/model.py

@@ -168,20 +168,23 @@ def preprocess_observation(
         if train:
             # Convert from [-1, 1] to [0, 1] for augmax.
             image = image / 2.0 + 0.5
+            sub_rngs = jax.random.split(rng, image.shape[0])
 
-            transforms = []
+            # Stage 1: Apply spatial transforms (base cameras only)
             if "wrist" not in key:
                 height, width = image.shape[1:3]
-                transforms += [
+                spatial_rngs = jax.vmap(lambda k: jax.random.fold_in(k, 0))(sub_rngs)
+                image = jax.vmap(augmax.Chain(
                     augmax.RandomCrop(int(width * 0.95), int(height * 0.95)),
                     augmax.Resize(width, height),
                     augmax.Rotate((-5, 5)),
-                ]
-            transforms += [
+                ))(spatial_rngs, image)
+
+            # Stage 2: Apply color transforms (all cameras, consistent RNG)
+            color_rngs = jax.vmap(lambda k: jax.random.fold_in(k, 1))(sub_rngs)
+            image = jax.vmap(augmax.Chain(
                 augmax.ColorJitter(brightness=0.3, contrast=0.4, saturation=0.5),
-            ]
-            sub_rngs = jax.random.split(rng, image.shape[0])
-            image = jax.vmap(augmax.Chain(*transforms))(sub_rngs, image)
+            ))(color_rngs, image)
 
             # Back to [-1, 1].
             image = image * 2.0 - 1.0