1import json, math, time
  2from pathlib import Path
  3import numpy as np
  4import torch
  5from torch import nn
  6
  7SEED = 640
  8np.random.seed(SEED)
  9torch.manual_seed(SEED)
 10torch.set_num_threads(4)
 11device = "cuda" if torch.cuda.is_available() else "cpu"
 12try:
 13    if device == "cuda":
 14        torch.cuda.set_device(0)
 15        torch.zeros(1, device=device)
 16except Exception:
 17    device = "cpu"
 18
 19BETA, TAU, KTRAIN, KEVAL = 0.92, 1.4, 8, 256
 20
 21def logmeanexp(x, dim=-1):
 22    return torch.logsumexp(x, dim=dim) - math.log(x.shape[dim])
 23
 24def transition(s, a, noise):
 25    return torch.tanh(0.72*s + 0.35*a + noise)
 26
 27def reward(s, a):
 28    return 1.0 - 0.55*s.square() - 0.12*a.square()
 29
 30def mlp(out=1):
 31    return nn.Sequential(nn.Linear(2, 48), nn.Tanh(), nn.Linear(48, 48), nn.Tanh(), nn.Linear(48, out)).to(device)
 32
 33def sample_batch(n, k):
 34    s = (torch.rand(n, 1, device=device)*2-1)
 35    a = (torch.rand(n, 1, device=device)*2-1)
 36    noise = torch.randn(n, k, 1, device=device)*0.38
 37    sn = transition(s[:,None,:], a[:,None,:], noise)
 38    return s, a, reward(s,a), sn
 39
 40@torch.no_grad()
 41def high_k_target(vtarget, s, a, k=KEVAL):
 42    noise = torch.randn(s.shape[0], k, 1, device=device)*0.38
 43    sn = transition(s[:,None,:], a[:,None,:], noise)
 44    inp = torch.cat([sn.reshape(-1,1), a[:,None,:].expand(-1,k,-1).reshape(-1,1)], 1)
 45    vals = vtarget(inp).reshape(s.shape[0], k)
 46    return -logmeanexp(-TAU*vals, 1)/TAU
 47
 48def run(seed, mode, steps=1200):
 49    torch.manual_seed(seed)
 50    v = mlp(); vt = mlp(); vt.load_state_dict(v.state_dict())
 51    m = mlp() if mode == 'aux' else None
 52    opt = torch.optim.Adam(list(v.parameters()) + ([] if m is None else list(m.parameters())), lr=1e-3)
 53    records=[]; t0=time.perf_counter()
 54    for step in range(steps):
 55        s,a,r,sn = sample_batch(64,KTRAIN)
 56        if mode == 'direct':
 57            with torch.no_grad():
 58                inp=torch.cat([sn.reshape(-1,1), a[:,None,:].expand(-1,KTRAIN,-1).reshape(-1,1)],1)
 59                vv=vt(inp).reshape(64,KTRAIN)
 60                ce=-logmeanexp(-TAU*vv,1)
 61                target=r+BETA*ce
 62            pred=v(torch.cat([s,a],1)).squeeze(1)
 63            loss=(pred-target).square().mean()
 64        else:
 65            with torch.no_grad():
 66                inp=torch.cat([sn.reshape(-1,1), a[:,None,:].expand(-1,KTRAIN,-1).reshape(-1,1)],1)
 67                vv=vt(inp).reshape(64,KTRAIN)
 68                ce=-logmeanexp(-TAU*vv,1)
 69            ma=m(torch.cat([s,a],1)).squeeze(1)
 70            pred=v(torch.cat([s,a],1)).squeeze(1)
 71            bell=(r+BETA*ma-pred).square().mean()
 72            cel=(ma-ce).square().mean()
 73            loss=bell+0.7*cel
 74        opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(list(v.parameters())+([] if m is None else list(m.parameters())), 10); opt.step()
 75        with torch.no_grad():
 76            for p,q in zip(vt.parameters(),v.parameters()): p.mul_(0.97).add_(q,alpha=0.03)
 77        if step in (steps//2, steps-1):
 78            s0,a0,_,_=sample_batch(256,1)
 79            true=high_k_target(vt,s0,a0)
 80            with torch.no_grad():
 81                if mode=='direct':
 82                    est=high_k_target(vt,s0,a0,KTRAIN) # direct target estimator distribution
 83                    auxerr=float('nan')
 84                    predv=v(torch.cat([s0,a0],1)).squeeze(1)
 85                    # Bellman residual using high-K CE
 86                    br=(reward(s0,a0)+BETA*true-predv)
 87                else:
 88                    ma=m(torch.cat([s0,a0],1)).squeeze(1)
 89                    auxerr=(ma-true).abs().mean().item()
 90                    br=(reward(s0,a0)+BETA*ma-v(torch.cat([s0,a0],1)).squeeze(1))
 91                records.append({'step':step+1,'ce_abs_error':auxerr,'bellman_rmse':float(br.square().mean().sqrt()),'loss':float(loss)})
 92    # estimator variance at fixed states, over repeated K samples
 93    with torch.no_grad():
 94        s,a,_,_=sample_batch(128,1)
 95        vals=[]
 96        for _ in range(30): vals.append(high_k_target(vt,s,a,KTRAIN).cpu().numpy())
 97        estvar=float(np.mean(np.var(np.stack(vals),axis=0)))
 98        true=high_k_target(vt,s,a,KEVAL)
 99        vpred=v(torch.cat([s,a],1)).squeeze(1)
100        value_rmse=float((vpred- (reward(s,a)+BETA*true)).square().mean().sqrt())
101        if m is not None: ce_rmse=float((m(torch.cat([s,a],1)).squeeze(1)-true).square().mean().sqrt())
102        else: ce_rmse=float('nan')
103    return {'mode':mode,'seed':seed,'device':device,'seconds':time.perf_counter()-t0,'estimator_variance':estvar,'value_target_rmse':value_rmse,'ce_rmse':ce_rmse,'checkpoints':records}
104
105def math_check():
106    # Exact equality for constant continuation values, and convergence of sample CE.
107    x=torch.tensor([[1.,2.,3.]])
108    got=float((-logmeanexp(-TAU*x,1)/TAU).item())
109    expected=float(-math.log(np.mean(np.exp(-TAU*x.numpy())))/TAU)
110    rng=np.random.default_rng(SEED)
111    errs=[]
112    for k in (2,8,32,128):
113        z=rng.normal(size=(1000*k,))
114        mx0=np.max(-TAU*z)
115        exact=-(mx0+math.log(np.mean(np.exp(-TAU*z-mx0))))/TAU
116        e=[]
117        for i in range(1000):
118            q=z[i*k:(i+1)*k]
119            mx=np.max(-TAU*q)
120            e.append((-(mx+math.log(np.mean(np.exp(-TAU*q-mx))))/TAU)-exact)
121        errs.append({'K':k,'abs_error_mean':float(np.mean(np.abs(e))),'bias':float(np.mean(e))})
122    return {'logmeanexp_absolute_difference':abs(got-expected),'sample_scaling':errs}
123
124if __name__=='__main__':
125    out={'math_check':math_check(),'runs':[]}
126    for seed in (11,22,33):
127        for mode in ('direct','aux'): out['runs'].append(run(seed,mode))
128    Path('results.json').write_text(json.dumps(out,indent=2))
129    print(json.dumps(out,indent=2))