Increase arc-gen fitting to 200 + add color shuffling augmentation for conv generalization

own-solver/neurogolf_solver/data_loader.py

@@ -56,8 +56,22 @@ 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)."""
+    """Get examples for fitting with ARC-GEN augmentation (fixed shape).
+    Uses up to 200 arc-gen examples + color-shuffled augmentations."""
     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:
@@ -72,11 +86,20 @@ 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))
-    return base_exs + ag_exs[:50]
+    # 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
 
 
 def get_exs_for_fitting_variable(td):
-    """Get examples for variable-shape fitting."""
+    """Get examples for variable-shape fitting.
+    Uses up to 200 arc-gen examples + color-shuffled augmentations."""
     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 = []
@@ -85,7 +108,16 @@ 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))
-    return base_exs + ag_exs[:50]
+    # 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
 
 
 def fixed_shapes(td):