Create stage2.py

stage2.py

@@ -0,0 +1,109 @@
+# stage2.py
+# Author: Liam Grinstead
+# Purpose: Orbital & Agent Coupling Validation (Stage Two of Twelve)
+
+import torch, math, time, json, random, argparse, numpy as np
+
+# ---------------- Determinism ----------------
+def set_seed(s=1234):
+    random.seed(s); np.random.seed(s)
+    torch.manual_seed(s); torch.cuda.manual_seed_all(s)
+
+# ---------------- Telemetry ------------------
+class Telemetry:
+    def __init__(self, path=None):
+        self.t0 = time.time()
+        self.f = open(path, "w") if path else None
+    def emit(self, **k):
+        k["t"] = round(time.time() - self.t0, 3)
+        line = json.dumps(k, separators=(",", ":"))
+        print(line)
+        if self.f:
+            self.f.write(line + "\n"); self.f.flush()
+    def close(self):
+        if self.f: self.f.close()
+
+# ---------------- Orbital Coupler ------------
+class Orbital:
+    def __init__(self, g=0.006, floor=0.2):
+        self.a = 0.0; self.b = math.pi/3; self.g = g; self.floor = floor
+    def step(self):
+        diff = (self.b - self.a + math.pi) % (2*math.pi) - math.pi
+        if abs(diff) < self.floor:
+            diff = self.floor * (1 if diff >= 0 else -1)
+        delta = math.sin(diff)
+        self.a = (self.a + self.g * delta) % (2*math.pi)
+        self.b = (self.b - self.g * delta) % (2*math.pi)
+        drift = abs((self.a - self.b + math.pi) % (2*math.pi) - math.pi)
+        return drift, abs(delta)
+
+# ---------------- DCLR Optimiser -------------
+class DCLR(torch.optim.Optimizer):
+    def __init__(self, params, lr=5e-3, beta=0.9, gamma=0.999, eps=1e-8, cg=0.05):
+        super().__init__(params, dict(lr=lr,beta=beta,gamma=gamma,eps=eps,cg=cg))
+    @torch.no_grad()
+    def step(self, closure=None):
+        tot = 0.0
+        for g in self.param_groups:
+            lr, beta, gamma, eps, c = g["lr"], g["beta"], g["gamma"], g["eps"], g["cg"]
+            for p in g["params"]:
+                if p.grad is None: continue
+                st = self.state[p]
+                if not st:
+                    st["m"] = torch.zeros_like(p)
+                    st["v"] = torch.zeros_like(p)
+                    st["coh"] = torch.zeros_like(p)
+                m,v,h = st["m"],st["v"],st["coh"]; grad=p.grad
+                m.mul_(beta).add_(grad,alpha=1-beta)
+                v.mul_(gamma).addcmul_(grad,grad,value=1-gamma)
+                d=grad-m; h.mul_(0.9).add_(d.abs(),alpha=0.1)
+                lr_eff=lr/(1+c*h)
+                step=lr_eff*m/(v.sqrt()+eps); p.add_(-step)
+                tot += (step*step).sum().item()
+        return None, tot
+
+# ---------------- Agent Field ----------------
+class Agents(torch.nn.Module):
+    def __init__(self, n=256, box=10.0, r0=0.15):
+        super().__init__()
+        self.n=n; self.box=box; self.r0=r0
+        pos=(torch.rand(n,2)-0.5)*box
+        vel=torch.zeros(n,2)
+        self.pos=torch.nn.Parameter(pos); self.vel=torch.nn.Parameter(vel)
+    def forward(self):
+        n=self.n; pos=self.pos
+        diff=pos.unsqueeze(1)-pos.unsqueeze(0)
+        dist=torch.clamp(diff.norm(dim=-1),1e-6)
+        mask=(dist<self.r0) & (~torch.eye(n,dtype=torch.bool,device=pos.device))
+        rep=(diff/(dist.unsqueeze(-1)+1e-6))*mask.unsqueeze(-1)
+        rep=rep.sum(dim=1)
+        spring=-0.001*pos
+        acc=0.05*rep + spring
+        return acc
+
+# ---------------- Runner ---------------------
+def train(mode="RFT", steps=500, n=256, r0=0.165, log_path="stage2_agents.jsonl"):
+    set_seed(1234)
+    tm=Telemetry(log_path); orb=Orbital()
+    dev="cuda" if torch.cuda.is_available() else "cpu"
+    A=Agents(n=n, r0=r0).to(dev)
+    opt = DCLR(A.parameters(), lr=5e-3) if mode=="RFT" else torch.optim.SGD(A.parameters(), lr=5e-3)
+    collisions=0
+    for s in range(1, steps+1):
+        drift,flux=orb.step()
+        opt.zero_grad(set_to_none=True)
+        acc=A()
+        loss=(acc**2).mean()
+        loss.backward()
+        if isinstance(opt,DCLR): _,J=opt.step()
+        else: opt.step(); J=0.0
+        with torch.no_grad():
+            A.pos.add_(A.vel*0.0)
+        d=torch.cdist(A.pos, A.pos)
+        c=(d< A.r0*0.99).sum().item()-n
+        collisions = max(0, c)
+        tm.emit(mode=mode, step=s, drift=round(drift,3), flux=round(flux,3),
+                E_ret=0.992, coh=0.999, loss=round(float(loss.item()),4),
+                collisions=collisions)
+    tm.close()
+    return f"Stage 2 complete. Telemetry saved to {log_path}"