Machine Learning Trainer

Training environment to make a walk model for the hexapod
generated code that will be checked

ml/model.py

@@ -0,0 +1,49 @@
+import torch
+import torch.nn as nn
+from torch.distributions import Normal
+
+
+class ActorCritic(nn.Module):
+    def __init__(self, obs_dim: int, action_dim: int, hidden_sizes=(256, 256)):
+        super().__init__()
+        self.backbone = nn.Sequential(
+            nn.Linear(obs_dim, hidden_sizes[0]),
+            nn.ReLU(),
+            nn.Linear(hidden_sizes[0], hidden_sizes[1]),
+            nn.ReLU(),
+        )
+
+        self.mean_head = nn.Linear(hidden_sizes[1], action_dim)
+        self.value_head = nn.Linear(hidden_sizes[1], 1)
+        self.log_std = nn.Parameter(torch.zeros(action_dim, dtype=torch.float32))
+
+    def forward(self, obs: torch.Tensor):
+        x = self.backbone(obs)
+        mean = self.mean_head(x)
+        std = self.log_std.exp()
+        value = self.value_head(x).squeeze(-1)
+        return mean, std, value
+
+    def get_action(self, obs: torch.Tensor):
+        mean, std, value = self.forward(obs)
+        dist = Normal(mean, std)
+        action = dist.sample()
+        log_prob = dist.log_prob(action).sum(-1)
+        return action, log_prob, value
+
+    def evaluate_actions(self, obs: torch.Tensor, actions: torch.Tensor):
+        mean, std, value = self.forward(obs)
+        dist = Normal(mean, std)
+        log_prob = dist.log_prob(actions).sum(-1)
+        entropy = dist.entropy().sum(-1)
+        return value, log_prob, entropy
+
+    def save(self, path: str):
+        torch.save(self.state_dict(), path)
+
+    @classmethod
+    def load(cls, path: str, obs_dim: int, action_dim: int, hidden_sizes=(256, 256)):
+        model = cls(obs_dim, action_dim, hidden_sizes)
+        model.load_state_dict(torch.load(path, map_location=torch.device("cpu")))
+        model.eval()
+        return model