from __future__ import annotations import math import time import numpy as np from td_repro.proof_replay import replay_theorem_3_4 N_STATES = 512 N_FEATURES = 64 N_SEEDS = 64 HORIZONS = (5_000, 20_000, 80_000, 250_000) GAMMA = 0.5 RHO = 0.6 NOISE = 0.2 ETA0 = (1.0 - GAMMA) / 8.0 def make_problem(scale_floor: float) -> dict: rng = np.random.default_rng(260302577) q, _ = np.linalg.qr(rng.normal(size=(N_STATES, N_FEATURES))) scales = np.linspace(1.0, scale_floor, N_FEATURES) features = q * scales features /= np.linalg.norm(features, axis=1).max() reward_direction = rng.normal(size=N_FEATURES) raw_reward = features @ reward_direction rewards = 0.2 * raw_reward / np.abs(raw_reward).max() stationary_feature = features.mean(axis=0) expected_next = RHO * features + (1.0 - RHO) * stationary_feature covariance = features.T @ features / N_STATES td_matrix = features.T @ (features - GAMMA * expected_next) / N_STATES td_rhs = features.T @ rewards / N_STATES optimum = np.linalg.solve(td_matrix, td_rhs) omega = float(np.linalg.eigvalsh(covariance)[0]) return { "features": features, "rewards": rewards, "optimum": optimum, "omega": omega, "max_feature_norm": float(np.linalg.norm(features, axis=1).max()), "rank": int(np.linalg.matrix_rank(features)), "stationary_residual": 0.0, "max_abs_reward": float(np.abs(rewards).max() + NOISE), } def simulate(problem: dict, constant_step: bool = False) -> list[dict]: features = problem["features"] rewards = problem["rewards"] optimum = problem["optimum"] snapshots = [] for horizon in HORIZONS: rng = np.random.default_rng((314159 if constant_step else 271828) + horizon) weights = np.zeros((N_SEEDS, N_FEATURES)) alpha = horizon ** (-1.0 / horizon) for t in range(1, horizon + 1): states = rng.integers(N_STATES, size=N_SEEDS) next_states = states.copy() refresh = rng.random(N_SEEDS) >= RHO next_states[refresh] = rng.integers(N_STATES, size=int(refresh.sum())) reward_noise = np.where(rng.random(N_SEEDS) < 0.5, -NOISE, NOISE) phi = features[states] phi_next = features[next_states] prediction = np.einsum("ij,ij->i", phi, weights) next_prediction = np.einsum("ij,ij->i", phi_next, weights) td_error = rewards[states] + reward_noise + GAMMA * next_prediction - prediction eta = ETA0 if constant_step else ETA0 * alpha**t weights += eta * td_error[:, None] * phi squared_errors = np.square(weights - optimum).sum(axis=1) snapshots.append( { "T": horizon, "mean_squared_parameter_error": float(squared_errors.mean()), "ci95_half_width": float(1.96 * squared_errors.std(ddof=1) / math.sqrt(N_SEEDS)), "median_squared_parameter_error": float(np.median(squared_errors)), "rate_normalized_error": float(squared_errors.mean() * horizon / math.log(horizon) ** 2), } ) return snapshots def exact_sigma_squared(problem: dict) -> float: features = problem["features"] rewards = problem["rewards"] optimum = problem["optimum"] values = features @ optimum same_delta = rewards + (GAMMA - 1.0) * values refresh_delta = rewards[:, None] - values[:, None] + GAMMA * values[None, :] expected_delta_squared = ( RHO * np.square(same_delta) + (1.0 - RHO) * np.square(refresh_delta).mean(axis=1) + NOISE**2 ) return float(np.mean(np.square(features).sum(axis=1) * expected_delta_squared)) def theorem_bound(problem: dict, T: int) -> float: alpha = T ** (-1.0 / T) omega = problem["omega"] initial_error = float(np.square(problem["optimum"]).sum()) sigma_squared = exact_sigma_squared(problem) bias = initial_error * math.e * math.exp( -ETA0 * omega * (1.0 - GAMMA) * alpha * T / math.log(T) ) variance = ( 8.0 * sigma_squared / (math.e * (omega * (1.0 - GAMMA)) ** 2) * math.log(T) ** 2 / (alpha**2 * T) ) return bias + variance def scalar_independent_checker() -> dict: seeds = 20_000 horizon = 100_000 gamma = GAMMA eta0 = ETA0 alpha = horizon ** (-1.0 / horizon) mean_reward = 0.2 noise = 0.2 optimum = mean_reward / (1.0 - gamma) mean_error = -optimum second_moment = optimum**2 weights = np.zeros(seeds) rng = np.random.default_rng(1618033) for t in range(1, horizon + 1): eta = eta0 * alpha**t contraction = 1.0 - eta * (1.0 - gamma) mean_error *= contraction second_moment = contraction**2 * second_moment + eta**2 * noise**2 sample_noise = np.where(rng.random(seeds) < 0.5, -noise, noise) weights += eta * (mean_reward + sample_noise - (1.0 - gamma) * weights) errors = np.square(weights - optimum) observed = float(errors.mean()) standard_error = float(errors.std(ddof=1) / math.sqrt(seeds)) return { "method": "independent closed second-moment recursion for scalar TD compared with 20,000 trajectories", "horizon": horizon, "analytic_mse": second_moment, "monte_carlo_mse": observed, "monte_carlo_standard_error": standard_error, "absolute_z_score": abs(observed - second_moment) / standard_error, "passed": abs(observed - second_moment) <= 5.0 * standard_error, } def log_slope(rows: list[dict]) -> float: x = np.log([row["T"] for row in rows]) y = np.log([row["mean_squared_parameter_error"] for row in rows]) return float(np.polyfit(x, y, 1)[0]) def verify_claim1() -> dict: started = time.perf_counter() problems = [] for scale_floor in (1.0, 0.8, 0.6): problem = make_problem(scale_floor) rows = simulate(problem) for row in rows: row["theorem_upper_bound"] = theorem_bound(problem, row["T"]) row["bound_satisfied"] = ( row["mean_squared_parameter_error"] + row["ci95_half_width"] <= row["theorem_upper_bound"] ) problems.append( { "scale_floor": scale_floor, "omega": problem["omega"], "assumptions": { "n_states": N_STATES, "n_features": N_FEATURES, "feature_rank": problem["rank"], "max_feature_norm": problem["max_feature_norm"], "max_abs_reward": problem["max_abs_reward"], "irreducible": True, "aperiodic": True, "stationary_residual": problem["stationary_residual"], }, "rows": rows, "log_log_slope": log_slope(rows), "all_theorem_bounds_satisfied": all(row["bound_satisfied"] for row in rows), } ) control_problem = make_problem(1.0) control_rows = simulate(control_problem, constant_step=True) control_slope = log_slope(control_rows) independent = scalar_independent_checker() proof_replay = replay_theorem_3_4() assumptions_pass = all( item["assumptions"]["feature_rank"] == N_FEATURES and item["assumptions"]["max_feature_norm"] <= 1.0 + 1e-12 and item["assumptions"]["max_abs_reward"] <= 1.0 for item in problems ) primary_pass = all( item["all_theorem_bounds_satisfied"] and item["rows"][-1]["mean_squared_parameter_error"] < item["rows"][0]["mean_squared_parameter_error"] for item in problems ) negative_control_pass = control_slope > -0.25 passed = assumptions_pass and primary_pass and independent["passed"] and proof_replay["passed"] and negative_control_pass return { "claim_id": "claim1", "verdict": "VERIFIED" if passed else "BLOCKED", "passed": passed, "contract": "Theorem 3.4 exact last-iterate schedule, eta0=(1-gamma)/8 and alpha=T^(-1/T), is omega-free and its explicit finite-T bound covers observed iid TD(0) error on the predeclared full-dimensional family.", "scope": "computational theorem reproduction over three 512-state/64-feature MDPs; not a universal formal proof", "schedule": {"gamma": GAMMA, "eta0": ETA0, "uses_omega": False}, "deterministic_seeds": [260302577, 271828, 314159, 1618033], "problems": problems, "independent_checker": independent, "proof_replay": proof_replay, "negative_control": { "mutation": "replace exponential schedule by constant eta_t=eta0", "expected": "variance floor; final log-log slope greater than -0.25", "rows": control_rows, "slope": control_slope, "passed_as_control": negative_control_pass, }, "non_circularity": "horizons and feature spectra were predeclared; pass thresholds do not use the paper's target slope; the exact bound is evaluated after simulation, and the constant-step control uses identical data dimensions and noise.", "runtime_seconds": time.perf_counter() - started, }