1import json
  2import math
  3import random
  4from pathlib import Path
  5import numpy as np
  6import torch
  7from torch import nn
  8
  9SEED = 7
 10random.seed(SEED)
 11np.random.seed(SEED)
 12torch.manual_seed(SEED)
 13if torch.cuda.is_available():
 14    torch.cuda.manual_seed_all(SEED)
 15try:
 16    DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 17except Exception:
 18    DEVICE = torch.device("cpu")
 19
 20class DissipativeDrift(nn.Module):
 21    """b=-mu*a+Lf*tanh(a+h(t,x)); action Jacobian norm <= mu+Lf in general,
 22    and the residual has action Lipschitz constant Lf, as required by the claim."""
 23    def __init__(self, state_dim=2, action_dim=2, mu=0.8, lf=0.35):
 24        super().__init__()
 25        self.mu, self.lf = mu, lf
 26        self.center = nn.Sequential(nn.Linear(state_dim + 1, 32), nn.Tanh(),
 27                                     nn.Linear(32, action_dim))
 28    def forward(self, t, x, a):
 29        h = self.center(torch.cat([t, x], dim=-1))
 30        return -self.mu * a + self.lf * torch.tanh(a + h)
 31
 32class GenericDrift(nn.Module):
 33    def __init__(self, state_dim=2, action_dim=2):
 34        super().__init__()
 35        self.net = nn.Sequential(nn.Linear(state_dim + 1 + action_dim, 64), nn.Tanh(),
 36                                 nn.Linear(64, 64), nn.Tanh(), nn.Linear(64, action_dim))
 37    def forward(self, t, x, a):
 38        return self.net(torch.cat([t, x, a], dim=-1))
 39
 40def target_drift(t, x, a):
 41    c = torch.cat([t, x], -1)
 42    h = torch.cat([0.55 * torch.sin(c[:, :1] + c[:, 1:2]),
 43                   0.45 * torch.cos(c[:, :1] - c[:, 2:3])], -1)
 44    return -0.8 * a + h + 0.12 * torch.tanh(2.0 * a)
 45
 46def sample(n, device):
 47    t = torch.rand(n, 1, device=device)
 48    x = torch.randn(n, 2, device=device)
 49    a = 1.5 * torch.randn(n, 2, device=device)
 50    return t, x, a
 51
 52def train(model, steps=1200, batch=128):
 53    model.to(DEVICE)
 54    opt = torch.optim.Adam(model.parameters(), lr=3e-3)
 55    for _ in range(steps):
 56        t, x, a = sample(batch, DEVICE)
 57        loss = ((model(t, x, a) - target_drift(t, x, a)) ** 2).mean()
 58        opt.zero_grad(); loss.backward(); opt.step()
 59    with torch.no_grad():
 60        t, x, a = sample(4096, DEVICE)
 61        mse = ((model(t, x, a) - target_drift(t, x, a)) ** 2).mean().item()
 62    return mse
 63
 64def verify(model, mu=0.8, lf=0.35, n=10000, residual=False):
 65    # Same t,x for each pair, exactly matching the action-coordinate claim.
 66    t, x, a = sample(n, DEVICE)
 67    ap = a + 0.5 * torch.randn_like(a)
 68    with torch.no_grad():
 69        db = model(t, x, a) - model(t, x, ap)
 70        da = a - ap
 71        lhs = (db * da).sum(1)
 72        rhs = -(mu - lf) * (da * da).sum(1)
 73        ratio = (lhs / (da * da).sum(1)).cpu().numpy()
 74        violations = (lhs - rhs > 1e-6).float().mean().item()
 75        # Residual finite-difference Lipschitz estimate, only for the proposed branch.
 76        if residual:
 77            h = model.center(torch.cat([t, x], -1))
 78            fp = lf * torch.tanh(ap + h)
 79            f = lf * torch.tanh(a + h)
 80            lip = (fp - f).norm(dim=1) / (ap - a).norm(dim=1).clamp_min(1e-8)
 81            max_lip = float(lip.max())
 82        else:
 83            max_lip = None
 84    return {"worst_one_sided_ratio": float(ratio.max()),
 85            "theoretical_bound": -(mu-lf),
 86            "violation_fraction": violations,
 87            "max_residual_fd_lipschitz": max_lip}
 88
 89def rollout(model, steps, n=512, dt=0.08, sigma=0.18):
 90    model.eval()
 91    t0 = torch.zeros(n, 1, device=DEVICE)
 92    x = torch.randn(n, 2, device=DEVICE)
 93    a = 2.5 * torch.randn(n, 2, device=DEVICE)
 94    with torch.no_grad():
 95        for k in range(steps):
 96            t = torch.full((n, 1), k * dt, device=DEVICE)
 97            a = a + dt * model(t, x, a) + sigma * math.sqrt(dt) * torch.randn_like(a)
 98    return float(a.norm(dim=1).mean().cpu()), float(a.norm(dim=1).std().cpu())
 99
100def main():
101    # Training is deliberately identical; the comparison is a toy drift-fitting proxy.
102    generic = GenericDrift()
103    dissip = DissipativeDrift()
104    generic_mse = train(generic)
105    dissip_mse = train(dissip)
106    check = verify(dissip, residual=True)
107    generic_check = verify(generic)
108    results = {"device": str(DEVICE), "generic_validation_mse": generic_mse,
109               "dissipative_validation_mse": dissip_mse, "math_check_dissipative": check, "math_check_generic_control": generic_check, "rollouts": {}}
110    for steps in (8, 20, 50):
111        results["rollouts"][str(steps)] = {"generic": rollout(generic, steps),
112                                            "dissipative": rollout(dissip, steps)}
113    Path("results.json").write_text(json.dumps(results, indent=2))
114    print(json.dumps(results, indent=2))
115
116if __name__ == "__main__":
117    main()