Merge pull request #1 from tylerbessire/codex/review-code-for-arc-2025-readiness

Harden submission runner with deterministic budgets

Author: tylerbessire

arc_solver/canonical.py

@@ -0,0 +1,99 @@
+"""Canonicalisation utilities for ARC grids.
+
+This module provides functions to normalise grids under the D4 symmetry group
+(rotations and reflections) and canonicalise colour labels. Canonicalisation
+reduces the search space by treating symmetric grids as identical.
+"""
+from __future__ import annotations
+
+from typing import Dict, Iterable, Tuple
+
+import numpy as np
+
+Array = np.ndarray
+
+# Precompute the eight transformations in the D4 symmetry group.
+D4: Tuple[callable, ...] = (
+    lambda g: g,
+    lambda g: np.rot90(g, 1),
+    lambda g: np.rot90(g, 2),
+    lambda g: np.rot90(g, 3),
+    lambda g: np.flipud(g),
+    lambda g: np.fliplr(g),
+    lambda g: np.rot90(np.flipud(g), 1),
+    lambda g: np.rot90(np.fliplr(g), 1),
+)
+
+
+def canonicalize_colors(grid: Array) -> Tuple[Array, Dict[int, int]]:
+    """Relabel colours in ``grid`` in descending frequency order.
+
+    Parameters
+    ----------
+    grid:
+        Input array containing integer colour labels.
+
+    Returns
+    -------
+    canonical:
+        Array with colours mapped to ``0..n-1`` in frequency order.
+    mapping:
+        Dictionary mapping original colours to canonical labels.
+
+    Raises
+    ------
+    TypeError
+        If ``grid`` is not a ``numpy.ndarray`` or is not of integer type.
+    """
+    if not isinstance(grid, np.ndarray):
+        raise TypeError("grid must be a numpy array")
+    if not np.issubdtype(grid.dtype, np.integer):
+        raise TypeError("grid dtype must be integer")
+
+    vals, counts = np.unique(grid, return_counts=True)
+    order = [int(v) for v, _ in sorted(zip(vals, counts), key=lambda t: (-t[1], t[0]))]
+    mapping = {c: i for i, c in enumerate(order)}
+    vect_map = np.vectorize(mapping.get)
+    canonical = vect_map(grid)
+    return canonical.astype(np.int16), mapping
+
+
+def canonicalize_D4(grid: Array) -> Array:
+    """Return the lexicographically smallest grid under D4 symmetries.
+
+    The grid is first transformed by each D4 element, then colour-canonicalised.
+    The transformation with the smallest shape and byte representation is chosen
+    as the canonical representative.
+
+    Parameters
+    ----------
+    grid:
+        Input array to canonicalise.
+
+    Returns
+    -------
+    np.ndarray
+        Canonicalised grid.
+
+    Raises
+    ------
+    TypeError
+        If ``grid`` is not a ``numpy.ndarray`` or is not of integer type.
+    """
+    if not isinstance(grid, np.ndarray):
+        raise TypeError("grid must be a numpy array")
+    if not np.issubdtype(grid.dtype, np.integer):
+        raise TypeError("grid dtype must be integer")
+
+    best: Array | None = None
+    best_key: Tuple[Tuple[int, int], bytes] | None = None
+    for transform in D4:
+        transformed = transform(grid)
+        canonical, _ = canonicalize_colors(transformed)
+        key = (canonical.shape, canonical.tobytes())
+        if best_key is None or key < best_key:
+            best, best_key = canonical, key
+    if best is None:
+        # This should not occur because D4 contains identity, but guard anyway.
+        return grid.copy()
+    return best

arc_solver/enhanced_search.py

@@ -236,20 +236,19 @@ def save_components(self):
         self.sketch_miner.save_sketches("sketches.json")
 
 
-def predict_two_enhanced(progs: List[List[Tuple[str, Dict[str, int]]]], 
-                        test_inputs: List[Array]) -> List[List[Array]]:
+def predict_two_enhanced(
+    progs: List[List[Tuple[str, Dict[str, int]]]],
+    test_inputs: List[Array],
+    prefer_diverse: bool = False,
+) -> List[List[Array]]:
     """Enhanced prediction with better fallback strategies."""
-    if len(progs) == 0:
-        # No programs found, use identity
+    if not progs:
         picks = [[("identity", {})], [("identity", {})]]
-    elif len(progs) == 1:
-        # Only one program, use it twice with slight variation if possible
-        main_prog = progs[0]
-        picks = [main_prog, main_prog]
+    elif prefer_diverse and len(progs) > 1:
+        picks = [progs[0], progs[1]]
     else:
