class GaussianHead(nn.Module):
"""Mean network plus a state-independent learned log standard deviation."""
def __init__(self, obs_dim, act_dim, hidden):
super().__init__()
self.mean = nn.Sequential(nn.Linear(obs_dim, hidden), nn.Tanh(),
nn.Linear(hidden, act_dim))
self.log_std = nn.Parameter(torch.zeros(act_dim))
def forward(self, obs):
return self.mean(obs), self.log_std.exp()
class GaussianPolicy(d2l.ActorCritic):
"""The same interface over a Normal instead of a softmax; nothing that
consumes the interface changes."""
def __init__(self, obs_dim, act_dim, hidden=64, lr=1e-2):
super().__init__(GaussianHead(obs_dim, act_dim, hidden),
nn.Sequential(nn.Linear(obs_dim, hidden), nn.Tanh(),
nn.Linear(hidden, 1)), lr)
def log_prob(self, obs, act):
mean, std = self.policy(obs)
return torch.distributions.Normal(mean, std).log_prob(act).sum(-1)
def act(self, obs, rng):
with torch.no_grad():
mean, std = self.policy(torch.as_tensor(obs))
return mean.numpy() + std.numpy() * rng.standard_normal(
mean.shape, dtype=np.float32)
def act_greedy(self, obs, rng=None):
with torch.no_grad():
return self.policy(torch.as_tensor(obs))[0].numpy()