Update app.py

app.py

@@ -1,11 +1,10 @@
 import json
 import math
-import time
 from dataclasses import dataclass, asdict
 from typing import Dict, List, Tuple, Optional
 
 import numpy as np
-from PIL import Image, ImageDraw, ImageFont
+from PIL import Image, ImageDraw
 
 import gradio as gr
 
@@ -18,7 +17,7 @@ import gradio as gr
 # - Branching timelines (fork from any previous step)
 # - Click-to-edit map tiles
 #
-# Minimal philosophy: explicit rules, no hidden weights, replayable.
+# Compatible with older Gradio versions by avoiding fn_kwargs in .click()
 # ============================================================
 
 # -----------------------------
@@ -53,7 +52,7 @@ TILE_NAMES = {
     TELE: "Teleporter",
 }
 
-# Palette (kept simple; inspectable)
+# Palette (simple + inspectable)
 SKY = np.array([14, 16, 26], dtype=np.uint8)
 FLOOR_NEAR = np.array([24, 26, 40], dtype=np.uint8)
 FLOOR_FAR = np.array([10, 11, 18], dtype=np.uint8)
@@ -70,8 +69,6 @@ AGENT_COLORS = {
 # Deterministic RNG helper
 # -----------------------------
 def rng_for(seed: int, step: int, stream: int = 0) -> np.random.Generator:
-    # Stable stream keyed by (seed, step, stream)
-    # Using PCG64 bitgen for reproducibility.
     mix = (seed * 1_000_003) ^ (step * 9_999_937) ^ (stream * 97_531)
     return np.random.default_rng(mix & 0xFFFFFFFFFFFFFFFF)
 
@@ -84,22 +81,22 @@ class Agent:
     x: int
     y: int
     ori: int  # 0..3
-    energy: int = 100  # mainly for prey, food, etc.
+    energy: int = 100
 
 @dataclass
 class WorldState:
     seed: int
     step: int
-    grid: List[List[int]]  # ints
+    grid: List[List[int]]
     agents: Dict[str, Agent]
-    controlled: str  # which agent receives manual control
-    pov: str         # which agent camera is showing
+    controlled: str
+    pov: str
     autorun: bool
     speed_hz: float
     overlay: bool
     event_log: List[str]
     caught: bool
-    branches: Dict[str, int]  # branch_name -> step_index in history
+    branches: Dict[str, int]
 
 @dataclass
 class Snapshot:
@@ -119,12 +116,12 @@ def default_grid() -> List[List[int]]:
         g[y][0] = WALL
         g[y][GRID_W - 1] = WALL
 
-    # Some interior structure
+    # Interior structure
     for x in range(4, 17):
         g[7][x] = WALL
-    g[7][10] = DOOR  # a door gap
+    g[7][10] = DOOR
 
-    # Toys
+    # Items
     g[3][4] = FOOD
     g[11][15] = FOOD
     g[4][14] = NOISE
@@ -155,23 +152,21 @@ def init_state(seed: int) -> WorldState:
     )
 
 # -----------------------------
-# Per-agent belief memory
+# Belief memory
 # -----------------------------
 def init_belief() -> Dict[str, np.ndarray]:
-    # -1 unknown, else tile id
     b = {}
     for name in ["Predator", "Prey", "Scout"]:
         b[name] = -1 * np.ones((GRID_H, GRID_W), dtype=np.int16)
     return b
 
 # -----------------------------
-# Utility: movement + collision
+# Movement + collision
 # -----------------------------
 def in_bounds(x: int, y: int) -> bool:
     return 0 <= x < GRID_W and 0 <= y < GRID_H
 
 def is_blocking(tile: int) -> bool:
-    # door is passable (for drama); wall blocks; tele is passable
     return tile == WALL
 
 def move_forward(state: WorldState, a: Agent) -> None:
@@ -181,17 +176,14 @@ def move_forward(state: WorldState, a: Agent) -> None:
         return
     if is_blocking(state.grid[ny][nx]):
         return
-    # Door toggle mechanic: if you step onto a door, it becomes empty (door opens)
     if state.grid[ny][nx] == DOOR:
         state.grid[ny][nx] = EMPTY
         state.event_log.append(f"t={state.step}: {a.name} opened a door.")
     a.x, a.y = nx, ny
 
