Add TRM solver NN executor with 14 transform types and ONNX export

trm_solver/executor.py

@@ -0,0 +1,330 @@
+"""
+TRM Solver — NN Executor for ARC-AGI NeuroGolf
+
+Takes a parsed transform (from Kilo/DeepSeek) and executes it as a
+tiny neural network. Each transform is implemented as a minimal NN
+that can be exported to ONNX.
+
+Architecture:
+- Each transform is a PyTorch nn.Module with frozen weights
+- Weights encode the transform parameters (not learned — set directly)
+- ONNX export produces a tiny model per task
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from typing import Dict, List, Tuple, Optional
+from dataclasses import dataclass
+import json
+
+
+# ─── Data Structures ───────────────────────────────────────────
+
+@dataclass
+class TransformSpec:
+    """Parsed output from Kilo/DeepSeek."""
+    name: str
+    params: Dict
+    objects: List[Dict] = None
+
+    def __post_init__(self):
+        if self.objects is None:
+            self.objects = []
+
+
+# ─── Base NN Transform ─────────────────────────────────────────
+
+class BaseTransformNN(nn.Module):
+    """Base class for all transform NNs. Subclasses implement _forward_impl."""
+    
+    def __init__(self, spec: TransformSpec):
+        super().__init__()
+        self.spec = spec
+        self.max_size = 30  # ARC max grid size
+        
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: Input grid [B, 1, H, W] or [1, H, W] — values 0-9
+        Returns:
+            Output grid [B, 1, H_out, W_out] — values 0-9
+        """
+        if x.dim() == 3:
+            x = x.unsqueeze(0)
+        if x.dim() == 2:
+            x = x.unsqueeze(0).unsqueeze(0)
+        return self._forward_impl(x)
+    
+    def _forward_impl(self, x: torch.Tensor) -> torch.Tensor:
+        raise NotImplementedError
+    
+    def count_params(self) -> int:
+        return sum(p.numel() for p in self.parameters())
+
+
+# ─── Identity ──────────────────────────────────────────────────
+
+class IdentityNN(BaseTransformNN):
+    """Output equals input."""
+    def _forward_impl(self, x):
+        return x
+
+
+# ─── Color Map ─────────────────────────────────────────────────
+
+class ColorMapNN(BaseTransformNN):
+    """Per-pixel color remapping. 100 params."""
+    
+    def __init__(self, spec: TransformSpec):
+        super().__init__(spec)
+        lut = spec.params.get("color_map", list(range(10)))
+        self.lut = nn.Conv2d(10, 10, kernel_size=1, bias=False)
+        weight = torch.zeros(10, 10, 1, 1)
+        for i, j in enumerate(lut):
+            weight[j, i, 0, 0] = 1.0
+        self.lut.weight = nn.Parameter(weight, requires_grad=False)
+    
+    def _forward_impl(self, x):
+        B, _, H, W = x.shape
+        x_flat = x.long().squeeze(1).clamp(0, 9)
+        onehot = F.one_hot(x_flat, num_classes=10).permute(0, 3, 1, 2).float()
+        out = self.lut(onehot)
+        return out.argmax(dim=1, keepdim=True).float()
+
+
+# ─── Geometric ─────────────────────────────────────────────────
+
+class FlipNN(BaseTransformNN):
+    def _forward_impl(self, x):
+        direction = self.spec.params.get("direction", "horizontal")
+        dim = 3 if direction == "horizontal" else 2
+        return torch.flip(x, [dim])
+
+
+class TransposeNN(BaseTransformNN):
+    def _forward_impl(self, x):
+        return x.transpose(2, 3)
+
+
+class RotateNN(BaseTransformNN):
+    def _forward_impl(self, x):
+        k = self.spec.params.get("k", 1)
+        return torch.rot90(x, k, [2, 3])
+
+
+# ─── Upscale ───────────────────────────────────────────────────
+
+class UpscaleNN(BaseTransformNN):
+    """Nearest-neighbor upscaling. ~scale**2 params."""
+    
+    def __init__(self, spec: TransformSpec):
+        super().__init__(spec)
+        self.scale = spec.params.get("scale", 2)
+        th = spec.params.get("output_shape", [30, 30])[0]
+        tw = spec.params.get("output_shape", [30, 30])[1]
+        self.target_h = th
+        self.target_w = tw
+    
+    def _forward_impl(self, x):
+        x = x.repeat_interleave(self.scale, dim=2).repeat_interleave(self.scale, dim=3)
+        return x[:, :, :self.target_h, :self.target_w]
+
+
+# ─── Kronecker Self-Similar ────────────────────────────────────
+
+class KronSelfSimilarNN(BaseTransformNN):
+    """output = kron((input != 0), input). 0 learnable params."""
+    
+    def _forward_impl(self, x):
+        mask = (x != 0).float()
+        B, _, H_in, W_in = x.shape
+        inp_e = x.unsqueeze(2).unsqueeze(2)
+        mask_e = mask.unsqueeze(4).unsqueeze(4)
+        result = (mask_e * inp_e).float()
+        result = result.permute(0, 1, 2, 4, 3, 5).contiguous()
+        H_out, W_out = H_in * H_in, W_in * W_in
+        return result.view(B, 1, H_out, W_out)
+
+
+class TileRepeatNN(BaseTransformNN):
+    def _forward_impl(self, x):
+        hr = self.spec.params.get("h_repeat", 2)
+        wr = self.spec.params.get("w_repeat", 2)
+        return x.repeat(1, 1, hr, wr)
+
+
+# ─── Concat Patterns ───────────────────────────────────────────
+
+class ConcatPatternsNN(BaseTransformNN):
+    """Concatenate transformed copies horizontally/vertically."""
+    
+    def _forward_impl(self, x):
+        axis = self.spec.params.get("axis", "horizontal")
+        ops = self.spec.params.get("operations", ["identity", "identity"])
+        pieces = []
+        for op in ops:
+            if op == "flip_h": pieces.append(torch.flip(x, [3]))
+            elif op == "flip_v": pieces.append(torch.flip(x, [2]))
+            elif op == "transpose": pieces.append(x.transpose(2, 3))
+            elif op == "rot90": pieces.append(torch.rot90(x, 1, [2, 3]))
+            elif op == "rot180": pieces.append(torch.rot90(x, 2, [2, 3]))
+            elif op == "rot270": pieces.append(torch.rot90(x, 3, [2, 3]))
+            else: pieces.append(x)
+        dim = 3 if axis == "horizontal" else 2
+        return torch.cat(pieces, dim=dim)
+
+
+# ─── Position Color LUT ────────────────────────────────────────
+
+class PositionColorLUTNN(BaseTransformNN):
+    """Per-position color lookup. H*W params."""
+    
+    def __init__(self, spec: TransformSpec):
+        super().__init__(spec)
+        lut = spec.params.get("lut", {})
+        self.h_o = spec.params.get("output_shape", [30, 30])[0]
+        self.w_o = spec.params.get("output_shape", [30, 30])[1]
+        self.lut = nn.Parameter(torch.zeros(1, 1, self.h_o, self.w_o), requires_grad=False)
+        with torch.no_grad():
+            for k, v in lut.items():
+                h, w = map(int, k.split(","))
+                if h < self.h_o and w < self.w_o:
+                    self.lut[0, 0, h, w] = float(v)
+    
+    def _forward_impl(self, x):
+        B = x.shape[0]
+        out = self.lut.expand(B, -1, -1, -1)
+        mask = (x[:, :, :self.h_o, :self.w_o] != 0).float()
+        return mask * out
+
+
+# ─── Spatial Gather ────────────────────────────────────────────
+
+class SpatialGatherNN(BaseTransformNN):
+    """Rearrange pixels via gather map. H*W*2 params."""
+    
+    def __init__(self, spec: TransformSpec):
+        super().__init__(spec)
+        gmap = spec.params.get("gather_map", {})
+        self.h_o = spec.params.get("output_shape", [30, 30])[0]
+        self.w_o = spec.params.get("output_shape", [30, 30])[1]
+        self.gh = nn.Parameter(torch.zeros(self.h_o, self.w_o, dtype=torch.long), requires_grad=False)
+        self.gw = nn.Parameter(torch.zeros(self.h_o, self.w_o, dtype=torch.long), requires_grad=False)
+        with torch.no_grad():
+            for k, v in gmap.items():
+                h, w = map(int, k.split(","))
+                sh, sw = map(int, v.split(","))
+                if h < self.h_o and w < self.w_o:
+                    self.gh[h, w] = sh
+                    self.gw[h, w] = sw
+    
+    def _forward_impl(self, x):
+        B, C, Hi, Wi = x.shape
+        gh = self.gh.clamp(0, Hi - 1)
+        gw = self.gw.clamp(0, Wi - 1)
+        return x[:, :, gh, gw]
+
+
+# ─── One-Hot Convolution ───────────────────────────────────────
+
+class OneHotConvNN(BaseTransformNN):
+    """One-hot encode, convolve, argmax decode. K^2*100 params."""
+    
+    def __init__(self, spec: TransformSpec):
+        super().__init__(spec)
+        kh = spec.params.get("kernel_h", 3)
+        kw = spec.params.get("kernel_w", 3)
+        self.conv = nn.Conv2d(10, 10, kernel_size=(kh, kw), padding='same', bias=False)
+        if "weights" in spec.params:
+            w = torch.tensor(spec.params["weights"], dtype=torch.float32)
+            self.conv.weight = nn.Parameter(w.view(10, 10, kh, kw), requires_grad=False)
+    
+    def _forward_impl(self, x):
+        B, _, H, W = x.shape
+        onehot = F.one_hot(x.long().squeeze(1).clamp(0, 9), 10).permute(0, 3, 1, 2).float()
+        return self.conv(onehot).argmax(dim=1, keepdim=True).float()
+
+
+class OneHotLinearNN(BaseTransformNN):
+    """One-hot encode, linear, argmax. 100 params."""
+    
+    def __init__(self, spec: TransformSpec):
+        super().__init__(spec)
+        self.linear = nn.Linear(10, 10, bias=False)
+        if "weights" in spec.params:
+            self.linear.weight = nn.Parameter(
+                torch.tensor(spec.params["weights"], dtype=torch.float32), requires_grad=False)
+    
+    def _forward_impl(self, x):
+        onehot = F.one_hot(x.long().squeeze(1).clamp(0, 9), 10).float()
+        return self.linear(onehot).argmax(dim=-1).unsqueeze(1).float()
+
+
+# ─── Factory & Parser ──────────────────────────────────────────
+
+TRANSFORM_REGISTRY = {
+    "identity": IdentityNN,
+    "color_map": ColorMapNN,
+    "flip": FlipNN,
+    "transpose": TransposeNN,
+    "rotate": RotateNN,
+    "upscale": UpscaleNN,
+    "kron_self_similar": KronSelfSimilarNN,
+    "tile_repeat": TileRepeatNN,
+    "concat_patterns": ConcatPatternsNN,
+    "pos_color_lut": PositionColorLUTNN,
+    "spatial_gather": SpatialGatherNN,
+    "onehot_conv": OneHotConvNN,
+    "onehot_linear": OneHotLinearNN,
+}
+
+
+def create_transform_nn(spec: TransformSpec) -> BaseTransformNN:
+    cls = TRANSFORM_REGISTRY.get(spec.name)
+    if cls is None:
+        raise ValueError(f"Unknown transform: {spec.name}")
+    return cls(spec)
+
+
+def parse_kilo_output(md: str) -> TransformSpec:
+    """Parse Kilo markdown into TransformSpec."""
+    lines = md.strip().split('\n')
+    name, params, section = None, {}, None
+    for line in lines:
+        line = line.strip()
+        if line.startswith('## '):
+            section = line[3:].strip().lower()
+            continue
+        if section == 'transform' and line.startswith('name:'):
+            name = line.split(':', 1)[1].strip()
+        elif section == 'parameters' and line.startswith('- '):
+            kv = line[2:].split(':', 1)
+            if len(kv) == 2:
+                k, v = kv[0].strip(), kv[1].strip()
+                try:
+                    import ast
+                    params[k] = ast.literal_eval(v)
+                except (ValueError, SyntaxError):
+                    params[k] = v
+    if not name:
+        raise ValueError("No transform name in Kilo output")
+    return TransformSpec(name=name, params=params)
+
+
+# ─── ONNX Export ───────────────────────────────────────────────
+
+def export_to_onnx(model: BaseTransformNN, input_shape: Tuple[int, int],
+                   output_path: str, opset: int = 17):
+    model.eval()
+    H, W = input_shape
+    dummy = torch.zeros(1, 1, H, W)
+    torch.onnx.export(model, dummy, output_path,
+        input_names=["input"], output_names=["output"],
+        dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}},
+        opset_version=opset, do_constant_folding=True)
+    import os
+    kb, p = os.path.getsize(output_path) / 1024, model.count_params()
+    print(f"Exported {output_path}: {kb:.1f} KB, {p} params")
+    return output_path