1"""Bellman stopping controller MVP and mechanism checks.
  2
  3Run with: /home/maxwelhelp/main/bin/python3 bellman_stopping_experiment.py
  4"""
  5import json
  6import math
  7import random
  8from dataclasses import dataclass
  9from pathlib import Path
 10import numpy as np
 11
 12SEED = 2092
 13rng = np.random.default_rng(SEED)
 14
 15@dataclass
 16class ToyEnv:
 17    n: int = 21                 # score states 0,.05,...,1
 18    max_steps: int = 5
 19
 20    def p_improve(self, s):
 21        # Refinement has diminishing returns at high verifier score.
 22        return 0.10 + 0.75 * (1.0 - s)
 23
 24    def transition(self, i):
 25        s = i / (self.n - 1)
 26        p = self.p_improve(s)
 27        return [(i, 1-p), (min(i+1, self.n-1), p)]
 28
 29    def payoff(self, i):
 30        # Normalized terminal task utility / calibrated pass probability.
 31        return i / (self.n - 1)
 32
 33
 34def bellman(env, cost):
 35    """Exact finite-horizon DP. V[k,i] is value with k refinements left."""
 36    H, n = env.max_steps, env.n
 37    V = np.zeros((H + 1, n), dtype=float)
 38    policy = np.zeros((H, n), dtype=bool)  # True = continue
 39    V[0] = np.arange(n) / (n - 1)
 40    q_all = np.zeros_like(policy, dtype=float)
 41    for left in range(1, H + 1):
 42        for i in range(n):
 43            expected = sum(prob * V[left-1][j] for j, prob in env.transition(i))
 44            q = -cost + expected
 45            q_all[left-1, i] = q
 46            policy[left-1, i] = q > env.payoff(i) + 1e-12
 47            V[left, i] = max(env.payoff(i), q)
 48    return V, policy, q_all
 49
 50
 51def simulate(env, cost, adaptive=True, fixed_steps=None, trials=100000):
 52    V, policy, q = bellman(env, cost)
 53    n_calls, qualities, objectives, stops = [], [], [], []
 54    for _ in range(trials):
 55        i = 0
 56        calls = 0
 57        while calls < env.max_steps:
 58            if adaptive:
 59                # policy indexed by remaining refinements.
 60                if not policy[env.max_steps - calls - 1, i]:
 61                    break
 62            elif calls >= fixed_steps:
 63                break
 64            s = i / (env.n - 1)
 65            if rng.random() < env.p_improve(s):
 66                i = min(i + 1, env.n - 1)
 67            calls += 1
 68        quality = env.payoff(i)
 69        n_calls.append(calls)
 70        qualities.append(quality)
 71        objectives.append(quality - cost * calls)
 72        stops.append(i)
 73    return {
 74        "calls": float(np.mean(n_calls)),
 75        "quality": float(np.mean(qualities)),
 76        "objective": float(np.mean(objectives)),
 77        "pass_at_0.75": float(np.mean(np.asarray(qualities) >= .75)),
 78        "stop_state_mean": float(np.mean(stops)),
 79    }
 80
 81
 82def threshold(policy_row, env):
 83    # Return first score at which policy stops; score monotonicity can be checked.
 84    stop = np.flatnonzero(~policy_row)
 85    return float(stop[0] / (env.n - 1)) if len(stop) else None
 86
 87
 88def main():
 89    env = ToyEnv()
 90    # Prediction 1: continuation at the initial state decreases monotonically with c.
 91    costs = [0.00, 0.02, 0.03, 0.035, 0.038, 0.040, 0.045, 0.10]
 92    initial_continue = []
 93    initial_thresholds = []
 94    boundary_monotone = []
 95    for c in costs:
 96        V, pol, q = bellman(env, c)
 97        initial_continue.append(bool(pol[env.max_steps-1, 0]))
 98        # At the first decision, policy row has max remaining horizon.
 99        row = pol[env.max_steps-1]
