Port objective-reward benchmark from claude/test-app-objective-rewards

Brings the objective-reward harness from branch
claude/test-app-objective-rewards-rjcquh (commit 3f3a7a2) onto master.
That branch forked before the Mojo 1.0 migration, so it was written in
pre-1.0 syntax (fn/alias/old imports, the removed Tensor type) and could
not compile or merge cleanly. Rather than a raw merge, the genuinely-new
value is reimplemented in 1.0 idioms:

- arc_io.mojo: add exact_match(), a SIMD discrete reward (fraction of
  cells correct after rounding) mirroring ARC's per-cell scoring.
- synth_tasks.py: deterministic task generator (flip/transpose/recolor/
  shift) reusing _save_grid so the .bin layout stays the single source
  of truth (pure Python, copied as-is).
- tests/test_es_convergence.mojo: asserts reward climbs toward a known
  target and reaches an exact match (auto-run via the test_*.mojo glob).
- benchmark.mojo: end-to-end harness reporting MSE reward, exact-match %,
  and aggregate solve rate over generated tasks.
- run_tests.sh: generate a small task set and run the benchmark.

Deliberately dropped: the branch's esper_evolution.mojo change (constant
epsilon -> Gaussian noise). Master's migration already replaced that with
real Gaussian noise *plus* antithetic mirrored sampling, which strictly
supersedes it. Convergence test now solves the target (exact-match 1.0)
and the benchmark solves 4/4 generated tasks.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: tacio

run_tests.sh

@@ -26,4 +26,18 @@ done
 echo "Running src/main.mojo (end-to-end driver)..."
 mojo run -I src src/main.mojo
 
+# Objective-reward benchmark: generate a small deterministic task set, then fit
+# each target with the ES loop and report the aggregate solve rate.
+echo "Running objective-reward benchmark (src/benchmark.mojo)..."
+BENCH_DIR="$(mktemp -d)"
+trap 'rm -rf "$BENCH_DIR"' EXIT
+python - "$BENCH_DIR" <<'PY'
+import sys
+sys.path.insert(0, "src")
+from synth_tasks import generate_tasks
+generate_tasks("flip_h", sys.argv[1], count=4, rows=4, cols=4, seed=0)
+print("Generated benchmark tasks in", sys.argv[1])
+PY
+mojo run -I src src/benchmark.mojo "$BENCH_DIR"/flip_h_*_out.bin
+
 echo "All tests passed successfully."

src/arc_io.mojo

@@ -1,5 +1,6 @@
 from std.sys import simd_width_of, size_of
 from std.memory import memcpy, UnsafePointer
+from std.math import round
 from hope import ArcGrid
 
 comptime nelts = simd_width_of[DType.float32]()
