1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5
  6SEED = 2345
  7np.set_printoptions(precision=6, suppress=True)
  8
  9def seed_all(seed):
 10    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 11
 12
 13def normalize(A, eps=1e-8):
 14    mu = A.mean(axis=0, keepdims=True)
 15    sig = A.std(axis=0, keepdims=True)
 16    return (A-mu)/(sig+eps), mu, sig
 17
 18
 19def toy_verification():
 20    # Two equally-sized return groups with controlled separation d. Every timestep
 21    # has the same return ordering, so the conditional-margin prediction is exact.
 22    B, T = 4000, 5
 23    rng = np.random.default_rng(SEED)
 24    feasible = np.zeros(B, dtype=bool); feasible[:B//2] = True
 25    rng.shuffle(feasible)
 26    noise = rng.normal(size=(B,T))
 27    results = {}
 28
 29    # Prediction 1: normalized variance is (sigma/(sigma+eps))^2, approximately one.
 30    var_rows=[]
 31    for scale in [0.05, 0.2, 1.0, 5.0, 20.0]:
 32        A = scale * noise
 33        At, _, sig = normalize(A)
 34        observed = float(At.var(axis=0).mean())
 35        predicted = float(np.mean((sig/(sig+1e-8))**2))
 36        var_rows.append({'scale':scale, 'observed_variance':observed, 'predicted':predicted})
 37
 38    # Prediction 2/3: with equal groups and feasible-return mean higher by d,
 39    # E[At|F]-E[At|V] = d/sigma; weighted margin is (wF+wV)d/(2 sigma).
 40    margin_rows=[]
 41    base_noise = rng.normal(scale=0.25, size=(B,T))
 42    for d in [0.0, 0.2, 0.5, 1.0, 2.0]:
 43        G = base_noise.copy()
 44        G[feasible] += d/2; G[~feasible] -= d/2
 45        At, _, sig = normalize(G)
 46        mf = float(At[feasible].mean()); mv = float(At[~feasible].mean())
 47        sigma = float(sig.mean())
 48        for ratio in [1.0, 2.0, 4.0]:
 49            wf, wv = ratio, 1.0
 50            observed = wf*mf - wv*mv
 51            predicted = (wf+wv)*d/(2*sigma) if sigma > 0 else 0.0
 52            margin_rows.append({'separation':d, 'wF':wf, 'wV':wv,
 53                                'observed_margin':observed, 'predicted_margin':predicted})
 54    # Fit observed margin slope at wF=2,wV=1, compare formula slope.
 55    selected=[r for r in margin_rows if r['wF']==2.0]
 56    x=np.array([r['separation'] for r in selected]); y=np.array([r['observed_margin'] for r in selected])
 57    slope=float(np.polyfit(x,y,1)[0])
 58    sigma_ref=float(np.mean(np.abs(base_noise))) # only informational
 59    return {'normalization_sweep':var_rows, 'margin_sweep':margin_rows,
 60            'margin_observed_slope_wF2_wV1':slope,
 61            'margin_prediction_slope_note':'(wF+wV)/(2*sigma), with sigma measured per sweep'}
 62
 63
 64def rollout(theta, B=64, T=5, device='cpu'):
 65    # action 0 is feasible/safe and pays .8; action 1 is risky and pays 1.2.
 66    p=torch.sigmoid(theta)
 67    acts=torch.bernoulli(p.expand(B,T))
 68    rewards=torch.where(acts < .5, torch.tensor(.8,device=device), torch.tensor(1.2,device=device))
 69    feasible=(acts.sum(dim=1)==0)
 70    G=rewards.sum(dim=1)
 71    return acts, rewards, feasible, G
 72
 73
 74def update(theta, old_logp, acts, G, feasible, ranked, clip_eps=.2, wf=3., wv=.25):
 75    # Group-relative trajectory advantage, repeated across time, then timestep norm.
 76    raw=(G-G.mean()).unsqueeze(1).expand_as(acts)
 77    mu=raw.mean(dim=0,keepdim=True); sd=raw.std(dim=0,keepdim=True,unbiased=False)
 78    A=(raw-mu)/(sd+1e-8)
 79    p=torch.sigmoid(theta)
 80    logp=acts*torch.log(p+1e-8)+(1-acts)*torch.log(1-p+1e-8)
 81    ratio=torch.exp(logp-old_logp)
 82    clipped=torch.clamp(ratio,1-clip_eps,1+clip_eps)
 83    weights=torch.where(feasible, torch.tensor(wf), torch.tensor(wv)) if ranked else torch.ones_like(feasible,dtype=torch.float)
 84    obj=torch.minimum(ratio*A,clipped*A)*weights.unsqueeze(1)
 85    loss=-obj.mean()
 86    loss.backward()
 87    return float(loss.detach()), float(A.var(dim=0,unbiased=False).mean()), float(feasible.float().mean())
 88
 89
 90def mini_experiment():
 91    # Same sampled-rollout setup and seeds for standard normalized GRPO vs ranking.
 92    device='cuda' if torch.cuda.is_available() else 'cpu'
 93    rows=[]
 94    for ranked in [False, True]:
 95        finals=[]; rewards=[]; curves=[]
 96        for s in range(8):
 97            seed_all(SEED+s)
 98            theta=torch.tensor(0.0,device=device,requires_grad=True)
 99            opt=torch.optim.SGD([theta],lr=.18)
100            hist=[]
101            for it in range(180):
102                opt.zero_grad()
103                acts,rews,feas,G=rollout(theta,device=device)
104                p=torch.sigmoid(theta)
105                oldlp=acts*torch.log(p+1e-8)+(1-acts)*torch.log(1-p+1e-8)
106                loss,var,fb=update(theta,oldlp.detach(),acts,G,feas,ranked)
107                opt.step()
108                hist.append((fb,float(rews.mean())))
109            with torch.no_grad():
110                p=float(torch.sigmoid(theta)); final_feas=(1-p)**5
111                finals.append(final_feas); rewards.append(.8*5*(1-p)+1.2*5*p)
112                curves.append(hist)
113        rows.append({'method':'ranked' if ranked else 'standard',
114                     'final_feasibility_mean':float(np.mean(finals)),
115                     'final_feasibility_std':float(np.std(finals)),
116                     'policy_risky_probability_mean':float(1-np.mean(np.array(finals)**(1/5))),
117                     'expected_reward_mean':float(np.mean(rewards)),
118                     'training_rollout_feasibility_last20':float(np.mean([np.mean(c[-20:],axis=0)[0] for c in curves]))})
119    return {'device':device,'rows':rows}
120
121
122def main():
123    seed_all(SEED)
124    try:
125        mini = mini_experiment()
126    except Exception as exc:
127        print('CUDA/accelerator run failed; retrying on CPU:', repr(exc))
128        if torch.cuda.is_available():
129            torch.cuda.empty_cache()
130        old = torch.cuda.is_available
131        torch.cuda.is_available = lambda: False
132        try:
133            mini = mini_experiment()
134        finally:
135            torch.cuda.is_available = old
136    out={'toy_verification':toy_verification(),'mini_experiment':mini}
137    Path('results.json').write_text(json.dumps(out,indent=2))
138    print(json.dumps(out,indent=2))
139
140if __name__=='__main__': main()