1import json, math, random
 2import numpy as np
 3import torch
 4import torch.nn as nn
 5import torch.nn.functional as F
 6
 7SEED = 548
 8random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
 9torch.set_num_threads(4)
10try:
11    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
12    if device.type == 'cuda': torch.zeros(1, device=device)
13except Exception:
14    device = torch.device('cpu')
15
16A = torch.tensor([[0.93, 0.10], [-0.08, 0.88]], dtype=torch.float32)
17B = torch.tensor([[0.12], [0.22]], dtype=torch.float32)
18K = torch.tensor([[0.72, 0.42]], dtype=torch.float32)
19
20class Policy(nn.Module):
21    def __init__(self, innovation=True):
22        super().__init__(); self.innovation = innovation
23        self.gru = nn.GRUCell(5, 24)
24        self.z = nn.Linear(24, 2); self.oh = nn.Linear(2, 2)
25        self.nom = nn.Linear(2, 1); self.L = nn.Linear(2, 2); self.lout = nn.Linear(2, 1)
26        self.unc = nn.Linear(2, 1)
27    def forward(self, obs, mask, uprev, h=None, c=2.0):
28        if h is None: h = torch.zeros(obs.shape[0], 24, device=obs.device)
29        h = self.gru(torch.cat([obs * mask, mask, uprev], 1), h)
30        z = self.z(h); pred = self.oh(z); nu = obs - pred
31        uncertainty = F.softplus(self.unc(z)) + 0.02
32        alpha = (1.0 / (1.0 + c * uncertainty)).clamp(0, 1)
33        nominal = self.nom(z)
34        correction = self.lout(torch.tanh(self.L(z)) * nu)
35        u = nominal - (alpha * correction if self.innovation else 0.0)
36        return u.clamp(-1, 1), pred, nu, uncertainty, h
37
38def batch_rollout(policy, batch=96, horizon=35, corrupt=False, train=True):
39    Adev, Bdev = A.to(device), B.to(device)
40    x = torch.randn(batch, 2, device=device) * 0.7
41    h = None; up = torch.zeros(batch, 1, device=device)
42    xs=[]; obs=[]; masks=[]; us=[]; preds=[]; nus=[]; uncs=[]
43    for t in range(horizon):
44        m = torch.ones(batch, 2, device=device)
45        if corrupt and 12 <= t < 18: m.zero_()
46        noise = torch.randn(batch,2,device=device)*0.10
47        o = (x + noise) * m
48        u,p,n,q,h = policy(o,m,up,h)
49        # reference is zero; controller should dissipate state.
50        xnext = x @ Adev.T + u @ Bdev.T + torch.randn(batch,2,device=device)*0.018
51        xs.append(x); obs.append(o); masks.append(m); us.append(u); preds.append(p); nus.append(n); uncs.append(q)
52        x=xnext; up=u
53    return map(torch.stack, (xs,obs,masks,us,preds,nus,uncs))
54
55def train(policy, steps=220):
56    opt=torch.optim.Adam(policy.parameters(),lr=3e-3)
57    policy.train()
58    for step in range(steps):
59        x,o,m,u,p,n,q = batch_rollout(policy, batch=64, horizon=28, train=True)
60        # Regulation, action effort, and estimator prediction objectives.
61        loss = x.pow(2).sum(-1).mean() + 0.035*u.pow(2).mean() + 0.18*(o-p).pow(2).mean()
62        opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(policy.parameters(),1.0); opt.step()
63    return float(loss.detach().cpu())
64
65def evaluate(policy, corrupt=False, episodes=192):
66    policy.eval()
67    with torch.no_grad():
68        x,o,m,u,p,n,q = batch_rollout(policy,episodes,45,corrupt)
69        state_cost=x.pow(2).sum(-1)
70        rms=float(torch.sqrt(state_cost.mean()).cpu())
71        transient=float(torch.sqrt(state_cost[:,12:20].mean()).cpu())
72        recovery=float(torch.sqrt(state_cost[:,20:].mean()).cpu())
73        action_std=float(u.std().cpu())
74        pred=float((o-p).pow(2).mean().sqrt().cpu())
75    return {'rms_state_error':rms,'transient_error':transient,'post_corruption_error':recovery,'action_std':action_std,'innovation_rmse':pred}
76
77def math_check():
78    # The gate alpha=1/(1+c sigma) is monotone and gain-limiting; verify numerically.
79    sig=torch.linspace(0,10,1001); c=2.0; alpha=1/(1+c*sig)
80    monotone=bool(torch.all(alpha[1:] <= alpha[:-1]))
81    bounds=bool(torch.all((alpha>=0)&(alpha<=1)))
82    # For a scalar innovation correction, increasing uncertainty strictly shrinks it.
83    correction=torch.tensor(1.7)*alpha
84    shrink=bool(correction[0] > correction[-1] and correction[0] <= 1.7)
85    return {'gate_monotone':monotone,'gate_bounds':bounds,'correction_shrinks':shrink,'alpha_sigma0':float(alpha[0]),'alpha_sigma10':float(alpha[-1])}
86
87def main():
88    check=math_check(); results={'device':str(device),'math_check':check}
89    base=Policy(False).to(device); idea=Policy(True).to(device)
90    results['train_loss_baseline']=train(base); results['train_loss_idea']=train(idea)
91    results['clean_baseline']=evaluate(base,False); results['clean_idea']=evaluate(idea,False)
92    results['corrupt_baseline']=evaluate(base,True); results['corrupt_idea']=evaluate(idea,True)
93    with open('results.json','w') as f: json.dump(results,f,indent=2)
94    print(json.dumps(results,indent=2))
95
96if __name__=='__main__': main()