I don't suspect it was an intended behavior, because it would probably be documented somewhere. Below please find a minimal script that visualizes the problem. While the potential fix seems to be trivial, it might impact the training dynamics, so please have a deeper look (e.g. do we want heavier color-reliance between base and wrists cameras? Or maybe this "bug" helps to regularize the model?)
"""
Reproduce augmetnation RNG bug: base and wrist cameras get different ColorJitter.
augmax.Chain splits RNG by transform count. Base uses 4 transforms, wrist uses 1,
so ColorJitter gets different sub-RNGs despite same input seed.
"""
import argparse
import augmax
import cv2
import jax
import jax.numpy as jnp
import numpy as np
W, H = 64, 40
N = 10
SEP = 40
def buggy(rng, img, is_wrist, hue):
t = [] if is_wrist else [
augmax.RandomCrop(int(W * 0.95), int(H * 0.95)),
augmax.Resize(W, H),
augmax.Rotate((-5, 5)),
]
t += [augmax.ColorJitter(brightness=0.3, contrast=0.4, saturation=0.5, hue=hue)]
return augmax.Chain(*t)(rng, img)
def fixed(rng, img, is_wrist, hue):
rng_s, rng_c = jax.random.split(rng)
if not is_wrist:
img = augmax.Chain(
augmax.RandomCrop(int(W * 0.95), int(H * 0.95)),
augmax.Resize(W, H),
augmax.Rotate((-5, 5)),
)(rng_s, img)
return augmax.Chain(augmax.ColorJitter(brightness=0.3, contrast=0.4, saturation=0.5, hue=hue))(rng_c, img)
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--hue", type=float, default=0.1)
args = parser.parse_args()
# R/G/B stripes
img = np.zeros((H, W, 3), dtype=np.float32)
img[:, :W//3, 0] = 0.8
img[:, W//3:2*W//3, 1] = 0.8
img[:, 2*W//3:, 2] = 0.8
cams = ["base", "l_wrist", "r_wrist"]
u8 = lambda x: (np.array(x).clip(0, 1) * 255).astype(np.uint8)
vsep = np.full((H, 1, 3), SEP, dtype=np.uint8)
hsep = lambda w: np.full((1, w, 3), SEP, dtype=np.uint8)
rows_b, rows_f = [], []
for seed in range(N):
rng = jax.random.key(seed * 100)
rb, rf = [u8(img)], [u8(img)]
for c in cams:
w = "wrist" in c
rb += [vsep, u8(buggy(rng, jnp.array(img), w, args.hue))]
rf += [vsep, u8(fixed(rng, jnp.array(img), w, args.hue))]
rows_b.append(np.hstack(rb))
rows_f.append(np.hstack(rf))
def stack(rows):
out = [rows[0]]
for r in rows[1:]:
out += [hsep(r.shape[1]), r]
return np.vstack(out)
grid_b = cv2.cvtColor(stack(rows_b), cv2.COLOR_RGB2BGR)
grid_f = cv2.cvtColor(stack(rows_f), cv2.COLOR_RGB2BGR)
cv2.imwrite("augmax_buggy.png", grid_b)
cv2.imwrite("augmax_fixed.png", grid_f)
print("Saved: augmax_buggy.png, augmax_fixed.png")
if __name__ == "__main__":
main()
In
src/openpi/models/model.pyI noticed a silent bug, which might produce inconsistent color semantics between base and wrist cameras, mainly due to hue component of augmax.ColorJitter transform and different RNG used within the same frame sample. Because of that, VLA has to rely more on spatial relation of objects, rather than on their colors (VLM sees corrupted colors). Please note that this problem doesn't impact wrist-to-wrist colors, because they are processed with the same ColorJitter + RNG op.I don't suspect it was an intended behavior, because it would probably be documented somewhere. Below please find a minimal script that visualizes the problem. While the potential fix seems to be trivial, it might impact the training dynamics, so please have a deeper look (e.g. do we want heavier color-reliance between base and wrists cameras? Or maybe this "bug" helps to regularize the model?)
Buggy transform (row: original, base, wrist1, wrist2)
Fixed transform (row: original, base, wrist1, wrist2)