Spaces:
Running
Running
| from __future__ import annotations | |
| import math | |
| from concurrent.futures import ProcessPoolExecutor | |
| from dataclasses import dataclass | |
| import numpy as np | |
| REPETITIONS = 160 | |
| MASTER_SEED = 20260802 | |
| class Cell: | |
| family: str | |
| d: int | |
| horizon: int | |
| sigma: float | |
| radius: float | |
| prior_scale: float | |
| anisotropy: float | |
| def source_bound(d: int, horizon: int, sigma: float, radius: float, eigenvalues: np.ndarray) -> dict[str, float]: | |
| q = 24.0 * math.log(horizon) / d | |
| c1 = math.sqrt(1.0 + max(q, math.sqrt(q))) | |
| c2 = 2.0 * c1 * math.sqrt(math.log(1.0 + radius**2 * float(eigenvalues.max()) * horizon / (d * sigma**2))) | |
| minimax = d * sigma * math.sqrt(horizon) * c2 | |
| trace_sqrt = float(np.sqrt(eigenvalues).sum()) | |
| trace = float(eigenvalues.sum()) | |
| prior = 3.0 * math.sqrt(2.0) * radius * math.sqrt(d) * trace_sqrt * c1 | |
| tail = math.sqrt(2.0 * radius**2 * trace) | |
| return { | |
| "C1": c1, | |
| "C2": c2, | |
| "minimax_term": minimax, | |
| "prior_trace_term": prior, | |
| "tail_term": tail, | |
| "theorem_1_bound": minimax + prior + tail, | |
| "corollary_burnin_scale": d * radius * math.sqrt(trace), | |
| "legacy_multiplicative_burnin_scale": d * radius * math.sqrt(trace) * math.sqrt(horizon), | |
| "trace_sqrt": trace_sqrt, | |
| "trace": trace, | |
| "operator_norm": float(eigenvalues.max()), | |
| } | |
| def covariance_eigenvalues(cell: Cell) -> np.ndarray: | |
| if cell.anisotropy == 1.0: | |
| return np.full(cell.d, cell.prior_scale**2) | |
| return cell.prior_scale**2 * np.geomspace(1.0, 1.0 / cell.anisotropy, cell.d) | |
| def simulate_cell(payload: tuple[int, Cell]) -> dict[str, float | int | str | list[float]]: | |
| cell_index, cell = payload | |
| rng = np.random.default_rng(np.random.SeedSequence([MASTER_SEED, cell_index])) | |
| eigenvalues = covariance_eigenvalues(cell) | |
| covariance = np.diag(eigenvalues) | |
| regrets = np.empty(REPETITIONS) | |
| frozen_regrets = np.empty(REPETITIONS) | |
| for repetition in range(REPETITIONS): | |
| theta = rng.multivariate_normal(np.zeros(cell.d), covariance) | |
| mean = np.zeros(cell.d) | |
| posterior = covariance.copy() | |
| regret = 0.0 | |
| frozen_regret = 0.0 | |
| optimal = cell.radius * float(np.linalg.norm(theta)) | |
| for _ in range(cell.horizon): | |
| sampled = rng.multivariate_normal(mean, posterior, check_valid="raise") | |
| sampled_norm = float(np.linalg.norm(sampled)) | |
| action = cell.radius * sampled / max(sampled_norm, 1e-300) | |
| regret += optimal - float(theta @ action) | |
| frozen_sample = rng.multivariate_normal(np.zeros(cell.d), covariance, check_valid="raise") | |
| frozen_action = cell.radius * frozen_sample / max(float(np.linalg.norm(frozen_sample)), 1e-300) | |
| frozen_regret += optimal - float(theta @ frozen_action) | |
| reward = float(theta @ action + rng.normal(0.0, cell.sigma)) | |
| cov_action = posterior @ action | |
| innovation = cell.sigma**2 + float(action @ cov_action) | |
| mean = mean + cov_action * (reward - float(action @ mean)) / innovation | |
| posterior = posterior - np.outer(cov_action, cov_action) / innovation | |
| posterior = 0.5 * (posterior + posterior.T) | |
| regrets[repetition] = regret | |
| frozen_regrets[repetition] = frozen_regret | |
| theorem = source_bound(cell.d, cell.horizon, cell.sigma, cell.radius, eigenvalues) | |
| mean_regret = float(regrets.mean()) | |
| se = float(regrets.std(ddof=1) / math.sqrt(REPETITIONS)) | |
| frozen_mean = float(frozen_regrets.mean()) | |
| frozen_se = float(frozen_regrets.std(ddof=1) / math.sqrt(REPETITIONS)) | |
| return { | |
| "family": cell.family, | |
| "d": cell.d, | |
| "T": cell.horizon, | |
| "sigma": cell.sigma, | |
| "radius": cell.radius, | |
| "prior_scale": cell.prior_scale, | |
| "anisotropy": cell.anisotropy, | |
| "prior_eigenvalues": [float(value) for value in eigenvalues], | |
| "repetitions": REPETITIONS, | |
| "mean_bayesian_regret": mean_regret, | |
| "regret_standard_error": se, | |
| "regret_ci95_low": mean_regret - 1.96 * se, | |
| "regret_ci95_high": mean_regret + 1.96 * se, | |
| "frozen_posterior_regret": frozen_mean, | |
| "frozen_standard_error": frozen_se, | |
| "regret_to_bound_ratio": mean_regret / theorem["theorem_1_bound"], | |
| **theorem, | |
| } | |
| def experiment_cells() -> list[Cell]: | |
| cells = [] | |
| for scale in (0.25, 1.0, 4.0): | |
| for horizon in (64, 128, 256, 512, 1024, 2048): | |
| cells.append(Cell("horizon", 8, horizon, 0.5, 1.0, scale, 1.0)) | |
| for dimension in (2, 4, 8, 16): | |
| cells.append(Cell("dimension", dimension, 512, 0.5, 1.0, 1.0, 1.0)) | |
| for sigma in (0.125, 0.25, 0.5, 1.0, 2.0): | |
| cells.append(Cell("noise", 8, 512, sigma, 1.0, 1.0, 1.0)) | |
| for dimension in (8, 16): | |
| for horizon in (128, 512, 2048): | |
| cells.append(Cell("anisotropic", dimension, horizon, 0.5, 1.0, 4.0, 64.0)) | |
| return cells | |
| def log_slope(rows: list[dict], value: str) -> float: | |
| return float(np.polyfit(np.log([row["T"] for row in rows]), np.log([row[value] for row in rows]), 1)[0]) | |
| def run(worker_count: int) -> dict[str, object]: | |
| cells = experiment_cells() | |
| with ProcessPoolExecutor(max_workers=worker_count) as executor: | |
| rows = list(executor.map(simulate_cell, enumerate(cells))) | |
| horizon_slopes = {} | |
| frozen_slopes = {} | |
| for scale in (0.25, 1.0, 4.0): | |
| subset = [row for row in rows if row["family"] == "horizon" and row["prior_scale"] == scale] | |
| horizon_slopes[str(scale)] = log_slope(subset, "mean_bayesian_regret") | |
| frozen_slopes[str(scale)] = log_slope(subset, "frozen_posterior_regret") | |
| diffuse_ratios = [] | |
| for horizon in (64, 128, 256, 512, 1024, 2048): | |
| diffuse = next(row for row in rows if row["family"] == "horizon" and row["prior_scale"] == 4.0 and row["T"] == horizon) | |
| tight = next(row for row in rows if row["family"] == "horizon" and row["prior_scale"] == 0.25 and row["T"] == horizon) | |
| excess = max(0.0, diffuse["mean_bayesian_regret"] - tight["mean_bayesian_regret"]) | |
| diffuse_ratios.append({"T": horizon, "excess_to_additive_burnin": excess / diffuse["corollary_burnin_scale"]}) | |
| summary = { | |
| "master_seed": MASTER_SEED, | |
| "cells": len(rows), | |
| "paths": len(rows) * REPETITIONS, | |
| "total_simulated_steps": sum(int(row["T"]) * REPETITIONS for row in rows), | |
| "horizon_log_slopes": horizon_slopes, | |
| "frozen_control_log_slopes": frozen_slopes, | |
| "max_regret_to_exact_theorem_1_bound": max(float(row["regret_to_bound_ratio"]) for row in rows), | |
| "max_relative_ci95_halfwidth": max(1.96 * float(row["regret_standard_error"]) / float(row["mean_bayesian_regret"]) for row in rows), | |
| "diffuse_prior_excess_ratios": diffuse_ratios, | |
| "max_diffuse_excess_to_additive_burnin": max(item["excess_to_additive_burnin"] for item in diffuse_ratios), | |
| "min_frozen_control_slope": min(frozen_slopes.values()), | |
| "max_adaptive_slope": max(horizon_slopes.values()), | |
| } | |
| return {"summary": summary, "rows": rows} | |