-    # Teleporter: stepping onto TELE sends you to the other TELE (deterministically)
     if state.grid[ny][nx] == TELE:
         teles = [(x, y) for y in range(GRID_H) for x in range(GRID_W) if state.grid[y][x] == TELE]
         if len(teles) >= 2:
-            # choose destination as "the other tele" based on sorted list
             teles_sorted = sorted(teles)
             idx = teles_sorted.index((nx, ny))
             dest = teles_sorted[(idx + 1) % len(teles_sorted)]
@@ -205,10 +197,9 @@ def turn_right(a: Agent) -> None:
     a.ori = (a.ori + 1) % 4
 
 # -----------------------------
-# Perception: LOS + FOV on grid
+# LOS + FOV visibility
 # -----------------------------
 def los_clear(grid: List[List[int]], x0: int, y0: int, x1: int, y1: int) -> bool:
-    # Bresenham line-of-sight; walls block
     dx = abs(x1 - x0)
     dy = abs(y1 - y0)
     sx = 1 if x0 < x1 else -1
@@ -230,15 +221,12 @@ def los_clear(grid: List[List[int]], x0: int, y0: int, x1: int, y1: int) -> bool
             y += sy
 
 def within_fov(observer: Agent, tx: int, ty: int, fov_deg: float = 78.0) -> bool:
-    # vector from observer to target in observer's local frame
     dx = tx - observer.x
     dy = ty - observer.y
     if dx == 0 and dy == 0:
         return True
-    # absolute angle of target
     angle = math.degrees(math.atan2(dy, dx)) % 360
     facing = ORI_DEG[observer.ori]
-    # smallest signed difference
     diff = (angle - facing + 540) % 360 - 180
     return abs(diff) <= (fov_deg / 2)
 
@@ -248,29 +236,26 @@ def visible(observer: Agent, target: Agent, grid: List[List[int]]) -> bool:
 # -----------------------------
 # Raycast pseudo-3D render
 # -----------------------------
-def raycast_view(state: WorldState, observer: Agent, belief: Optional[np.ndarray] = None) -> np.ndarray:
-    # Returns RGB uint8 image
+def raycast_view(state: WorldState, observer: Agent) -> np.ndarray:
     img = np.zeros((VIEW_H, VIEW_W, 3), dtype=np.uint8)
     img[:, :] = SKY
 
