Apply the same augmentation chain to all camera images

src/openpi/models/model.py

@@ -169,15 +169,13 @@ def preprocess_observation(
             # Convert from [-1, 1] to [0, 1] for augmax.
             image = image / 2.0 + 0.5
 
-            transforms = []
-            if "wrist" not in key:
-                height, width = image.shape[1:3]
-                transforms += [
-                    augmax.RandomCrop(int(width * 0.95), int(height * 0.95)),
-                    augmax.Resize(width, height),
-                    augmax.Rotate((-5, 5)),
-                ]
-            transforms += [
+            height, width = image.shape[1:3]
+            # Follows pi0.5 Appendix E; the shared rng keeps per-frame augmentation
+            # parameters consistent across cameras.
+            transforms = [
+                augmax.RandomCrop(int(width * 0.95), int(height * 0.95)),
+                augmax.Resize(width, height),
+                augmax.Rotate((-5, 5)),
                 augmax.ColorJitter(brightness=0.3, contrast=0.4, saturation=0.5),
             ]
             sub_rngs = jax.random.split(rng, image.shape[0])

src/openpi/models/model_test.py

@@ -1,5 +1,7 @@
 from flax import nnx
 import jax
+import jax.numpy as jnp
+import numpy as np
 import pytest
 
 from openpi.models import model as _model
@@ -9,6 +11,28 @@
 from openpi.shared import nnx_utils
 
 
+def test_preprocess_observation_augmentation_consistent_across_cameras():
+    batch_size = 2
+    image = jnp.linspace(-1.0, 1.0, batch_size * 224 * 224 * 3, dtype=jnp.float32).reshape(batch_size, 224, 224, 3)
+    obs = _model.Observation(
+        images=dict.fromkeys(_model.IMAGE_KEYS, image),
+        image_masks={key: jnp.ones((batch_size,), dtype=jnp.bool) for key in _model.IMAGE_KEYS},
+        state=jnp.zeros((batch_size, 1), dtype=jnp.float32),
+    )
+
+    augmented = _model.preprocess_observation(jax.random.key(0), obs, train=True)
+
+    np.testing.assert_array_equal(augmented.images["base_0_rgb"], augmented.images["left_wrist_0_rgb"])
+    np.testing.assert_array_equal(augmented.images["base_0_rgb"], augmented.images["right_wrist_0_rgb"])
+    assert not np.array_equal(np.asarray(augmented.images["base_0_rgb"]), np.asarray(image))
+    assert jnp.min(augmented.images["base_0_rgb"]) >= -1.0 - 1e-6
+    assert jnp.max(augmented.images["base_0_rgb"]) <= 1.0 + 1e-6
+
+    not_augmented = _model.preprocess_observation(jax.random.key(0), obs, train=False)
+    for key in _model.IMAGE_KEYS:
+        np.testing.assert_array_equal(not_augmented.images[key], image)
+
+
 def test_pi0_model():
     key = jax.random.key(0)
     config = pi0_config.Pi0Config()