Create models/action_agent.py

models/action_agent.py

@@ -0,0 +1,15 @@
+import torch
+import torch.nn as nn
+
+class ActionAgent(nn.Module):
+    def __init__(self, config):
+        super(ActionAgent, self).__init__()
+        self.fc_layers = nn.Sequential(
+            nn.Linear(config["decision_output_size"], 128),
+            nn.ReLU(),
+            nn.Linear(128, config["action_output_size"])
+        )
+
+    def forward(self, x):
+        x = self.fc_layers(x)
+        return x