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()