100        initial_thresholds.append(threshold(row, env))
101        # Stop region should be an upper score interval.
102        stop_indices = np.flatnonzero(~row)
103        boundary_monotone.append(bool(len(stop_indices) == 0 or np.all(np.diff(stop_indices) == 1)))
104
105    # Prediction 2: the stopping boundary moves to lower score as cost rises.
106    # Ignore c=0's no-stop case when calculating finite boundaries.
107    finite_thresholds = [x for x in initial_thresholds if x is not None]
108    threshold_nonincreasing = all(a >= b for a, b in zip(finite_thresholds, finite_thresholds[1:]))
109
110    # Prediction 3: Bellman value is non-increasing in cost and equals max(stop, continue).
111    bellman_residuals, value_by_cost = [], []
112    for c in costs:
113        V, pol, q = bellman(env, c)
114        residual = np.max(np.abs(V[env.max_steps] - np.maximum(env.payoff(np.arange(env.n)), q[env.max_steps-1])))
115        bellman_residuals.append(float(residual))
116        value_by_cost.append(float(V[env.max_steps, 0]))
117    value_nonincreasing = all(a >= b - 1e-12 for a, b in zip(value_by_cost, value_by_cost[1:]))
118
119    # Monte Carlo comparison at moderate cost, with identical starting state.
120    c_eval = 0.038
121    adaptive = simulate(env, c_eval, adaptive=True)
122    fixed = {str(k): simulate(env, c_eval, adaptive=False, fixed_steps=k) for k in [1, 3, 5]}
123
124    # A direct numerical check of Delta > c versus policy on every state/horizon.
125    V, pol, q = bellman(env, c_eval)
126    decisions_agree = bool(np.all(pol == (q > np.tile(env.payoff(np.arange(env.n)), (env.max_steps, 1)))))
127
128    report = {
129        "seed": SEED,
130        "environment": {"score_states": env.n, "max_refinements": env.max_steps,
131                         "transition": "score increases one grid step with p=0.10+0.75*(1-score)",
132                         "payoff": "score"},
133        "math_check": {
134            "max_bellman_residual": max(bellman_residuals),
135            "delta_gt_cost_matches_policy": decisions_agree,
136            "value_at_initial_state_by_cost": dict(zip(map(str, costs), value_by_cost)),
137        },
138        "predictions": {
139            "P1_initial_continue_decreases_with_cost": {
140                "predicted": "monotone non-increasing",
141                "costs": costs,
142                "observed_continue": initial_continue,
143                "observed_threshold_score": initial_thresholds,
144                "confirmed": bool(all(a >= b for a, b in zip([int(x) for x in initial_continue], [int(x) for x in initial_continue][1:])))
145            },
146            "P2_stopping_boundary_is_upper_score_region_and_moves_lower_with_cost": {
147                "predicted": "stop for sufficiently high score; boundary non-increasing in cost",
148                "boundary_monotone_each_cost": boundary_monotone,
149                "finite_boundaries": finite_thresholds,
150                "confirmed": bool(all(boundary_monotone) and threshold_nonincreasing),
151            },
152            "P3_bellman_value_decreases_with_cost": {
153                "predicted": "non-increasing value and exact max recursion",
154                "values": value_by_cost,
155                "residuals": bellman_residuals,
156                "confirmed": bool(value_nonincreasing and max(bellman_residuals) < 1e-12),
157            },
158        },
159        "simulation_cost_0.038": {"adaptive": adaptive, "fixed": fixed},
160        "simulation_cost_sweep": {
161            str(c): {"adaptive": simulate(env, c, adaptive=True, trials=30000)}
162            for c in costs
163        },
164    }
165    report["worked"] = all(x["confirmed"] for x in report["predictions"].values()) and decisions_agree
166    Path("results.json").write_text(json.dumps(report, indent=2))
167    print(json.dumps(report, indent=2))
168
169if __name__ == "__main__":
170    main()