-        # Use top 2 programs
-        picks = progs[:2]
-    
+        picks = progs[:2] if len(progs) >= 2 else [progs[0], progs[0]]
+
     attempts: List[List[Array]] = []
     for program in picks:
         outs: List[Array] = []
@@ -258,25 +257,26 @@ def predict_two_enhanced(progs: List[List[Tuple[str, Dict[str, int]]]],
                 result = apply_program(ti, program)
                 outs.append(result)
             except Exception:
-                # Fallback to identity on failure
                 outs.append(ti)
         attempts.append(outs)
-    
+
     return attempts
 
 
 # Integration function to use enhanced search in the main solver
-def synthesize_with_enhancements(train_pairs: List[Tuple[Array, Array]], 
-                               max_programs: int = 256) -> List[List[Tuple[str, Dict[str, int]]]]:
+def synthesize_with_enhancements(
+    train_pairs: List[Tuple[Array, Array]],
+    max_programs: int = 256,
+    force_alt: bool = False,
+) -> List[List[Tuple[str, Dict[str, int]]]]:
     """Main function to synthesize programs with all enhancements."""
-    
-    # Initialize enhanced search (this will be cached across calls in practice)
+
     enhanced_search = EnhancedSearch()
-    
-    # Try enhanced synthesis
     programs = enhanced_search.synthesize_enhanced(train_pairs, max_programs)
-    
-    # Save learned components periodically
+
+    if force_alt and len(programs) > 1:
+        programs = programs[1:]
+
     enhanced_search.save_components()
-    
+
     return programs

arc_solver/enhanced_solver.py

@@ -8,19 +8,22 @@
 
 from __future__ import annotations
 
-from typing import Dict, List, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 import numpy as np
 import os
 
 from .grid import to_array, to_list, Array
-from .search import synthesize, predict_two  # Keep original as fallback
+from .search import (
+    synthesize as synth_baseline,
+    predict_two as predict_two_baseline,
+)
 from .enhanced_search import synthesize_with_enhancements, predict_two_enhanced
 
 
 class ARCSolver:
     """Enhanced ARC solver with neural components and episodic memory."""
 
-    def __init__(self, use_enhancements: bool = True, 
+    def __init__(self, use_enhancements: bool = True,
                  guidance_model_path: str = None,
                  episode_db_path: str = "episodes.json"):
         self.use_enhancements = use_enhancements
@@ -32,6 +35,7 @@ def __init__(self, use_enhancements: bool = True,
             'enhancement_success_rate': 0.0,
             'fallback_used': 0,
         }
+        self._last_outputs: Optional[Tuple[List[List[List[int]]], List[List[List[int]]]]] = None
 
     def solve_task(self, task: Dict[str, List[Dict[str, List[List[int]]]]]) -> Dict[str, List[List[List[int]]]]:
         """Solve a single ARC task using enhanced or baseline methods."""
@@ -71,14 +75,77 @@ def solve_task(self, task: Dict[str, List[Dict[str, List[List[int]]]]]) -> Dict[
 
         except Exception:
             # Fall back to baseline approach
-            progs = synthesize(train_pairs)
-            attempts = predict_two(progs, test_inputs)
+            progs = synth_baseline(train_pairs)
+            attempts = predict_two_baseline(progs, test_inputs)
 
         # Convert outputs back to nested lists
         return {
             "attempt_1": [to_list(arr) for arr in attempts[0]],
             "attempt_2": [to_list(arr) for arr in attempts[1]],
         }
+
+    def solve_task_two_attempts(
+        self, task: Dict[str, List[Dict[str, List[List[int]]]]]
+    ) -> Tuple[List[List[List[int]]], List[List[List[int]]]]:
+        """Solve a task and ensure two diverse attempts.
+
+        Args:
+            task: ARC task specification.
+
+        Returns:
+            A tuple ``(attempt1, attempt2)`` each being a list of output grids
+            corresponding to the test inputs.
+        """
+
+        result = self.solve_task(task)
+        attempt1 = result["attempt_1"]
+        attempt2 = result["attempt_2"]
+
+        if attempt1 == attempt2:
+            alt = self._second_pass_diversified(task)
+            if alt is not None:
+                attempt2 = alt
+
+        self._last_outputs = (attempt1, attempt2)
+        return attempt1, attempt2
+
+    def _second_pass_diversified(
+        self, task: Dict[str, List[Dict[str, List[List[int]]]]]
+    ) -> Optional[List[List[List[int]]]]:
+        """Run a diversified second search pass to obtain an alternative output."""
+
+        train_pairs = [
+            (to_array(p["input"]), to_array(p["output"])) for p in task["train"]
+        ]
+        test_inputs = [to_array(p["input"]) for p in task["test"]]
+
+        try:
+            programs = synthesize_with_enhancements(train_pairs, force_alt=True)
+            attempts = predict_two_enhanced(programs, test_inputs, prefer_diverse=True)
+            return [to_list(x) for x in attempts[0]]
+        except Exception:
+            try:
+                programs = synth_baseline(train_pairs)
+                attempts = predict_two_baseline(
+                    programs, test_inputs, prefer_diverse=True
+                )
+                return [to_list(x) for x in attempts[0]]
+            except Exception:
+                return None
+
+    def best_so_far(
+        self, task: Dict[str, List[Dict[str, List[List[int]]]]]
+    ) -> List[List[List[int]]]:
+        """Return the best outputs computed so far for the current task.
+
+        If the solver has produced at least one attempt, that attempt is
+        returned. Otherwise, the identity transformation of the first test
+        input is used as a safe fallback.
+        """
+
+        if self._last_outputs is not None:
+            return self._last_outputs[0]
+        return [task["test"][0]["input"]]
 
     def _validate_solution(self, attempts: List[List[Array]], test_inputs: List[Array]) -> bool:
         """Basic validation to check if solution seems reasonable."""

arc_solver/features.py

@@ -13,6 +13,7 @@
 
 from .grid import Array, histogram, bg_color, eq
 from .objects import connected_components
+from .canonical import canonicalize_D4
 
 
 def extract_task_features(train_pairs: List[Tuple[Array, Array]]) -> Dict[str, Any]:
@@ -21,7 +22,15 @@ def extract_task_features(train_pairs: List[Tuple[Array, Array]]) -> Dict[str, A
     These features capture task-level properties that can help predict which
     DSL operations are likely to be relevant for solving the task.
     """
-    features = {}
+    try:
+        train_pairs = [
+            (canonicalize_D4(inp), canonicalize_D4(out))
+            for inp, out in train_pairs
+        ]
+    except TypeError as exc:
+        raise ValueError(f"invalid grid in train_pairs: {exc}") from exc
+
+    features: Dict[str, Any] = {}
 
     # Basic grid statistics
     input_shapes = [inp.shape for inp, _ in train_pairs]

arc_solver/objects.py

@@ -13,6 +13,7 @@
 from typing import List, Tuple, Dict, Any
 
 from .grid import Array, bg_color
+from .canonical import canonicalize_D4
 
 
 def neighbors4(y: int, x: int) -> List[Tuple[int, int]]:
@@ -21,14 +22,21 @@ def neighbors4(y: int, x: int) -> List[Tuple[int, int]]:
 
 
 def connected_components(a: Array) -> List[Dict[str, Any]]:
-    """Find all 4-connected components in the grid, grouping by exact color.
+    """Find all 4-connected components in a canonicalised grid.
 
-    Each component dictionary contains:
+    The input grid is first normalised under D4 symmetries and colour
+    relabelling to ensure deterministic component extraction. Each component
+    dictionary contains:
       - color: the color value of the component
       - bbox: (top, left, height, width) of the bounding box
       - mask: a 2D array of shape (height, width) with the component values
       - pixels: list of (row, col) indices in original grid
     """
+    try:
+        a = canonicalize_D4(a)
+    except TypeError as exc:
+        raise ValueError(f"invalid grid: {exc}") from exc
+
     h, w = a.shape
     visited = np.zeros_like(a, dtype=bool)
     comps: List[Dict[str, Any]] = []
@@ -65,11 +73,13 @@ def connected_components(a: Array) -> List[Dict[str, Any]]:
 
 
 def infer_symmetries(a: Array) -> Dict[str, bool]:
-    """Return a simple dictionary of possible symmetries in the grid.
+    """Return a dictionary of potential symmetries for a canonicalised grid."""
+    try:
+        a = canonicalize_D4(a)
+    except TypeError as exc:
+        raise ValueError(f"invalid grid: {exc}") from exc
 
-    For speed, this function does not check each symmetry but sets flags to True.
-    More precise symmetry detection can be incorporated later based on heuristics.
-    """
+    # Placeholder flags; symmetry detection can be refined with heuristics.
     return {
         "rot90": True,
         "rot180": True,