diff --git a/ml/model.py b/ml/model.py deleted file mode 100644 index cb3aabd..0000000 --- a/ml/model.py +++ /dev/null @@ -1,49 +0,0 @@ -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