1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5
  6SEED = 1217
  7np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
  8torch.set_num_threads(4)
  9
 10# One-step stochastic point-mass abstraction. z is a severity/initial-state feature.
 11# The policy is u=theta*z; task optimum theta=1, while safety requires theta*z <= c.
 12# Thus g(theta,z)=theta*z-c and v=[g]_+ (scale s=1).
 13C = 0.62
 14B = 32
 15ROUNDS = 18
 16UPDATES = 12
 17BETA = 10.0
 18LR = 0.045
 19T = 0.035
 20
 21def soft_penalty(g):
 22    return torch.nn.functional.softplus(g / T).pow(2) * T*T
 23
 24def update(theta, z, picked=False):
 25    # z is a tensor of trajectories in the current uniform batch or constraint buffer.
 26    theta = theta.detach().clone().requires_grad_(True)
 27    g = theta * z - C
 28    task = (theta - 1.0).pow(2)
 29    loss = task + BETA * soft_penalty(g).mean()
 30    loss.backward()
 31    with torch.no_grad():
 32        theta -= LR * theta.grad
 33    return theta.detach(), float(loss.detach())
 34
 35def run(method, seed=SEED, record=False):
 36    rng = np.random.default_rng(seed)
 37    theta = torch.tensor(1.0)
 38    buffer = []
 39    history = []
 40    probe = torch.linspace(0, 1, 4001)
 41    for m in range(ROUNDS):
 42        z_np = rng.random(B)
 43        z = torch.tensor(z_np, dtype=torch.float32)
 44        v_np = np.maximum(theta.item() * z_np - C, 0)
 45        if method == 'pick':
 46            # The exact argmax trajectory is retained, as prescribed by C_{m+1}.
 47            j = int(np.argmax(v_np))
 48            buffer.append(float(z_np[j]))
 49            train_z = torch.tensor(buffer, dtype=torch.float32)
 50        else:
 51            train_z = z
 52        for _ in range(UPDATES):
 53            theta, _ = update(theta, train_z)
 54        probe_v = torch.relu(theta * probe - C)
 55        history.append({'round': m+1, 'theta': float(theta),
 56                        'batch_max_v': float(v_np.max()),
 57                        'probe_max_v': float(probe_v.max())})
 58    return theta.item(), history
 59
 60def certify(theta, seed=991):
 61    rng = np.random.default_rng(seed)
 62    out = []
 63    delta = .05
 64    for n in [50, 100, 200, 400, 600, 1200, 2400]:
 65        z = rng.random(n)
 66        v = np.maximum(theta*z-C, 0)
 67        phat = float(np.mean(v > 0))
 68        eps = math.sqrt(math.log(1/delta)/(2*n))
 69        out.append({'N': n, 'phat': phat, 'upper': min(1., phat+eps), 'epsilon': eps})
 70    # The correction itself is the claimed O(N^-1/2) object.
 71    x = np.log([r['N'] for r in out]); y = np.log([r['epsilon'] for r in out])
 72    slope = float(np.polyfit(x, y, 1)[0])
 73    return out, slope
 74
 75def main():
 76    # Prediction 1: for z~U[0,1], E[max z]=B/(B+1), hence selected max
 77    # should approach 1 and exceed the random-batch mean ~1/2.
 78    order_rows = []
 79    for b in [4, 8, 16, 32, 64, 128]:
 80        rng = np.random.default_rng(700+b)
 81        reps = 20000
 82        maxima = rng.random((reps,b)).max(axis=1)
 83        means = rng.random((reps,b)).mean(axis=1)
 84        order_rows.append({'B': b, 'observed_max': float(maxima.mean()),
 85                           'predicted_max': b/(b+1),
 86                           'observed_batch_mean': float(means.mean()),
 87                           'predicted_batch_mean': .5,
 88                           'max_minus_mean': float((maxima-means).mean())})
 89    # Prediction 2: the adaptive penalty should lower the worst-case probe margin
 90    # more than uniform empirical-risk updates at equal rollout budget.
 91    pick_theta, pick_hist = run('pick')
 92    uni_theta, uni_hist = run('uniform')
 93    ztest = np.linspace(0,1,200001)
 94    def eval_theta(th):
 95        v = np.maximum(th*ztest-C,0)
 96        return {'theta': th, 'violation_rate': float(np.mean(v>0)),
 97                'max_violation': float(v.max()), 'mean_violation': float(v.mean())}
 98    # Prediction 3: fixed-policy certificate correction has log-log slope -1/2.
 99    cert, cert_slope = certify(pick_theta)
100    # Independent held-out certification requested by the idea.
101    rng = np.random.default_rng(4242)
102    zh = rng.random(600); vh = np.maximum(pick_theta*zh-C,0)
103    heldout = {'N': 600, 'violations': int(np.sum(vh>0)),
104               'phat': float(np.mean(vh>0)),
105               'upper_95': float(min(1., np.mean(vh>0)+math.sqrt(math.log(20)/(1200))))}
106    result = {
107      'config': {'C': C, 'batch': B, 'rounds': ROUNDS, 'updates': UPDATES,
108                 'beta': BETA, 'lr': LR, 'temperature': T, 'seed': SEED},
109      'prediction_1_order_statistics': order_rows,
110      'prediction_2_equal_budget': {'pick': eval_theta(pick_theta),
111                                     'uniform': eval_theta(uni_theta),
112                                     'pick_history': pick_hist,
113                                     'uniform_history': uni_hist,
114                                     'rollouts_each': B*ROUNDS},
115      'prediction_3_certificate': {'rows': cert, 'observed_loglog_slope': cert_slope,
116                                   'predicted_slope': -0.5, 'heldout_600': heldout},
117      'notes': 'Violation is g=[theta*z-C], with U[0,1] severity and fixed policy during certification.'
118    }
119    Path('results.json').write_text(json.dumps(result, indent=2))
120    print(json.dumps(result, indent=2))
121
122if __name__ == '__main__': main()