From 0724cd399de69c659f986fd01b6601b52af91a0d Mon Sep 17 00:00:00 2001
From: Advait Jayant <advaitjk@gmail.com>
Date: Fri, 12 Jun 2026 01:18:03 +0100
Subject: [PATCH] Apply the same augmentation chain to all camera images

augmax.Chain splits its rng once per sub-transform, so with 4 transforms
on the base camera and 1 on wrist cameras, ColorJitter drew different
parameters for base vs wrist views of the same frame even though the
same rng is passed for every camera. The pi0.5 paper (Appendix E)
applies the full crop/resize/rotate/jitter chain to all input images,
so this change applies the full chain to every camera, which also makes
augmentation parameters consistent across cameras within a frame. Adds
a regression test for cross-camera consistency.

Fixes #859
---
 src/openpi/models/model.py      | 16 +++++++---------
 src/openpi/models/model_test.py | 24 ++++++++++++++++++++++++
 2 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/src/openpi/models/model.py b/src/openpi/models/model.py
index 29618b4945..59cbb3e190 100644
--- a/src/openpi/models/model.py
+++ b/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])
diff --git a/src/openpi/models/model_test.py b/src/openpi/models/model_test.py
index 495dc18b5f..6f6b798d91 100644
--- a/src/openpi/models/model_test.py
+++ b/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()