1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5from torch import nn
  6
  7SEED = 3082
  8random.seed(SEED)
  9np.random.seed(SEED)
 10torch.manual_seed(SEED)
 11torch.set_num_threads(4)
 12try:
 13    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 14    if device.type == "cuda":
 15        torch.cuda.set_device(0)
 16        _ = torch.zeros(1, device=device)
 17except Exception:
 18    device = torch.device("cpu")
 19
 20# A small robust-control Bellman envelope.  The action is a in [-1,1],
 21# and the adversarial disturbance is w in {-d,+d}.
 22def q_exact(x, a, d=0.18):
 23    z1 = -(x + a - d) ** 2 - 0.1 * a ** 2
 24    z2 = -(x + a + d) ** 2 - 0.1 * a ** 2
 25    return torch.minimum(z1, z2)
 26
 27def exact_target(x, n=4001):
 28    a = torch.linspace(-1., 1., n, device=x.device)
 29    return q_exact(x[:, None], a[None, :]).amax(1)
 30
 31def math_check():
 32    x = torch.linspace(-1.2, 1.2, 401)
 33    a = torch.linspace(-1., 1., 10001)
 34    vals = q_exact(x[:, None], a[None, :])
 35    da = float(a[1] - a[0])
 36    L = float((vals[:, 1:] - vals[:, :-1]).abs().max() / da)
 37    rows = []
 38    for M in [5, 9, 17, 33, 65]:
 39        grid = torch.linspace(-1., 1., M)
 40        disc = q_exact(x[:, None], grid[None, :]).amax(1)
 41        err = float((vals.amax(1) - disc).abs().max())
 42        rows.append({"M": M, "delta": 2.0/(M-1), "max_error": err,
 43                     "L_delta": L * 2.0/(M-1)})
 44    # max is monotone in every branch: increasing one branch cannot decrease output
 45    branches = torch.randn(200, 11)
 46    base = branches.max(1).values
 47    bumped = branches.clone(); bumped[:, 4] += torch.rand(200) * 2.0
 48    monotone_violations = int((bumped.max(1).values < base - 1e-7).sum())
 49    slopes = np.array([r["max_error"] for r in rows])
 50    deltas = np.array([r["delta"] for r in rows])
 51    slope = float(np.polyfit(np.log(deltas), np.log(slopes + 1e-12), 1)[0])
 52    return {"action_lipschitz_estimate": L, "discretization": rows,
 53            "loglog_error_slope": slope,
 54            "max_monotonicity_violations": monotone_violations}
 55
 56class Baseline(nn.Module):
 57    def __init__(self, width=64):
 58        super().__init__()
 59        self.net = nn.Sequential(nn.Linear(1, width), nn.Tanh(),
 60                                 nn.Linear(width, width), nn.Tanh(),
 61                                 nn.Linear(width, 1))
 62    def forward(self, x):
 63        return self.net(x[:, None]).squeeze(1)
 64
 65class Envelope(nn.Module):
 66    def __init__(self, M=17, width=32, tau=0.03):
 67        super().__init__()
 68        self.M, self.tau = M, tau
 69        self.actions = torch.linspace(-1., 1., M)
 70        self.branches = nn.ModuleList([
 71            nn.Sequential(nn.Linear(1, width), nn.Tanh(), nn.Linear(width, width),
 72                          nn.Tanh(), nn.Linear(width, 1)) for _ in range(M)])
 73        # Learned branch penalties, initialized near the known running cost.
 74        self.penalty = nn.Parameter(0.1 * self.actions.square())
 75    def forward(self, x, hard=False):
 76        ys = torch.stack([net(x[:, None]).squeeze(1) for net in self.branches], 1)
 77        q = ys - self.penalty[None, :]
 78        return q.max(1).values if hard else self.tau * torch.logsumexp(q / self.tau, 1)
 79
 80def train(model, xtr, ytr, steps=1400):
 81    model.to(device)
 82    opt = torch.optim.Adam(model.parameters(), lr=2e-3)
 83    for step in range(steps):
 84        ix = torch.randint(0, len(xtr), (128,), device=device)
 85        pred = model(xtr[ix])
 86        loss = (pred-ytr[ix]).square().mean()
 87        opt.zero_grad(); loss.backward(); opt.step()
 88    return model
 89
 90def experiment():
 91    gen = torch.Generator(device=device).manual_seed(SEED)
 92    xtr = (torch.rand(1024, generator=gen, device=device)*2.4-1.2)
 93    ytr = exact_target(xtr)
 94    xte = torch.linspace(-1.2, 1.2, 801, device=device)
 95    yte = exact_target(xte)
 96    results = {}
 97    for name, model in [("baseline", Baseline()), ("envelope", Envelope())]:
 98        train(model, xtr, ytr)
 99        with torch.no_grad():
100            pred = model(xte, hard=True) if name == "envelope" else model(xte)
101            mse = float((pred-yte).square().mean())
102            mae = float((pred-yte).abs().mean())
103            worst = float((pred-yte).abs().max())
104        results[name] = {"mse": mse, "mae": mae, "max_abs_error": worst}
105    # Refinement uses the exact finite action envelope, isolating the claimed grid effect.
106    refine = {}
107    for M in [5, 9, 17, 33, 65]:
108        grid = torch.linspace(-1., 1., M, device=device)
109        with torch.no_grad():
110            p = q_exact(xte[:, None], grid[None, :]).amax(1)
111            refine[str(M)] = float((p-yte).abs().max())
112    results["refinement_max_error"] = refine
113    return results
114
115if __name__ == "__main__":
116    out = {"device": str(device), "math": math_check(), "experiment": experiment()}
117    Path("results.json").write_text(json.dumps(out, indent=2))
118    print(json.dumps(out, indent=2))