Update app.py

app.py

@@ -1,6 +1,10 @@
 """
-ARC-AGI-3 Live Puzzle Interface
+ARC-AGI-3 Agent Spectator
 Hugging Face Space: beanapologist/arc-agi
+
+Watch the Re/Im agent explore live ARC-AGI-3 games in real time.
+Enter your API key, pick a game, hit Watch — the agent plays and
+the grid updates every step.
 """
 
 import gradio as gr
@@ -9,15 +13,75 @@ import matplotlib
 matplotlib.use('Agg')
 import matplotlib.pyplot as plt
 from matplotlib.colors import ListedColormap
-import io, json, os
+import io, json, os, time, threading, queue
 from PIL import Image
 
+# ── Palette ───────────────────────────────────────────────────────────────────
+
 ARC_HEX = ['#000000','#1a6faf','#e03a3a','#3aa63a','#f5c400',
            '#c060c0','#d07030','#aaaaaa','#60b8d0','#874010']
 ARC_CMAP = ListedColormap(ARC_HEX)
 COLOR_NAMES = ['black','blue','red','green','yellow',
                'purple','orange','gray','azure','maroon']
 
+# ── Rendering ─────────────────────────────────────────────────────────────────
+
+def render_grid(grid, title='', highlight_diff=None):
+    if grid is None: return None
+    H, W = grid.shape
+    cell = max(28, min(60, 360 // max(H, W)))
+    fig, ax = plt.subplots(figsize=((W*cell+4)/72, (H*cell+22)/72), dpi=72)
+    fig.patch.set_facecolor('#1e1e2e'); ax.set_facecolor('#1e1e2e')
+    ax.imshow(grid, cmap=ARC_CMAP, vmin=0, vmax=9,
+              interpolation='nearest', aspect='equal')
+    for x in range(W+1): ax.axvline(x-.5, color='#444', lw=.5)
+    for y in range(H+1): ax.axhline(y-.5, color='#444', lw=.5)
+    for r in range(H):
+        for c in range(W):
+            v = int(grid[r, c])
+            col = 'white' if v in [0,1,2,3,5,6,9] else 'black'
+            ax.text(c, r, str(v), ha='center', va='center',
+                    fontsize=max(7, cell//5), color=col,
+                    fontweight='bold', fontfamily='monospace')
+            if highlight_diff is not None and highlight_diff[r, c]:
+                ax.add_patch(plt.Rectangle(
+                    (c-.5,r-.5), 1, 1, fill=True,
+                    facecolor='#ffffff', alpha=0.25, lw=0))
+    ax.set_xlim(-.5, W-.5); ax.set_ylim(H-.5, -.5); ax.axis('off')
+    if title:
+        ax.set_title(title, color='#cdd6f4', fontsize=10, pad=4)
+    plt.tight_layout(pad=.3)
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png', dpi=72, bbox_inches='tight',
+                facecolor='#1e1e2e')
+    buf.seek(0); img = Image.open(buf).copy(); plt.close()
+    return img
+
+def render_action_bar(action_counts, total):
+    """Horizontal bar chart of action frequency."""
+    if not action_counts or total == 0: return None
+    labels = [f"A{k}" for k in sorted(action_counts)]
+    vals   = [action_counts[k] for k in sorted(action_counts)]
+    fig, ax = plt.subplots(figsize=(5, 1.4))
+    fig.patch.set_facecolor('#1e1e2e'); ax.set_facecolor('#1e1e2e')
+    colors = ['#4a9eff','#e05050','#50c050','#f5c400','#c060c0','#d07030','#60b8d0']
+    bars = ax.barh(labels, vals, color=colors[:len(labels)], height=0.6)
+    for bar, v in zip(bars, vals):
+        ax.text(bar.get_width()+.3, bar.get_y()+bar.get_height()/2,
+                str(v), va='center', color='white', fontsize=8)
+    ax.set_xlim(0, max(vals)*1.25)
+    ax.tick_params(colors='#888', labelsize=8)
+    ax.spines[:].set_visible(False)
+    ax.set_facecolor('#1e1e2e')
+    plt.tight_layout(pad=.4)
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png', dpi=90, bbox_inches='tight',
+                facecolor='#1e1e2e')
+    buf.seek(0); img = Image.open(buf).copy(); plt.close()
+    return img
+
+# ── Minimal inline agent (no file import needed) ──────────────────────────────
+
 def _cc(mask):
     labels=np.zeros_like(mask,dtype=np.int32); cur=0; H,W=mask.shape
     for r in range(H):
@@ -57,340 +121,307 @@ def _boundary(grid):
     return ((p[1:-1,1:-1]!=p[:-2,1:-1])|(p[1:-1,1:-1]!=p[2:,1:-1])|
             (p[1:-1,1:-1]!=p[1:-1,:-2])|(p[1:-1,1:-1]!=p[1:-1,2:])).astype(np.float32)
 
-def compute_scores(grid):
-    h=float(_sym(grid,'h').max()); v=float(_sym(grid,'v').max())
-    b=int(_boundary(grid).sum())
-    gx,gy=_sobel(grid.astype(np.float32)/9)
-    edge=round(float(np.sqrt(gx**2+gy**2).mean()),4)
-    nc=sum(int(_cc(grid==c).max()) for c in range(10) if (grid==c).any())
-    return dict(h_sym=h,v_sym=v,boundary=b,edge=edge,components=nc)
-
-def render_grid(grid,title='',highlight_wrong=None):
-    H,W=grid.shape; cell=max(28,min(56,280//max(H,W)))
-    fig,ax=plt.subplots(figsize=((W*cell+4)/72,(H*cell+20)/72),dpi=72)
-    fig.patch.set_facecolor('#1e1e2e'); ax.set_facecolor('#1e1e2e')
-    ax.imshow(grid,cmap=ARC_CMAP,vmin=0,vmax=9,interpolation='nearest',aspect='equal')
-    for x in range(W+1): ax.axvline(x-.5,color='#555',lw=.5)
-    for y in range(H+1): ax.axhline(y-.5,color='#555',lw=.5)
-    for r in range(H):
-        for c in range(W):
-            v=int(grid[r,c])
-            col='white' if v in [0,1,2,3,5,6,9] else 'black'
-            ax.text(c,r,str(v),ha='center',va='center',
-                    fontsize=max(7,cell//5),color=col,fontweight='bold',fontfamily='monospace')
-            if highlight_wrong is not None and highlight_wrong[r,c]:
-                ax.add_patch(plt.Rectangle((c-.5,r-.5),1,1,fill=False,edgecolor='#ff4444',lw=2))
-    ax.set_xlim(-.5,W-.5); ax.set_ylim(H-.5,-.5); ax.axis('off')
-    if title: ax.set_title(title,color='#cdd6f4',fontsize=9,pad=3)
-    plt.tight_layout(pad=.3)
-    buf=io.BytesIO()
-    plt.savefig(buf,format='png',dpi=72,bbox_inches='tight',facecolor='#1e1e2e')
-    buf.seek(0); img=Image.open(buf).copy(); plt.close()
-    return img
+def extract_features_fast(grid, num_colours=10):
+    """Lightweight version of the 56-channel extractor for the Space demo."""
+    import torch, torch.nn.functional as F
+    H, W = grid.shape
+    one_hot = np.zeros((num_colours, H, W), dtype=np.float32)
+    for c in range(num_colours):
+        one_hot[c] = (grid==c).astype(np.float32)
+    gx, gy = _sobel(grid.astype(np.float32)/9)
+    h_sym  = _sym(grid,'h')[np.newaxis]
+    v_sym  = _sym(grid,'v')[np.newaxis]
+    bound  = _boundary(grid)[np.newaxis]
+    edge   = np.sqrt(gx**2+gy**2)[np.newaxis].astype(np.float32)
+    stacked = np.concatenate([one_hot, h_sym, v_sym, bound, edge], axis=0)
+    t = torch.from_numpy(stacked).float().unsqueeze(0)
+    if H != 64 or W != 64:
+        t = F.interpolate(t, size=(64,64), mode='bilinear', align_corners=False)
+    return t.squeeze(0)   # (14, 64, 64)
+
+class TinyAgent:
+    """Stripped-down CNN agent for the spectator demo."""
+    def __init__(self):
+        import torch, torch.nn as nn
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.model  = self._make_model().to(self.device)
+        self.opt    = torch.optim.Adam(self.model.parameters(), lr=1e-4)
+        self.buf    = []
+        self.prev_feat  = None
+        self.prev_action = None
+        self.step_count  = 0
+        self.action_counts = {}
+
+    def _make_model(self):
+        import torch.nn as nn
+        return nn.Sequential(
+            nn.Conv2d(14, 32, 3, padding=1), nn.ReLU(),
+            nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(),
+            nn.Conv2d(64, 128,3, padding=1), nn.ReLU(),
+            nn.AdaptiveAvgPool2d(8),
+            nn.Flatten(),
+            nn.Linear(128*8*8, 256), nn.ReLU(),
+            nn.Linear(256, 6),   # actions 1-6
+        )
+
+    def reset(self):
+        import torch, torch.nn as nn
+        self.model = self._make_model().to(self.device)
+        self.opt   = torch.optim.Adam(self.model.parameters(), lr=1e-4)
+        self.buf   = []; self.prev_feat = None; self.prev_action = None
+        self.action_counts = {}
+
+    def choose(self, grid, available_actions=None):
+        import torch, torch.nn.functional as F
+        feat = extract_features_fast(grid).to(self.device)
+
+        # Store experience
+        if self.prev_feat is not None:
+            changed = not np.array_equal(
+                self.prev_feat.cpu().numpy(),
+                feat.cpu().numpy())
+            self.buf.append((self.prev_feat, self.prev_action, 1.0 if changed else 0.0))
+            if len(self.buf) > 500: self.buf.pop(0)
+
+        # Train every 10 steps
+        if self.step_count % 10 == 0 and len(self.buf) >= 16:
+            self._train()
+
+        # Sample action
+        with torch.no_grad():
+            logits = self.model(feat.unsqueeze(0)).squeeze(0)
+            mask = list(range(1,7))
+            if available_actions:
+                mask = [int(a.value if hasattr(a,'value') else a) for a in available_actions
+                        if int(a.value if hasattr(a,'value') else a) <= 6]
+            indices = [m-1 for m in mask if 1 <= m <= 6]
+            masked = torch.full((6,), float('-inf'))
+            for i in indices: masked[i] = logits[i]
+            probs = torch.softmax(masked, dim=0).cpu().numpy()
+            probs = np.nan_to_num(probs, nan=1/len(indices))
+            if probs.sum() == 0: probs[indices] = 1/len(indices)
+            probs = probs / probs.sum()
+            action_idx = np.random.choice(6, p=probs)
+
+        self.prev_feat   = feat
+        self.prev_action = action_idx
+        self.step_count += 1
+        a_id = action_idx + 1
+        self.action_counts[a_id] = self.action_counts.get(a_id, 0) + 1
+
+        from arcengine import GameAction
+        return GameAction(a_id), dict(probs=probs.tolist())
+
+    def _train(self):
+        import torch, torch.nn.functional as F
+        import random
+        batch = random.sample(self.buf, min(16, len(self.buf)))
+        states  = torch.stack([b[0] for b in batch]).to(self.device)
+        actions = torch.tensor([b[1] for b in batch], dtype=torch.long, device=self.device)
+        rewards = torch.tensor([b[2] for b in batch], dtype=torch.float32, device=self.device)
+        self.opt.zero_grad()
+        logits = self.model(states)
+        loss = F.binary_cross_entropy_with_logits(
+            logits.gather(1, actions.unsqueeze(1)).squeeze(1), rewards)
+        loss.backward(); self.opt.step()
+
+
+# ── Session state ─────────────────────────────────────────────────────────────
+# One agent per Space instance — shared across users viewing the demo.
+_agent  = TinyAgent()
+_stop_flag  = threading.Event()
+_run_thread = None
+_frame_queue = queue.Queue(maxsize=30)
+
+def _run_agent(game_id, api_key, max_steps):
+    """Background thread: run agent, push frames to queue."""
+    import arc_agi
+    from arcengine import GameState
+    try:
+        arc = arc_agi.Arcade(arc_api_key=api_key)
+        env = arc.make(game_id, include_frame_data=True)
+        frame = env.reset()
+        _agent.reset()
+        prev_grid = None
+        step = 0
+
+        while not _stop_flag.is_set() and step < max_steps:
+            if frame is None: break
+            raw = np.array(frame.frame, dtype=np.int64)
+            grid = raw[-1] if raw.ndim == 3 else raw
+
+            avail = getattr(frame, 'available_actions', None)
+            action, info = _agent.choose(grid, avail)
+
+            diff = (grid != prev_grid) if prev_grid is not None else None
+            prev_grid = grid.copy()
+
+            state_str = str(getattr(frame, 'state', ''))
+            levels = getattr(frame, 'levels_completed', 0)
+
+            _frame_queue.put({
+                'grid':    grid,
+                'diff':    diff,
+                'step':    step,
+                'action':  int(action.value),
+                'levels':  levels,
+                'state':   state_str,
+                'probs':   info['probs'],
+                'counts':  dict(_agent.action_counts),
+            }, block=True, timeout=5)
+
+            if 'WIN' in state_str or 'GAME_OVER' in state_str:
+                break
+
+            frame = env.step(action)
+            step += 1
+            time.sleep(0.05)   # ~20 fps max
+
+        _frame_queue.put({'done': True, 'step': step})
+    except Exception as e:
+        _frame_queue.put({'error': str(e)})
 
-def render_feature_maps(grid):
-    gx,gy=_sobel(grid.astype(np.float32)/9)
-    maps=[('H-sym',_sym(grid,'h'),'YlOrRd'),('V-sym',_sym(grid,'v'),'YlOrRd'),
-          ('Boundary',_boundary(grid),'plasma'),('Edge',np.sqrt(gx**2+gy**2),'hot')]
-    fig,axes=plt.subplots(1,4,figsize=(12,2.8)); fig.patch.set_facecolor('#1e1e2e')
-    for ax,(t,d,cm) in zip(axes,maps):
-        ax.set_facecolor('#0d0d1a')
-        im=ax.imshow(d,cmap=cm,vmin=0,vmax=max(1,float(d.max())),
-                     interpolation='nearest',aspect='equal')
-        ax.set_title(t,color='white',fontsize=9,pad=3); ax.axis('off')
-        plt.colorbar(im,ax=ax,fraction=.046,pad=.04)
-    plt.tight_layout(pad=1)
-    buf=io.BytesIO()
-    plt.savefig(buf,format='png',dpi=100,bbox_inches='tight',facecolor='#1e1e2e')
-    buf.seek(0); img=Image.open(buf).copy(); plt.close()
-    return img
 
-def parse_grid(text):
-    try:
-        rows=[]
-        for line in text.strip().split('\n'):
-            line=line.strip()
-            if not line: continue
-            rows.append([int(x) for x in line.split()])
-        if not rows: return None,"Empty"
-        W=len(rows[0])
-        if any(len(r)!=W for r in rows): return None,"Unequal rows"
-        g=np.array(rows,dtype=np.int64)
-        if g.min()<0 or g.max()>9: return None,"Values 0-9 only"
-        return g,None
-    except Exception as e: return None,str(e)
-
-BUILTIN_PUZZLES = {
-    "Mirror complete": {
-        "desc": "Complete the right half — mirror the left half horizontally.",
-        "train_in":  "0 1 2 3 0 0 0 0\n0 4 0 3 0 0 0 0\n0 1 2 3 0 0 0 0",
-        "train_out": "0 1 2 3 3 2 1 0\n0 4 0 3 3 0 4 0\n0 1 2 3 3 2 1 0",
-        "test_in":   "0 2 1 0 0 0 0 0\n0 2 3 0 0 0 0 0\n0 0 1 0 0 0 0 0",
-        "answer":    "0 2 1 0 0 1 2 0\n0 2 3 0 0 3 2 0\n0 0 1 0 0 1 0 0",
-    },
-    "Color shift +1": {
-        "desc": "Each nonzero color increases by 1.",
-        "train_in":  "0 1 0\n2 0 3\n0 4 0",
-        "train_out": "0 2 0\n3 0 4\n0 5 0",
-        "test_in":   "0 3 1\n0 0 2\n4 0 0",
-        "answer":    "0 4 2\n0 0 3\n5 0 0",
-    },
-    "Border only": {
-        "desc": "Keep the outer border, fill interior with 0.",
-        "train_in":  "2 2 2 2 2\n2 2 2 2 2\n2 2 2 2 2\n2 2 2 2 2\n2 2 2 2 2",
-        "train_out": "2 2 2 2 2\n2 0 0 0 2\n2 0 0 0 2\n2 0 0 0 2\n2 2 2 2 2",
-        "test_in":   "3 3 3 3 3\n3 3 3 3 3\n3 3 3 3 3\n3 3 3 3 3\n3 3 3 3 3",
-        "answer":    "3 3 3 3 3\n3 0 0 0 3\n3 0 0 0 3\n3 0 0 0 3\n3 3 3 3 3",
-    },
-    "Gravity down": {
-        "desc": "Each colored pixel falls to the bottom of its column.",
-        "train_in":  "1 0 2\n0 0 0\n0 3 0",
-        "train_out": "0 0 0\n0 0 2\n1 3 0",
-        "test_in":   "4 0 0\n0 5 0\n0 0 6",
-        "answer":    "0 0 0\n0 0 0\n4 5 6",
-    },
-    "Count to block": {
-        "desc": "Count the 1s; output an NxN block of 2s at bottom-left.",
-        "train_in":  "0 1 0 1 0\n0 0 0 0 0\n0 0 0 0 0\n0 0 0 0 0\n0 0 0 0 0",
-        "train_out": "0 0 0 0 0\n0 0 0 0 0\n0 0 0 0 0\n2 2 0 0 0\n2 2 0 0 0",
-        "test_in":   "0 1 0 1 1\n0 0 0 0 0\n0 0 0 0 0\n0 0 0 0 0\n0 0 0 0 0",
-        "answer":    "0 0 0 0 0\n0 0 0 0 0\n2 2 2 0 0\n2 2 2 0 0\n2 2 2 0 0",
-    },
-}
-
-def make_grid_html(grid_list, selected_color=0):
-    H=len(grid_list); W=len(grid_list[0])
-    cell=min(52,max(24,300//max(H,W)))
-    swatches="".join(
-        f'<div onclick="selectColor({i})" title="{i}:{COLOR_NAMES[i]}" '
-        f'style="width:26px;height:26px;background:{ARC_HEX[i]};'
-        f'border:{"3px solid white" if i==selected_color else "2px solid #555"};'
-        f'border-radius:4px;cursor:pointer;display:flex;align-items:center;'
-        f'justify-content:center;color:{"white" if i in [0,1,2,3,5,6,9] else "black"};'
-        f'font-size:10px;font-weight:bold;">{i}</div>'
-        for i in range(10))
-    rows="".join(
-        "<tr>"+"".join(
-            f'<td onclick="paintCell({r},{c})" id="cell-{r}-{c}" '
-            f'style="width:{cell}px;height:{cell}px;background:{ARC_HEX[grid_list[r][c]]};'
-            f'color:{"white" if grid_list[r][c] in [0,1,2,3,5,6,9] else "black"};'
-            f'text-align:center;vertical-align:middle;font-size:{max(8,cell//5)}px;'
-            f'font-weight:bold;font-family:monospace;cursor:pointer;border:1px solid #444;'
-            f'user-select:none;">{grid_list[r][c]}</td>'
-            for c in range(W))+"</tr>"
-        for r in range(H))
-    return f"""
-<div style="font-family:sans-serif;background:#1e1e2e;padding:12px;border-radius:8px;">
-  <div style="font-size:11px;color:#888;margin-bottom:6px;">Pick color → click cells → Submit</div>
-  <div style="display:flex;gap:5px;margin-bottom:8px;flex-wrap:wrap;">{swatches}</div>
-  <table style="border-collapse:collapse;margin-bottom:8px;">{rows}</table>
-  <div style="display:flex;gap:8px;">
-    <button onclick="submitGrid()"
-      style="padding:7px 18px;background:#4a9eff;color:white;border:none;
-             border-radius:6px;cursor:pointer;font-size:13px;font-weight:bold;">Submit</button>
-    <button onclick="resetGrid()"
-      style="padding:7px 14px;background:#333;color:#ccc;border:1px solid #555;
-             border-radius:6px;cursor:pointer;font-size:13px;">Reset</button>
-  </div>
-</div>
-<script>
-var gridState={json.dumps(grid_list)};
-var selColor={selected_color};
-var HEX={json.dumps(ARC_HEX)};
-var DARK=[true,true,true,true,false,true,true,false,false,true];
-function selectColor(c){{
-  selColor=c;
-  document.querySelectorAll('[onclick^="selectColor"]').forEach(function(el,i){{
-    el.style.border=i===c?'3px solid white':'2px solid #555';
-  }});
-}}
-function paintCell(r,c){{
-  gridState[r][c]=selColor;
-  var td=document.getElementById('cell-'+r+'-'+c);
-  td.style.background=HEX[selColor];
-  td.style.color=DARK[selColor]?'white':'black';
-  td.textContent=selColor;
-}}
-function resetGrid(){{
-  for(var r=0;r<gridState.length;r++)
-    for(var c=0;c<gridState[r].length;c++){{
-      gridState[r][c]=0;
-      var td=document.getElementById('cell-'+r+'-'+c);
-      td.style.background=HEX[0];td.style.color='white';td.textContent='0';
-    }}
-}}
-function submitGrid(){{
-  var txt=gridState.map(function(row){{return row.join(' ');}}).join('\\n');
-  var ta=document.querySelector('#grid-state-box textarea');
-  if(ta){{ta.value=txt;ta.dispatchEvent(new Event('input',{{bubbles:true}}));}}
-}}
-</script>"""
-
-def get_arcade(api_key=""):
+# ── Gradio handlers ───────────────────────────────────────────────────────────
+
+def fetch_games(api_key):
     try:
         import arc_agi
-        key=api_key or os.environ.get("ARC_API_KEY","")
-        return arc_agi.Arcade(arc_api_key=key),None
-    except Exception as e: return None,str(e)
-
-def on_load_builtin(name):
-    p=BUILTIN_PUZZLES[name]
-    ti_img=render_grid(parse_grid(p["train_in"])[0],  title="Training input")
-    to_img=render_grid(parse_grid(p["train_out"])[0], title="Training output")
-    test_img=render_grid(parse_grid(p["test_in"])[0], title="Test input")
-    ans_g=parse_grid(p["answer"])[0]
-    blank=np.zeros_like(ans_g).tolist()
-    return (p["desc"], p["answer"], ti_img, to_img, test_img,
-            make_grid_html(blank), None, None, None,
-            "*Paint your answer and click Submit.*")
-
-def on_grid_submit(grid_txt, answer_hidden):
-    if not grid_txt or not grid_txt.strip():
-        return None,None,None,"*Paint something first.*"
-    g,err=parse_grid(grid_txt)
-    if err: return None,None,None,f"**Parse error:** {err}"
-    feat_img=render_feature_maps(g)
-    sc=compute_scores(g)
-    if answer_hidden and answer_hidden.strip():
-        target_g,terr=parse_grid(answer_hidden)
-        if terr is None and target_g.shape==g.shape:
-            correct=(g==target_g); pct=correct.mean()*100
-            solved=bool(correct.all()); wrong=~correct
-            answer_img=render_grid(g,title="Your answer",highlight_wrong=wrong)
-            target_img=render_grid(target_g,title="Target")
-            tsc=compute_scores(target_g)
-            emoji="✅" if solved else "❌"
-            return answer_img,target_img,feat_img,f"""### {emoji} {"Solved!" if solved else "Not yet!"}
-**Pixel accuracy: {pct:.1f}%** · {int(correct.sum())}/{g.size} cells correct
-
-| Feature | Yours | Target |
-|---|---|---|
-| H-symmetry | {sc['h_sym']:.2f} | {tsc['h_sym']:.2f} |
-| V-symmetry | {sc['v_sym']:.2f} | {tsc['v_sym']:.2f} |
-| Boundary px | {sc['boundary']} | {tsc['boundary']} |
-| Components | {sc['components']} | {tsc['components']} |
-"""
-    answer_img=render_grid(g,title="Your answer")
-    return answer_img,None,feat_img,f"""### Feature analysis
-| Feature | Value |
-|---|---|
-| H-symmetry | {sc['h_sym']:.2f} |
-| V-symmetry | {sc['v_sym']:.2f} |
-| Boundary pixels | {sc['boundary']} |
-| Components | {sc['components']} |
+        arc = arc_agi.Arcade(arc_api_key=api_key)
+        envs = arc.get_environments()
+        ids = [e.game_id for e in envs]
+        return gr.Dropdown(choices=ids, value=ids[0] if ids else None), \
+               f"Found **{len(ids)}** games."
+    except Exception as e:
+        return gr.Dropdown(choices=[]), f"**Error:** {e}"
+
+def start_agent(game_id, api_key, max_steps):
+    global _run_thread, _stop_flag
+    if not game_id:   return "Select a game first."
+    if not api_key:   return "Enter your API key first."
+    _stop_flag.set()
+    if _run_thread and _run_thread.is_alive():
+        _run_thread.join(timeout=3)
+    while not _frame_queue.empty():
+        try: _frame_queue.get_nowait()
+        except: break
+    _stop_flag.clear()
+    _run_thread = threading.Thread(
+        target=_run_agent,
+        args=(game_id, api_key, int(max_steps)),
+        daemon=True)
+    _run_thread.start()
+    return f"Agent started on **{game_id}** for {int(max_steps)} steps."
+
+def stop_agent():
+    _stop_flag.set()
+    return "Agent stopped."
+
+def stream_frames():
+    """Generator: yield (grid_img, bar_img, status_md) for each frame."""
+    while True:
+        try:
+            data = _frame_queue.get(timeout=1)
+        except queue.Empty:
+            yield None, None, "*Waiting for agent...*"
+            continue
+
+        if 'error' in data:
+            yield None, None, f"**Error:** {data['error']}"
+            return
+        if data.get('done'):
+            yield None, None, f"**Done** — {data['step']} steps completed."
+            return
+
+        grid  = data['grid']
+        diff  = data['diff']
+        step  = data['step']
+        action_id = data['action']
+        levels    = data['levels']
+        state_str = data['state']
+        counts    = data['counts']
+        probs     = data['probs']
+
+        grid_img = render_grid(grid,
+            title=f"Step {step} | Action {action_id} | Levels {levels}",
+            highlight_diff=diff)
+        bar_img  = render_action_bar(counts, sum(counts.values()))
+
+        action_names = {1:'A1',2:'A2',3:'A3',4:'A4',5:'A5',6:'A6(click)'}
+        prob_str = "  ".join(
+            f"**{action_names.get(i+1,str(i+1))}** {p:.2f}"
+            for i,p in enumerate(probs))
+
+        status = f"""**Step:** {step} &nbsp;|&nbsp; **Action:** {action_id} &nbsp;|&nbsp; **Levels:** {levels} &nbsp;|&nbsp; **State:** {state_str}
+
+Action probabilities: {prob_str}
 """
+        yield grid_img, bar_img, status
 
-def on_refresh_games(api_key):
-    arc,err=get_arcade(api_key)
-    if err: return gr.Dropdown(choices=[]),f"**Error:** {err}"
-    try:
-        envs=arc.get_environments()
-        ids=[e.game_id for e in envs]
-        return gr.Dropdown(choices=ids,value=ids[0] if ids else None),f"Found **{len(ids)}** games."
-    except Exception as e:
-        return gr.Dropdown(choices=[]),f"**Error:** {e}"
 
-def on_load_live(game_id,api_key):
-    if not game_id: return None,None,"Enter a game ID."
-    arc,err=get_arcade(api_key)
-    if err: return None,None,f"**API error:** {err}"
-    try:
-        env=arc.make(game_id,include_frame_data=True)
-        frame=env.reset()
-        raw=np.array(frame.frame,dtype=np.int64)
-        grid=raw[-1] if raw.ndim==3 else raw
-        blank=np.zeros_like(grid).tolist()
-        frame_img=render_grid(grid,title=f"{game_id} — current frame")
-        return frame_img,make_grid_html(blank),(
-            f"**Game:** `{game_id}` | Levels: {frame.levels_completed} | "
-            f"Grid: {grid.shape[0]}×{grid.shape[1]}")
-    except Exception as e:
-        return None,None,f"**Error:** {e}"
-
-with gr.Blocks(title="ARC-AGI-3 Puzzle Interface") as demo:
+# ── UI ────────────────────────────────────────────────────────────────────────
 
-    s_answer=gr.State("")
+with gr.Blocks(title="ARC-AGI-3 Agent Spectator") as demo:
 
     gr.Markdown("""
-# ARC-AGI-3 Puzzle Interface
-Paint cells to solve ARC puzzles. Use **Practice puzzles** to learn,
-or connect to the **Live ARC API** with your key from [docs.arcprize.org/api-keys](https://docs.arcprize.org/api-keys).
+# ARC-AGI-3 Agent Spectator
+Watch the Re/Im CNN agent explore live ARC-AGI-3 games in real time.
+Get your API key at [docs.arcprize.org/api-keys](https://docs.arcprize.org/api-keys).
 """)
 
-    with gr.Tabs():
-
-        with gr.Tab("Practice puzzles"):
-            with gr.Row():
-                with gr.Column(scale=2):
-                    builtin_dd=gr.Dropdown(choices=list(BUILTIN_PUZZLES.keys()),
-                                           value=list(BUILTIN_PUZZLES.keys())[0],
-                                           label="Select puzzle")
-                with gr.Column(scale=3):
-                    builtin_desc=gr.Markdown()
-
-            gr.Markdown("### Training example")
-            with gr.Row():
-                b_train_in =gr.Image(label="Input",  type="pil",interactive=False,height=160)
-                b_train_out=gr.Image(label="Output", type="pil",interactive=False,height=160)
-                b_test_in  =gr.Image(label="Test input — solve this",
-                                      type="pil",interactive=False,height=160)
-
-            gr.Markdown("### Your answer — click cells to paint")
-            b_html=gr.HTML()
-            b_grid_box=gr.Textbox(label="Grid state",elem_id="grid-state-box",
-                                   lines=2,placeholder="Paint above then click Submit")
-            gr.Markdown("### Score")
+    with gr.Row():
+        with gr.Column(scale=3):
+            api_box = gr.Textbox(
+                label="ARC API key",
+                type="password",
+                value=os.environ.get("ARC_API_KEY",""),
+                placeholder="arc-key-...  (or set ARC_API_KEY as HF Space secret)")
+        with gr.Column(scale=1):
+            fetch_btn = gr.Button("Fetch games")
+
+    with gr.Row():
+        with gr.Column(scale=2):
+            game_dd = gr.Dropdown(label="Game", choices=[])
+        with gr.Column(scale=1):
+            steps_sl = gr.Slider(label="Max steps", minimum=20,
+                                  maximum=500, value=100, step=10)
+        with gr.Column(scale=1):
             with gr.Row():
-                b_ans_img=gr.Image(label="Your answer",type="pil",interactive=False,height=180)
-                b_tgt_img=gr.Image(label="Target",     type="pil",interactive=False,height=180)
-            b_score=gr.Markdown("*Paint and click Submit to score.*")
-            b_feat =gr.Image(label="Re/Im feature maps",type="pil",interactive=False)
-
-            b_outs=[builtin_desc,s_answer,b_train_in,b_train_out,b_test_in,
-                    b_html,b_ans_img,b_tgt_img,b_feat,b_score]
-            builtin_dd.change(on_load_builtin,inputs=builtin_dd,outputs=b_outs)
-            b_grid_box.change(on_grid_submit,inputs=[b_grid_box,s_answer],
-                              outputs=[b_ans_img,b_tgt_img,b_feat,b_score])
-
-        with gr.Tab("Live ARC API"):
-            gr.Markdown("Enter your API key then click **Fetch games** to see available games.")
-            with gr.Row():
-                with gr.Column(scale=3):
-                    api_box=gr.Textbox(label="ARC API key",type="password",
-                                        value=os.environ.get("ARC_API_KEY",""),
-                                        placeholder="arc-key-...  (or set ARC_API_KEY as HF Secret)")
-                with gr.Column(scale=1):
-                    refresh_btn=gr.Button("Fetch games")
-            with gr.Row():
-                with gr.Column(scale=2):
-                    live_dd=gr.Dropdown(label="Game ID",choices=[])
-                with gr.Column(scale=1):
-                    load_btn=gr.Button("Load game",variant="primary")
-            api_status=gr.Markdown()
-            live_desc=gr.Markdown()
-
-            gr.Markdown("### Current frame")
-            l_frame=gr.Image(label="Live frame",type="pil",interactive=False,height=220)
-
-            gr.Markdown("### Paint your answer")
-            l_html=gr.HTML()
-            l_grid_box=gr.Textbox(label="Grid state",elem_id="grid-state-box",
-                                   lines=2,placeholder="Paint above then click Submit")
-            with gr.Row():
-                l_ans_img=gr.Image(label="Your answer",type="pil",interactive=False,height=180)
-                l_tgt_img=gr.Image(label="Target (if available)",
-                                    type="pil",interactive=False,height=180)
-            l_score=gr.Markdown()
-            l_feat =gr.Image(label="Re/Im feature maps",type="pil",interactive=False)
-
-            refresh_btn.click(on_refresh_games,inputs=api_box,
-                              outputs=[live_dd,api_status])
-            load_btn.click(on_load_live,inputs=[live_dd,api_box],
-                           outputs=[l_frame,l_html,live_desc])
-            l_grid_box.change(on_grid_submit,inputs=[l_grid_box,s_answer],
-                              outputs=[l_ans_img,l_tgt_img,l_feat,l_score])
-
-    demo.load(on_load_builtin,inputs=builtin_dd,outputs=b_outs)
+                start_btn = gr.Button("▶ Watch", variant="primary")
+                stop_btn  = gr.Button("■ Stop")
+
+    api_status = gr.Markdown()
+    run_status = gr.Markdown("*Press Fetch games, select a game, then Watch.*")
+
+    gr.Markdown("---")
+
+    with gr.Row():
+        grid_img = gr.Image(label="Current frame", type="pil",
+                             interactive=False, height=320)
+        bar_img  = gr.Image(label="Action frequency", type="pil",
+                             interactive=False, height=320)
+
+    stream_btn = gr.Button("⟳ Refresh frame", variant="secondary")
+
+    fetch_btn.click(fetch_games, inputs=api_box,
+                    outputs=[game_dd, api_status])
+    start_btn.click(start_agent, inputs=[game_dd, api_box, steps_sl],
+                    outputs=run_status)
+    stop_btn.click(stop_agent, outputs=run_status)
+    stream_btn.click(
+        lambda: next(stream_frames()),
+        outputs=[grid_img, bar_img, run_status])
+
+    gr.Markdown("""
+---
+**How it works:** The agent encodes each frame as 14 feature channels
+(10 one-hot colors + H-symmetry + V-symmetry + boundary contour + edge magnitude)
+and feeds them through a tiny CNN. It learns online: reward = 1 if the action
+changed the frame, 0 if not. The action frequency chart shows which actions
+the CNN is favouring as it learns.
+
+For the full 56-channel extractor used in the Kaggle submission, see the agent code.
+""")
 
 if __name__ == "__main__":
     demo.launch()