Revert to 50 arc-gen, remove color shuffling (200+shuffling caused regression 67→38)

own-solver/neurogolf_solver/data_loader.py

@@ -56,22 +56,8 @@ def get_exs(td):
             for ex in td['train'] + td['test']]
 
 
-def _color_shuffle(inp, out, rng):
-    """Apply a random color permutation to both input and output.
-    Color 0 (background) is always preserved."""
-    perm = np.arange(10, dtype=np.int64)
-    # Shuffle colors 1-9 (keep 0 fixed)
-    non_bg = perm[1:].copy()
-    rng.shuffle(non_bg)
-    perm[1:] = non_bg
-    inp_shuffled = perm[inp]
-    out_shuffled = perm[out]
-    return inp_shuffled, out_shuffled
-
-
 def get_exs_for_fitting(td):
-    """Get examples for fitting with ARC-GEN augmentation (fixed shape).
-    Uses up to 200 arc-gen examples + color-shuffled augmentations."""
+    """Get examples for fitting with ARC-GEN augmentation (fixed shape)."""
     base_exs = [(np.array(ex['input'], dtype=np.int64), np.array(ex['output'], dtype=np.int64))
                 for ex in td['train'] + td['test']]
     if not base_exs:
@@ -86,20 +72,11 @@ def get_exs_for_fitting(td):
         out = np.array(ex['output'], dtype=np.int64)
         if inp.shape == base_shape and out.shape == base_exs[0][1].shape:
             ag_exs.append((inp, out))
-    # Use up to 200 arc-gen examples
-    fitting_exs = base_exs + ag_exs[:200]
-    # Add color-shuffled augmentations of base examples (helps learn color-invariant rules)
-    rng = np.random.RandomState(42)
-    for inp, out in base_exs:
-        for _ in range(3):
-            inp_s, out_s = _color_shuffle(inp, out, rng)
-            fitting_exs.append((inp_s, out_s))
-    return fitting_exs
+    return base_exs + ag_exs[:50]
 
 
 def get_exs_for_fitting_variable(td):
-    """Get examples for variable-shape fitting.
-    Uses up to 200 arc-gen examples + color-shuffled augmentations."""
+    """Get examples for variable-shape fitting."""
     base_exs = [(np.array(ex['input'], dtype=np.int64), np.array(ex['output'], dtype=np.int64))
                 for ex in td['train'] + td['test']]
     ag_exs = []
@@ -108,16 +85,7 @@ def get_exs_for_fitting_variable(td):
         out = np.array(ex['output'], dtype=np.int64)
         if inp.shape == out.shape and inp.shape[0] <= 30 and inp.shape[1] <= 30:
             ag_exs.append((inp, out))
-    # Use up to 200 arc-gen examples
-    fitting_exs = base_exs + ag_exs[:200]
-    # Add color-shuffled augmentations of base examples
-    rng = np.random.RandomState(42)
-    for inp, out in base_exs:
-        if inp.shape == out.shape:
-            for _ in range(3):
-                inp_s, out_s = _color_shuffle(inp, out, rng)
-                fitting_exs.append((inp_s, out_s))
-    return fitting_exs
+    return base_exs + ag_exs[:50]
 
 
 def fixed_shapes(td):