deleted leftover file
This commit is contained in:
-49
@@ -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
|
|
||||||
Reference in New Issue
Block a user