1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6
  7SEED = 1599
  8random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  9torch.set_num_threads(4)
 10DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 11try:
 12    if DEVICE == "cuda":
 13        torch.zeros(1, device="cuda").sum().item()
 14except Exception:
 15    DEVICE = "cpu"
 16
 17
 18def step(w, dt, s, admission):
 19    return torch.relu(w - dt) + admission * s
 20
 21
 22def deterministic_check():
 23    # Baseline admits a small first job and rejects a large second job.
 24    # Controlled does the opposite: it is initially lower, then exceeds baseline.
 25    dt, s1, s2 = 0.1, 2.0, 10.0
 26    wb = step(torch.tensor(0.), 0., s1, torch.tensor(1.))
 27    wc = step(torch.tensor(0.), 0., s1, torch.tensor(0.))
 28    wb = step(wb, dt, s2, torch.tensor(0.))
 29    wc = step(wc, dt, s2, torch.tensor(1.))
 30    return {"baseline_after_second": float(wb), "controlled_after_second": float(wc),
 31            "violation": float(torch.relu(wc-wb)), "order_loss": float(torch.relu(wc-wb)**2)}
 32
 33
 34def math_sweeps():
 35    # With first-job residual b and controlled residual zero, the next-job
 36    # violation is max(S2-b,0), hence onset at S2=b and quadratic loss above it.
 37    b, s2 = 1.9, np.linspace(0, 5, 101)
 38    observed = np.maximum(s2-b, 0.)
 39    predicted = np.maximum(s2-b, 0.)
 40    onset = float(s2[np.argmax(observed > 1e-9)])
 41    # Fit slope of loss vs (S2-b)^2 only in the active region.
 42    active = observed > 0
 43    x = (s2[active]-b)**2; y = observed[active]**2
 44    slope = float(np.dot(x,y)/np.dot(x,x))
 45    # Scaling sweep directly changes the excess gap.
 46    scales = np.array([0.5, 1., 1.5, 2.])
 47    gaps = np.maximum(scales*4-b, 0.)
 48    ratios = gaps**2 / (gaps[1]**2)
 49    return {"threshold_prediction": b, "threshold_observed_grid": onset,
 50            "quadratic_prediction_slope": 1.0, "quadratic_observed_slope": slope,
 51            "scale_factors": scales.tolist(), "loss_ratio_observed": ratios.tolist(),
 52            "loss_ratio_predicted": ((gaps/gaps[1])**2).tolist()}
 53
 54
 55def traces(n, device):
 56    # Paired arrivals and service costs. Baseline is a state-dependent reference
 57    # gate; the learned gate sees the same workload, size, and position.
 58    dt = torch.rand(n, 6, device=device) * .35 + .05
 59    s = torch.exp(torch.randn(n, 6, device=device)*.55 + .35).clamp(.2, 8.)
 60    return dt, s
 61
 62class Gate(nn.Module):
 63    def __init__(self):
 64        super().__init__()
 65        self.net = nn.Sequential(nn.Linear(4, 24), nn.Tanh(), nn.Linear(24, 1))
 66    def forward(self, x):
 67        return torch.sigmoid(self.net(x)).squeeze(-1)
 68
 69def baseline_decision(w, s):
 70    # Reference production policy: avoid admitting a job if current workload
 71    # is high or if the job is large. This creates the feedback counterexample.
 72    return ((w < 1.5) & (s < 3.0)).float()
 73
 74def rollout(model, dt, s, lam):
 75    n, k = s.shape
 76    wb = torch.zeros(n, device=dt.device); wc = torch.zeros(n, device=dt.device)
 77    violations=[]; probs=[]; accepted=[]; baseacc=[]
 78    for j in range(k):
 79        xb = torch.stack([wc, s[:,j], dt[:,j], torch.full_like(wc, j/(k-1))], 1)
 80        p = model(xb)
 81        ib = baseline_decision(wb, s[:,j])
 82        # Soft admission is the differentiable training surrogate.
 83        wb = step(wb, dt[:,j], s[:,j], ib)
 84        wc = step(wc, dt[:,j], s[:,j], p)
 85        v = torch.relu(wc-wb)
 86        violations.append(v); probs.append(p); accepted.append(p); baseacc.append(ib)
 87    v = torch.stack(violations,1); p = torch.stack(probs,1)
 88    # Throughput reward, mild workload regularization, plus proposed order term.
 89    objective = -p.mean() + .01*wc.mean() + lam*(v**2).mean()
 90    return objective, v.detach(), p.detach(), torch.stack(baseacc,1).detach(), wc.detach(), wb.detach()
 91
 92def train(lam, device):
 93    model=Gate().to(device); opt=torch.optim.Adam(model.parameters(),lr=.025)
 94    for epoch in range(350):
 95        dt,s=traces(96,device)
 96        loss,*_=rollout(model,dt,s,lam)
 97        opt.zero_grad(); loss.backward(); opt.step()
 98    with torch.no_grad():
 99        dt,s=traces(2048,device)
100        _,v,p,ba,wc,wb=rollout(model,dt,s,lam)
101        return {"violation_frequency":float((v>1e-5).any(1).float().mean()),
102                "mean_violation":float(v.mean()), "p99_violation":float(torch.quantile(v.flatten(),.99)),
103                "mean_admission_probability":float(p.mean()), "mean_final_workload":float(wc.mean()),
104                "baseline_mean_final_workload":float(wb.mean()),
105                "baseline_admission_rate":float(ba.mean())}
106
107def main():
108    checks=math_sweeps(); exact=deterministic_check()
109    results={}
110    for lam in (0.0, 0.5, 2.0):
111        results[str(lam)]=train(lam,DEVICE)
112    out={"device":DEVICE,"exact_counterexample":exact,"math_sweeps":checks,"training":results}
113    Path("results.json").write_text(json.dumps(out,indent=2))
114    print(json.dumps(out,indent=2))
115
116if __name__=="__main__": main()