Create policy.py

policy.py

@@ -0,0 +1,35 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+
+
+class PsiPolicy(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.model = nn.Sequential(
+            nn.Linear(dim, 32),
+            nn.ReLU(),
+            nn.Linear(32, dim)
+        )
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class PolicyController:
+    def __init__(self, system):
+        self.system = system
+        self.policy = PsiPolicy(system.n)
+        self.opt = optim.Adam(self.policy.parameters(), lr=0.001)
+
+    def policy_step(self):
+        state = torch.tensor(self.system.Xi, dtype=torch.float32)
+        Psi = self.policy(state)
+        Psi_np = Psi.detach().numpy()
+
+        self.system.step(Psi_np)
+
+        loss = torch.norm(Psi)
+        self.opt.zero_grad()
+        loss.backward()
+        self.opt.step()