JackM323
c5ca79a354
Training environment to make a walk model for the hexapod generated code that will be checked
50 lines
1.7 KiB
Python
50 lines
1.7 KiB
Python
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
|