@@ -63,3 +64,32 @@ def calculate_fitness(
             mse_sum += diff * diff
 
     return -(mse_sum / Float32(size))
+
+
+def exact_match(
+    pred_ptr: UnsafePointer[Float32, MutAnyOrigin],
+    target_ptr: UnsafePointer[Float32, MutAnyOrigin],
+    size: Int,
+) -> Float32:
+    """Discrete objective reward: fraction of cells that match exactly.
+
+    ARC grids hold integer colors (0-9) while the engine evolves continuous
+    weights, so each prediction is rounded to the nearest integer before
+    comparison. This mirrors ARC's all-or-nothing per-cell scoring and
+    complements the continuous negative-MSE signal from `calculate_fitness`.
+    Same SIMD main-loop + scalar-remainder shape as the rest of the hot path.
+    """
+    var matches = Float32(0.0)
+
+    for i in range(0, size - nelts + 1, nelts):
+        var p_vec = round(pred_ptr.load[width=nelts](i))
+        var t_vec = round(target_ptr.load[width=nelts](i))
+        matches += p_vec.eq(t_vec).cast[DType.float32]().reduce_add()
+
+    var remainder = size % nelts
+    if remainder > 0:
+        for i in range(size - remainder, size):
+            if round(pred_ptr[i]) == round(target_ptr[i]):
+                matches += 1.0
+
+    return matches / Float32(size)

src/benchmark.mojo

@@ -0,0 +1,95 @@
+from std.sys import argv
+from std.memory import alloc, UnsafePointer
+from std.random import seed
+
+from arc_io import load_arc_grid, calculate_fitness, exact_match
+from esper_evolution import ESWorkspace, evolve_fast_weights
+
+# ==========================================================================
+# Esper objective-reward benchmark harness.
+#
+# Each argument is a path to a compiled target grid (`*_out.bin`, produced by
+# `src/synth_tasks.py`). For every task we evolve a fresh fast-weight buffer to
+# fit that target, then report two objective rewards:
+#   * reward  = negative MSE   (continuous; closer to 0 is better)
+#   * match%  = exact_match    (discrete; fraction of cells correct after round)
+# A task counts as solved when match% >= SOLVE_THRESHOLD. The aggregate solve
+# rate is the headline objective metric for the engine.
+#
+# Run from the project root, e.g.:
+#   python src/synth_tasks.py --transform flip_h --out data_bin --count 8
+#   mojo run -I src src/benchmark.mojo data_bin/flip_h_*_out.bin
+# ==========================================================================
+
+comptime SOLVE_THRESHOLD = Float32(0.99)
+comptime ES_SAMPLES = 96
+comptime ES_SIGMA = Float32(0.1)
+comptime ES_ITERS = 150
+
+
+def solve_task(target_path: String) raises -> Float32:
+    """Evolve fast weights to fit one target grid; return its exact-match fraction.
+    """
+    var target = load_arc_grid(target_path)
+    var size = target.size()
+    var tptr = target.data
+
+    # Fresh fast weights start at zero (no memory carried between tasks).
+    var fast = alloc[Float32](size)
+    for i in range(size):
+        fast[i] = 0.0
+
+    var workspace = ESWorkspace(size)
+
+    # Scale the learning rate with grid size so the effective step
+    # (~2 * alpha / size) stays in a stable range across grid dimensions.
+    var alpha = Float32(0.1) * Float32(size)
+
+    for _ in range(ES_ITERS):
+        evolve_fast_weights(fast, workspace, tptr, ES_SAMPLES, alpha, ES_SIGMA)
+
+    var reward = calculate_fitness(fast, tptr, size)
+    var match_frac = exact_match(fast, tptr, size)
+
+    print("  task:", target_path)
+    print("    reward (-MSE):", reward, " match%:", match_frac * 100.0)
+
+    fast.free()
+    return match_frac
+
+
+def main() raises:
+    seed(0)
+
+    var args = argv()
+    if len(args) < 2:
+        print(
+            "Usage: mojo run -I src src/benchmark.mojo <target_grid.bin>"
+            " [more.bin ...]"
+        )
+        print("Generate tasks first, e.g.:")
+        print(
+            "  python src/synth_tasks.py --transform flip_h --out data_bin"
+            " --count 8"
+        )
+        return
+
+    var total = len(args) - 1
+    var solved = 0
+
+    print("Esper objective-reward benchmark over", total, "task(s)")
+    for idx in range(1, len(args)):
+        var match_frac = solve_task(String(args[idx]))
+        if match_frac >= SOLVE_THRESHOLD:
+            solved += 1
+
+    var solve_rate = Float32(solved) / Float32(total) * 100.0
+    print("--------------------------------------------------")
+    print("Solved", solved, "/", total, " (solve rate:", solve_rate, "% )")
+
+    # Non-zero exit (via raised error) if nothing solved, so the harness/CI
+    # treats a fully-failing benchmark as a failure.
+    if solved == 0:
+        raise Error(
+            "ERROR: benchmark solved 0 tasks; the ES loop is not learning."
+        )

src/synth_tasks.py

@@ -0,0 +1,140 @@
+"""
+Synthetic, deterministic grid-transformation task generator for Esper.
+
+Esper targets ARC-AGI, where every task has a unique, machine-checkable
+ground-truth output grid -- i.e. an *objective reward* (exact pixel match).
+Real ARC tasks also need a full forward pass and a primitive library, which the
+engine does not have yet. This generator instead emits *single deterministic
+transforms* (recolor, flip, transpose, shift). Each task therefore has an
+unambiguous target grid, so the Evolution-Strategy loop can be scored objectively
+(negative MSE via `calculate_fitness`, exact-match % via `exact_match`) and we can
+prove the learning loop actually converges before tackling full ARC reasoning.
+
+Grids are written in the exact `.bin` layout the Mojo engine reads
+(`load_arc_grid` in `arc_io.mojo`): two little-endian int64 (rows, cols) followed
+by the flattened grid as float32. We reuse `_save_grid` from `arc_compiler.py`
+so the writer stays the single source of truth for that format.
+
+Usage:
+    python src/synth_tasks.py --transform flip_h --out data_bin --count 8
+    python src/synth_tasks.py --transform recolor --rows 6 --cols 6 --seed 0
+"""
+
+import argparse
+import os
+import random as _random
+
+# Reuse the canonical .bin writer so the on-disk format cannot drift from the
+# compiler / engine contract.
+from arc_compiler import _save_grid
+
+# ARC grids use integer "colors" 0-9.
+NUM_COLORS = 10
+
+
+# ---------------------------------------------------------------------------
+# Deterministic transforms. Each takes a 2-D list-of-lists grid and returns a
+# new grid of the same shape (so input and target share dimensions, which keeps
+# the engine's fixed-size fast-weight buffer valid).
+# ---------------------------------------------------------------------------
+def _identity(grid):
+    return [row[:] for row in grid]
+
+
+def _flip_h(grid):
+    return [list(reversed(row)) for row in grid]
+
+
+def _flip_v(grid):
+    return [row[:] for row in reversed(grid)]
+
+
+def _transpose(grid):
+    # Square grids only keep the same shape under transpose; the generator
+    # enforces rows == cols when this transform is selected.
+    return [list(col) for col in zip(*grid)]
+
+
+def _recolor(grid):
+    # Fixed cyclic palette shift: color c -> (c + 1) % NUM_COLORS.
+    return [[(c + 1) % NUM_COLORS for c in row] for row in grid]
+
+
+def _shift(grid):
+    # Roll every row right by one (wrap-around). Deterministic and shape-stable.
+    return [[row[-1]] + row[:-1] for row in grid]
+
+
+TRANSFORMS = {
+    "identity": _identity,
+    "flip_h": _flip_h,
+    "flip_v": _flip_v,
+    "transpose": _transpose,
+    "recolor": _recolor,
+    "shift": _shift,
+}
+
+
+def _random_grid(rows, cols, rng):
+    return [[rng.randrange(NUM_COLORS) for _ in range(cols)] for _ in range(rows)]
+
+
+def generate_tasks(transform, out_dir, count, rows, cols, seed):
+    """
+    Generate `count` (input, target) grid pairs for `transform` and write them as
+    `.bin` files into `out_dir`. Returns the list of (input_path, target_path)
+    tuples written.
+    """
+    if transform not in TRANSFORMS:
+        raise ValueError(
+            "Unknown transform %r; choose from %s"
+            % (transform, ", ".join(sorted(TRANSFORMS)))
+        )
+    if transform == "transpose" and rows != cols:
+        raise ValueError("transpose requires a square grid (rows == cols)")
+
+    os.makedirs(out_dir, exist_ok=True)
+    fn = TRANSFORMS[transform]
+    rng = _random.Random(seed)
+
+    pairs = []
+    for i in range(count):
+        grid_in = _random_grid(rows, cols, rng)
+        grid_out = fn(grid_in)
+
+        in_path = os.path.join(out_dir, "%s_%d_in.bin" % (transform, i))
+        out_path = os.path.join(out_dir, "%s_%d_out.bin" % (transform, i))
+        _save_grid(grid_in, in_path)
+        _save_grid(grid_out, out_path)
+        pairs.append((in_path, out_path))
+
+    return pairs
+
+
+def _parse_args():
+    p = argparse.ArgumentParser(description=__doc__)
+    p.add_argument(
+        "--transform",
+        default="flip_h",
+        choices=sorted(TRANSFORMS),
+        help="Which deterministic transform maps input -> target.",
+    )
+    p.add_argument("--out", default="data_bin", help="Output directory for .bin files.")
+    p.add_argument("--count", type=int, default=8, help="Number of task pairs to emit.")
+    p.add_argument("--rows", type=int, default=6, help="Grid rows.")
+    p.add_argument("--cols", type=int, default=6, help="Grid cols.")
+    p.add_argument("--seed", type=int, default=0, help="RNG seed (reproducible tasks).")
+    return p.parse_args()
+
+
+if __name__ == "__main__":
+    args = _parse_args()
+    written = generate_tasks(
+        args.transform, args.out, args.count, args.rows, args.cols, args.seed
+    )
+    print(
+        "Generated %d %s task pair(s) (%dx%d) in %s/"
+        % (len(written), args.transform, args.rows, args.cols, args.out)
+    )
+    for in_path, out_path in written:
+        print("  %s -> %s" % (os.path.basename(in_path), os.path.basename(out_path)))

tests/test_es_convergence.mojo

@@ -0,0 +1,87 @@
+from std.memory import alloc, UnsafePointer
+from std.random import seed
+
+# Run from the project root with `mojo run -I src tests/test_es_convergence.mojo`
+# so these cross-directory imports resolve.
+from esper_evolution import ESWorkspace, evolve_fast_weights
+from arc_io import calculate_fitness, exact_match
+
+# ==========================================================================
+# Objective-reward convergence test.
+#
+# Master's ES uses real Gaussian noise with antithetic sampling, so the fast
+# weights should march toward a known target grid: the reward (negative MSE)
+# must climb toward 0 and the rounded prediction must reach an exact match.
+# This is the smallest end-to-end check that the arena/IO/ES/SIMD/reward
+# plumbing actually learns (it would fail for the old constant-epsilon RNG,
+# whose gradient estimate was degenerate).
+# ==========================================================================
+
+
+def main() raises:
+    seed(0)
+
+    comptime size = 64
+
+    # Build an integer-valued target grid, like an ARC grid (colors 0-9).
+    var target = alloc[Float32](size)
+    for i in range(size):
+        target[i] = Float32(i % 10)
+
+    # Fast weights start at zero, i.e. far from the target.
+    var fast = alloc[Float32](size)
+    for i in range(size):
+        fast[i] = 0.0
+
+    var workspace = ESWorkspace(size)
+
+    var N = 96
+    var alpha = Float32(6.0)
+    var sigma = Float32(0.1)
+    var iters = 60
+
+    var initial_reward = calculate_fitness(fast, target, size)
+    var prev = initial_reward
+    var improvements = 0
+
+    for _ in range(iters):
+        evolve_fast_weights(fast, workspace, target, N, alpha, sigma)
+        var r = calculate_fitness(fast, target, size)
+        if r > prev:
+            improvements += 1
+        prev = r
+
+    var final_reward = calculate_fitness(fast, target, size)
+    var match_frac = exact_match(fast, target, size)
+
+    print("Initial reward (-MSE):", initial_reward)
+    print("Final   reward (-MSE):", final_reward)
+    print("Exact-match fraction :", match_frac)
+    print("Improving iterations :", improvements, "/", iters)
+
+    # 1. Reward must improve overall.
+    if final_reward <= initial_reward:
+        fast.free()
+        target.free()
+        raise Error("ERROR: ES did not improve reward toward the target.")
+
+    # 2. Improvement must be sustained, not a single lucky step. We require at
+    #    least a third of iterations to improve: a degenerate loop plateaus
+    #    almost immediately, while a working optimizer climbs steadily before
+    #    jittering near the optimum.
+    if improvements * 3 < iters:
+        fast.free()
+        target.free()
+        raise Error(
+            "ERROR: ES reward did not improve steadily across iterations."
+        )
+
+    # 3. The objective (discrete) reward must be (near) solved.
+    if match_frac < 0.95:
+        fast.free()
+        target.free()
+        raise Error("ERROR: ES did not converge to an exact grid match.")
+
+    fast.free()
+    target.free()
+    print("ES convergence test passed.")