-    # floor gradient
     for y in range(VIEW_H // 2, VIEW_H):
         t = (y - VIEW_H // 2) / (VIEW_H // 2 + 1e-6)
         col = (1 - t) * FLOOR_NEAR + t * FLOOR_FAR
         img[y, :] = col.astype(np.uint8)
 
-    # ray setup
     fov = math.radians(FOV_DEG)
     half_fov = fov / 2
+
     for rx in range(RAY_W):
-        # camera plane: [-1, 1]
         cam_x = (2 * rx / (RAY_W - 1)) - 1
         ray_ang = math.radians(ORI_DEG[observer.ori]) + cam_x * half_fov
 
-        # DDA-like stepping
         ox, oy = observer.x + 0.5, observer.y + 0.5
         sin_a = math.sin(ray_ang)
         cos_a = math.cos(ray_ang)
+
         depth = 0.0
         hit_side = 0
 
@@ -283,19 +268,15 @@ def raycast_view(state: WorldState, observer: Agent, belief: Optional[np.ndarray
 
             tile = state.grid[ty][tx]
             if tile == WALL:
-                # side shading based on ray direction
-                # crude: if abs(cos)>abs(sin) consider "vertical" else "horizontal"
                 hit_side = 1 if abs(cos_a) > abs(sin_a) else 0
                 break
             if tile == DOOR:
-                # door is semi-visible in world; render thinner by treating as a hit but less dark
                 hit_side = 2
                 break
 
-        # project wall slice
         if depth >= MAX_DEPTH:
             continue
-        # fish-eye correction
+
         depth *= math.cos(ray_ang - math.radians(ORI_DEG[observer.ori]))
         depth = max(depth, 0.001)
 
@@ -308,30 +289,24 @@ def raycast_view(state: WorldState, observer: Agent, belief: Optional[np.ndarray
         elif hit_side == 1:
             col = WALL_SIDE.copy()
         else:
-            # door slice
             col = np.array([180, 210, 255], dtype=np.uint8)
 
-        # slight depth dim
         dim = max(0.25, 1.0 - (depth / MAX_DEPTH))
         col = (col * dim).astype(np.uint8)
 
-        # draw slice (scaled to full VIEW_W)
         x0 = int(rx * (VIEW_W / RAY_W))
         x1 = int((rx + 1) * (VIEW_W / RAY_W))
         img[y0:y1, x0:x1] = col
 
-    # overlay: draw "billboards" for visible agents
     for other_name, other in state.agents.items():
         if other_name == observer.name:
             continue
         if visible(observer, other, state.grid):
-            # place billboard at its relative angle
             dx = other.x - observer.x
             dy = other.y - observer.y
             ang = (math.degrees(math.atan2(dy, dx)) % 360)
             facing = ORI_DEG[observer.ori]
             diff = (ang - facing + 540) % 360 - 180
-            # map diff to screen x
             sx = int((diff / (FOV_DEG / 2)) * (VIEW_W / 2) + (VIEW_W / 2))
             dist = math.sqrt(dx * dx + dy * dy)
             h = int((VIEW_H * 0.65) / max(dist, 0.75))
@@ -345,7 +320,6 @@ def raycast_view(state: WorldState, observer: Agent, belief: Optional[np.ndarray
             img[y0:y1, x0:x1] = np.array(col, dtype=np.uint8)
 
     if state.overlay:
-        # reticle
         cx, cy = VIEW_W // 2, VIEW_H // 2
         img[cy - 1:cy + 2, cx - 10:cx + 10] = np.array([120, 190, 255], dtype=np.uint8)
         img[cy - 10:cy + 10, cx - 1:cx + 2] = np.array([120, 190, 255], dtype=np.uint8)
@@ -353,7 +327,7 @@ def raycast_view(state: WorldState, observer: Agent, belief: Optional[np.ndarray
     return img
 
 # -----------------------------
-# Top-down map render (truth or belief)
+# Top-down render
 # -----------------------------
 def render_topdown(grid: np.ndarray, agents: Dict[str, Agent], title: str, show_agents: bool = True) -> Image.Image:
     w = grid.shape[1] * TILE
@@ -361,12 +335,11 @@ def render_topdown(grid: np.ndarray, agents: Dict[str, Agent], title: str, show_
     im = Image.new("RGB", (w, h + 28), (10, 12, 18))
     draw = ImageDraw.Draw(im)
 
-    # tiles
     for y in range(grid.shape[0]):
         for x in range(grid.shape[1]):
             t = int(grid[y, x])
             if t == -1:
-                col = (18, 20, 32)  # unknown
+                col = (18, 20, 32)
             elif t == EMPTY:
                 col = (26, 30, 44)
             elif t == WALL:
@@ -385,7 +358,6 @@ def render_topdown(grid: np.ndarray, agents: Dict[str, Agent], title: str, show_
             x0, y0 = x * TILE, y * TILE + 28
             draw.rectangle([x0, y0, x0 + TILE - 1, y0 + TILE - 1], fill=col)
 
-    # grid lines
     for x in range(grid.shape[1] + 1):
         xx = x * TILE
         draw.line([xx, 28, xx, h + 28], fill=(12, 14, 22))
@@ -393,7 +365,6 @@ def render_topdown(grid: np.ndarray, agents: Dict[str, Agent], title: str, show_
         yy = y * TILE + 28
         draw.line([0, yy, w, yy], fill=(12, 14, 22))
 
-    # agents
     if show_agents:
         for name, a in agents.items():
             cx = a.x * TILE + TILE // 2
@@ -401,23 +372,19 @@ def render_topdown(grid: np.ndarray, agents: Dict[str, Agent], title: str, show_
             col = AGENT_COLORS.get(name, (220, 220, 220))
             r = TILE // 3
             draw.ellipse([cx - r, cy - r, cx + r, cy + r], fill=col)
-            # heading tick
             dx, dy = DIRS[a.ori]
             draw.line([cx, cy, cx + dx * r, cy + dy * r], fill=(10, 10, 10), width=3)
 
-    # title bar
     draw.rectangle([0, 0, w, 28], fill=(14, 16, 26))
     draw.text((8, 6), title, fill=(230, 230, 240))
-
     return im
 
 # -----------------------------
-# Autonomy policies (explicit rules)
+# Policies (explicit + deterministic)
 # -----------------------------
 def predator_policy(state: WorldState, step: int) -> str:
     pred = state.agents["Predator"]
     prey = state.agents["Prey"]
-    # If prey visible, chase: turn toward prey then forward
     if visible(pred, prey, state.grid):
         dx = prey.x - pred.x
         dy = prey.y - pred.y
@@ -429,28 +396,24 @@ def predator_policy(state: WorldState, step: int) -> str:
         if diff > 10:
             return "R"
         return "F"
-    # else wander deterministically
     r = rng_for(state.seed, step, stream=1)
     return r.choice(["F", "L", "R", "F", "F"])
 
 def prey_policy(state: WorldState, step: int) -> str:
     prey = state.agents["Prey"]
     pred = state.agents["Predator"]
-    # If predator visible, flee: turn away then forward
     if visible(prey, pred, state.grid):
         dx = pred.x - prey.x
         dy = pred.y - prey.y
         ang = (math.degrees(math.atan2(dy, dx)) % 360)
         facing = ORI_DEG[prey.ori]
         diff = (ang - facing + 540) % 360 - 180
-        # want to face opposite direction: add 180
         diff_away = ((diff + 180) + 540) % 360 - 180
         if diff_away < -10:
             return "L"
         if diff_away > 10:
             return "R"
         return "F"
-    # else seek food if adjacent, else wander
     for turn in [0, -1, 1, 2]:
         ori = (prey.ori + turn) % 4
         dx, dy = DIRS[ori]
@@ -462,36 +425,32 @@ def prey_policy(state: WorldState, step: int) -> str:
                 return "L"
             if turn == 1:
                 return "R"
-            return "R"  # 180 via two rights across ticks; keep simple
+            return "R"
     r = rng_for(state.seed, step, stream=2)
     return r.choice(["F", "L", "R", "F"])
 
 def scout_policy(state: WorldState, step: int) -> str:
-    # Scout tries to keep line-of-sight on predator without colliding
     scout = state.agents["Scout"]
     pred = state.agents["Predator"]
     if los_clear(state.grid, scout.x, scout.y, pred.x, pred.y):
-        # orbit-ish: if too close, turn away; else meander
         dist = abs(scout.x - pred.x) + abs(scout.y - pred.y)
         if dist <= 3:
             return "R"
         r = rng_for(state.seed, step, stream=3)
         return r.choice(["F", "L", "R", "F"])
-    else:
-        # seek predator direction
-        dx = pred.x - scout.x
-        dy = pred.y - scout.y
-        ang = (math.degrees(math.atan2(dy, dx)) % 360)
-        facing = ORI_DEG[scout.ori]
-        diff = (ang - facing + 540) % 360 - 180
-        if diff < -10:
-            return "L"
-        if diff > 10:
-            return "R"
-        return "F"
+    dx = pred.x - scout.x
+    dy = pred.y - scout.y
+    ang = (math.degrees(math.atan2(dy, dx)) % 360)
+    facing = ORI_DEG[scout.ori]
+    diff = (ang - facing + 540) % 360 - 180
+    if diff < -10:
+        return "L"
+    if diff > 10:
+        return "R"
+    return "F"
 
 # -----------------------------
-# Step simulation
+# Simulation step
 # -----------------------------
 def apply_action(state: WorldState, agent_name: str, action: str) -> None:
     a = state.agents[agent_name]
@@ -521,13 +480,10 @@ def tick(state: WorldState, manual_action: Optional[str] = None) -> None:
     if state.caught:
         return
 
-    # Manual action applies to controlled agent first (if provided)
     if manual_action:
         apply_action(state, state.controlled, manual_action)
 
-    # Autonomy for the others (and for controlled if autorun)
     step = state.step
-    # Controlled agent: if autorun and no manual action this tick, autopilot it
     if state.autorun and not manual_action:
         if state.controlled == "Predator":
             act = predator_policy(state, step)
@@ -537,7 +493,6 @@ def tick(state: WorldState, manual_action: Optional[str] = None) -> None:
             act = scout_policy(state, step)
         apply_action(state, state.controlled, act)
 
-    # Non-controlled always run their policy each tick
     for name in ["Predator", "Prey", "Scout"]:
         if name == state.controlled:
             continue
@@ -554,9 +509,9 @@ def tick(state: WorldState, manual_action: Optional[str] = None) -> None:
     state.step += 1
 
 # -----------------------------
-# History + branching
+# History
 # -----------------------------
-MAX_HISTORY = 3000  # keeps rewind practical on Spaces
+MAX_HISTORY = 3000
 
 def snapshot_of(state: WorldState) -> Snapshot:
     return Snapshot(
@@ -573,17 +528,13 @@ def restore_into(state: WorldState, snap: Snapshot) -> None:
     for k, d in snap.agents.items():
         state.agents[k] = Agent(**d)
     state.caught = snap.caught
-    # preserve full log, but annotate jump
     state.event_log.append(f"Jumped to t={snap.step} (rewind).")
 
 # -----------------------------
 # Belief updates
 # -----------------------------
 def update_belief_for_agent(state: WorldState, belief: np.ndarray, agent: Agent) -> None:
-    # Reveal tiles in a cone up to MAX_DEPTH using simple ray sampling
-    # plus always reveal own tile
     belief[agent.y, agent.x] = state.grid[agent.y][agent.x]
-
     base = math.radians(ORI_DEG[agent.ori])
     half = math.radians(FOV_DEG / 2)
     rays = 33 if agent.name != "Scout" else 45
@@ -605,29 +556,24 @@ def update_belief_for_agent(state: WorldState, belief: np.ndarray, agent: Agent)
                 break
 
 # -----------------------------
-# UI orchestration
+# Views + UI helpers
 # -----------------------------
 def build_views(state: WorldState, beliefs: Dict[str, np.ndarray]) -> Tuple[np.ndarray, Image.Image, Image.Image, Image.Image, str, str]:
     pov_agent = state.agents[state.pov]
 
-    # Update beliefs each frame (deterministic, based on current truth)
     for name, a in state.agents.items():
         update_belief_for_agent(state, beliefs[name], a)
 
-    # POV raycast
     pov_img = raycast_view(state, pov_agent)
 
-    # Truth map
     truth_np = np.array(state.grid, dtype=np.int16)
     truth_img = render_topdown(truth_np, state.agents, f"Truth Map — t={state.step}  seed={state.seed}", show_agents=True)
 
-    # Belief maps (two most interesting: controlled + other)
     ctrl = state.controlled
     other = "Prey" if ctrl == "Predator" else "Predator"
     ctrl_img = render_topdown(beliefs[ctrl], state.agents, f"{ctrl} Belief (Fog-of-War)", show_agents=True)
     other_img = render_topdown(beliefs[other], state.agents, f"{other} Belief (Fog-of-War)", show_agents=True)
 
-    # Status + log
     pred = state.agents["Predator"]
     prey = state.agents["Prey"]
     scout = state.agents["Scout"]
@@ -644,7 +590,6 @@ def build_views(state: WorldState, beliefs: Dict[str, np.ndarray]) -> Tuple[np.n
     return pov_img, truth_img, ctrl_img, other_img, status, log
 
 def grid_click_to_tile(evt: gr.SelectData, selected_tile: int, state: WorldState) -> WorldState:
-    # evt.index is pixel coords (x,y) on truth image; our truth image has 28px title bar
     x_px, y_px = evt.index
     y_px = y_px - 28
     if y_px < 0:
@@ -653,11 +598,8 @@ def grid_click_to_tile(evt: gr.SelectData, selected_tile: int, state: WorldState
     gy = int(y_px // TILE)
     if not in_bounds(gx, gy):
         return state
-
-    # Protect borders from accidental deletion (optional)
     if gx == 0 or gy == 0 or gx == GRID_W - 1 or gy == GRID_H - 1:
         return state
-
     state.grid[gy][gx] = selected_tile
     state.event_log.append(f"t={state.step}: Edited tile ({gx},{gy}) -> {TILE_NAMES.get(selected_tile, selected_tile)}.")
     return state
@@ -686,18 +628,14 @@ def import_run(txt: str) -> Tuple[WorldState, List[Snapshot], Dict[str, np.ndarr
     st.overlay = bool(data.get("overlay", False))
     st.branches = dict(data.get("branches", {"main": 0}))
 
-    history = []
-    for s in data.get("history", []):
-        history.append(Snapshot(**s))
-
-    beliefs = init_belief()
-    rewind_idx = min(len(history) - 1, len(history) - 1 if history else 0)
-
-    if history:
-        restore_into(st, history[-1])
+    hist = [Snapshot(**s) for s in data.get("history", [])]
+    bel = init_belief()
 
+    r_idx = min(len(hist) - 1, len(hist) - 1 if hist else 0)
+    if hist:
+        restore_into(st, hist[-1])
     st.event_log.append("Imported run.")
-    return st, history, beliefs, rewind_idx
+    return st, hist, bel, r_idx
 
 # -----------------------------
 # Gradio app
@@ -705,13 +643,12 @@ def import_run(txt: str) -> Tuple[WorldState, List[Snapshot], Dict[str, np.ndarr
 with gr.Blocks(title="ChronoSandbox — Agent Timeline Lab") as demo:
     gr.Markdown(
         "## ChronoSandbox — Agent Timeline Lab\n"
-        "Deterministic multi-agent POV sandbox with **time dilation, rewind, and branching timelines**.\n"
-        "Everything is explicit: no hidden weights, no magic state."
+        "Deterministic multi-agent POV sandbox with **time dilation, rewind, and branching**.\n"
+        "Explicit rules, replayable runs."
     )
 
-    # Persistent state
     st = gr.State(init_state(seed=1337))
-    history = gr.State([snapshot_of(init_state(seed=1337))])  # start with step 0
+    history = gr.State([snapshot_of(init_state(seed=1337))])
     beliefs = gr.State(init_belief())
     rewind_index = gr.State(0)
 
@@ -758,12 +695,11 @@ with gr.Blocks(title="ChronoSandbox — Agent Timeline Lab") as demo:
             import_box = gr.Textbox(label="Import JSON", lines=10)
             btn_import = gr.Button("Import Run")
 
-    timer = gr.Timer(0.12)  # base UI refresh; actual tick rate controlled by speed_hz + autorun gating
+    timer = gr.Timer(0.12)
 
     def refresh(state: WorldState, hist: List[Snapshot], bel: Dict[str, np.ndarray], r_idx: int):
-        # clamp rewind slider max
         r_max = max(0, len(hist) - 1)
-        r_idx = max(0, min(r_idx, r_max))
+        r_idx = max(0, min(int(r_idx), r_max))
         pov_np, truth_im, a_im, b_im, stxt, ltxt = build_views(state, bel)
         return (
             pov_np,
@@ -815,8 +751,7 @@ with gr.Blocks(title="ChronoSandbox — Agent Timeline Lab") as demo:
     def jump_fn(state: WorldState, hist: List[Snapshot], bel: Dict[str, np.ndarray], r_idx: int, idx: int):
         if not hist:
             return refresh(state, hist, bel, r_idx) + (state, hist, bel, r_idx)
-        idx = int(idx)
-        idx = max(0, min(idx, len(hist) - 1))
+        idx = max(0, min(int(idx), len(hist) - 1))
         restore_into(state, hist[idx])
         r_idx = idx
         return refresh(state, hist, bel, r_idx) + (state, hist, bel, r_idx)
@@ -827,8 +762,7 @@ with gr.Blocks(title="ChronoSandbox — Agent Timeline Lab") as demo:
         state.event_log.append(f"t={state.step}: Branched timeline '{nm}' at history idx={r_idx}.")
         return refresh(state, hist, bel, r_idx) + (state, hist, bel, r_idx)
 
-    def truth_click(evt: gr.SelectData, tile: int, state: WorldState, hist: List[Snapshot], bel: Dict[str, np.ndarray], r_idx: int):
-        # apply edit, snapshot after edit
+    def truth_click(tile: int, state: WorldState, hist: List[Snapshot], bel: Dict[str, np.ndarray], r_idx: int, evt: gr.SelectData):
         state = grid_click_to_tile(evt, int(tile), state)
         hist.append(snapshot_of(state))
         if len(hist) > MAX_HISTORY:
@@ -841,7 +775,6 @@ with gr.Blocks(title="ChronoSandbox — Agent Timeline Lab") as demo:
 
     def import_fn(txt: str):
         state, hist, bel, r_idx = import_run(txt)
-        # refresh outputs + return states
         pov_np, truth_im, a_im, b_im, stxt, ltxt = build_views(state, bel)
         r_max = max(0, len(hist) - 1)
         return (
@@ -850,27 +783,95 @@ with gr.Blocks(title="ChronoSandbox — Agent Timeline Lab") as demo:
             state, hist, bel, r_idx
         )
 
-    # Buttons
-    btn_L.click(do_action, inputs=[st, history, beliefs, rewind_index], outputs=[pov_img, truth, belief_a, belief_b, status, log, rewind, rewind_index, st, history, beliefs, rewind_index], api_name=False, queue=True, fn_kwargs={"act": "L"})
-    btn_F.click(do_action, inputs=[st, history, beliefs, rewind_index], outputs=[pov_img, truth, belief_a, belief_b, status, log, rewind, rewind_index, st, history, beliefs, rewind_index], api_name=False, queue=True, fn_kwargs={"act": "F"})
-    btn_R.click(do_action, inputs=[st, history, beliefs, rewind_index], outputs=[pov_img, truth, belief_a, belief_b, status, log, rewind, rewind_index, st, history, beliefs, rewind_index], api_name=False, queue=True, fn_kwargs={"act": "R"})
+    # --- CLICK HANDLERS (NO fn_kwargs; use lambdas for compatibility) ---
+    btn_L.click(
+        lambda s, h, b, r: do_action(s, h, b, r, "L"),
+        inputs=[st, history, beliefs, rewind_index],
+        outputs=[pov_img, truth, belief_a, belief_b, status, log, rewind, rewind_index, st, history, beliefs, rewind_index],
+        api_name=False,
+        queue=True,
+    )
+    btn_F.click(
+        lambda s, h, b, r: do_action(s, h, b, r, "F"),
+        inputs=[st, history, beliefs, rewind_index],
+        outputs=[pov_img, truth, belief_a, belief_b, status, log, rewind, rewind_index, st, history, beliefs, rewind_index],
+        api_name=False,
+        queue=True,
+    )
+    btn_R.click(
+        lambda s, h, b, r: do_action(s, h, b, r, "R"),
+        inputs=[st, history, beliefs, rewind_index],
+        outputs=[pov_img, truth, belief_a, belief_b, status, log, rewind, rewind_index, st, history, beliefs, rewind_index],
+        api_name=False,
+        queue=True,
+    )
 
-    btn_step.click(do_tick, inputs=[st, history, beliefs, rewind_index], outputs=[pov_img, truth, belief_a, belief_b, status, log, rewind, rewind_index, st, history, beliefs, rewind_index], queue=True)
+    btn_step.click(
+        do_tick,
+        inputs=[st, history, beliefs, rewind_index],
+        outputs=[pov_img, truth, belief_a, belief_b, status, log, rewind, rewind_index, st, history, beliefs, rewind_index],
+        queue=True
+    )
 
-    toggle_control.click(toggle_control_fn, inputs=[st, history, beliefs, rewind_index], outputs=[pov_img, truth, belief_a, belief_b, status, log, rewind, rewind_index, st, history, beliefs, rewind_index], queue=True)
-    toggle_pov.click(toggle_pov_fn, inputs=[st, history, beliefs, rewind_index], outputs=[pov_img, truth, belief_a, belief_b, status, log, rewind, rewind_index, st, history, beliefs, rewind_index], queue=True)
+    toggle_control.click(
+        toggle_control_fn,
+        inputs=[st, history, beliefs, rewind_index],
+        outputs=[pov_img, truth, belief_a, belief_b, status, log, rewind, rewind_index, st, history, beliefs, rewind_index],
+        queue=True
+    )
+    toggle_pov.click(
+        toggle_pov_fn,
+        inputs=[st, history, beliefs, rewind_index],
+        outputs=[pov_img, truth, belief_a, belief_b, status, log, rewind, rewind_index, st, history, beliefs, rewind_index],
+        queue=True
+    )
 
-    autorun.change(set_toggles, inputs=[st, history, beliefs, rewind_index, autorun, speed, overlay], outputs=[pov_img, truth, belief_a, belief_b, status, log, rewind, rewind_index, st, history, beliefs, rewind_index], queue=True)
-    speed.change(set_toggles, inputs=[st, history, beliefs, rewind_index, autorun, speed, overlay], outputs=[pov_img, truth, belief_a, belief_b, status, log, rewind, rewind_index, st, history, beliefs, rewind_index], queue=True)
-    overlay.change(set_toggles, inputs=[st, history, beliefs, rewind_index, autorun, speed, overlay], outputs=[pov_img, truth, belief_a, belief_b, status, log, rewind, rewind_index, st, history, beliefs, rewind_index], queue=True)
+    autorun.change(
+        set_toggles,
+        inputs=[st, history, beliefs, rewind_index, autorun, speed, overlay],
+        outputs=[pov_img, truth, belief_a, belief_b, status, log, rewind, rewind_index, st, history, beliefs, rewind_index],
+        queue=True
+    )
+    speed.change(
+        set_toggles,
+        inputs=[st, history, beliefs, rewind_index, autorun, speed, overlay],
+        outputs=[pov_img, truth, belief_a, belief_b, status, log, rewind, rewind_index, st, history, beliefs, rewind_index],
+        queue=True
+    )
+    overlay.change(
+        set_toggles,
+        inputs=[st, history, beliefs, rewind_index, autorun, speed, overlay],
+        outputs=[pov_img, truth, belief_a, belief_b, status, log, rewind, rewind_index, st, history, beliefs, rewind_index],
+        queue=True
+    )
 
-    btn_jump.click(jump_fn, inputs=[st, history, beliefs, rewind_index, rewind], outputs=[pov_img, truth, belief_a, belief_b, status, log, rewind, rewind_index, st, history, beliefs, rewind_index], queue=True)
-    btn_branch.click(branch_fn, inputs=[st, history, beliefs, rewind_index, branch_name], outputs=[pov_img, truth, belief_a, belief_b, status, log, rewind, rewind_index, st, history, beliefs, rewind_index], queue=True)
+    btn_jump.click(
+        jump_fn,
+        inputs=[st, history, beliefs, rewind_index, rewind],
+        outputs=[pov_img, truth, belief_a, belief_b, status, log, rewind, rewind_index, st, history, beliefs, rewind_index],
+        queue=True
+    )
+    btn_branch.click(
+        branch_fn,
+        inputs=[st, history, beliefs, rewind_index, branch_name],
+        outputs=[pov_img, truth, belief_a, belief_b, status, log, rewind, rewind_index, st, history, beliefs, rewind_index],
+        queue=True
+    )
 
-    truth.select(truth_click, inputs=[tile_pick, st, history, beliefs, rewind_index], outputs=[pov_img, truth, belief_a, belief_b, status, log, rewind, rewind_index, st, history, beliefs, rewind_index], queue=True)
+    truth.select(
+        truth_click,
+        inputs=[tile_pick, st, history, beliefs, rewind_index],
+        outputs=[pov_img, truth, belief_a, belief_b, status, log, rewind, rewind_index, st, history, beliefs, rewind_index],
+        queue=True
+    )
 
     btn_export.click(export_fn, inputs=[st, history], outputs=[export_box], queue=True)
-    btn_import.click(import_fn, inputs=[import_box], outputs=[pov_img, truth, belief_a, belief_b, status, log, rewind, st, history, beliefs, rewind_index], queue=True)
+    btn_import.click(
+        import_fn,
+        inputs=[import_box],
+        outputs=[pov_img, truth, belief_a, belief_b, status, log, rewind, st, history, beliefs, rewind_index],
+        queue=True
+    )
 
     # Timer-driven autorun
     def timer_fn(state: WorldState, hist: List[Snapshot], bel: Dict[str, np.ndarray], r_idx: int, ar: bool, sp: float):
@@ -880,8 +881,6 @@ with gr.Blocks(title="ChronoSandbox — Agent Timeline Lab") as demo:
         if not state.autorun or state.caught:
             return refresh(state, hist, bel, r_idx) + (state, hist, bel, r_idx)
 
-        # How many sim ticks per UI frame?
-        # timer runs ~8.33 Hz (0.12s). We convert desired Hz to ticks per frame.
         ticks_per_frame = max(1, int(round(state.speed_hz * 0.12)))
         for _ in range(ticks_per_frame):
             tick(state, manual_action=None)
@@ -899,7 +898,11 @@ with gr.Blocks(title="ChronoSandbox — Agent Timeline Lab") as demo:
         queue=True
     )
 
-    # Initial paint
-    demo.load(refresh, inputs=[st, history, beliefs, rewind_index], outputs=[pov_img, truth, belief_a, belief_b, status, log, rewind, rewind_index], queue=True)
+    demo.load(
+        refresh,
+        inputs=[st, history, beliefs, rewind_index],
+        outputs=[pov_img, truth, belief_a, belief_b, status, log, rewind, rewind_index],
+        queue=True
+    )
 
 demo.queue().launch()