class SquashedGaussianPolicy(nn.Module):
"""A state-dependent Gaussian squashed through a = c tanh(u)."""
def __init__(self, obs_dim, act_dim, hidden=64):
super().__init__()
self.trunk = nn.Sequential(nn.Linear(obs_dim, hidden), nn.ReLU(),
nn.Linear(hidden, hidden), nn.ReLU())
self.mu = nn.Linear(hidden, act_dim)
self.log_std = nn.Linear(hidden, act_dim)
def forward(self, obs):
h = self.trunk(obs)
return self.mu(h), self.log_std(h).clamp(-5, 2).exp()
def log_prob(self, u, mean, std):
"""log pi at a = c tanh(u), from the pre-squash u the sampler keeps."""
logdet = 2 * (np.log(2) - u - nn.functional.softplus(-2 * u))
return (torch.distributions.Normal(mean, std).log_prob(u)
- logdet - np.log(c)).sum(-1)
def sample(self, obs):
"""A reparameterized action and its log-probability, differentiable."""
mean, std = self(obs)
u = mean + std * torch.randn_like(std)
return c * torch.tanh(u), self.log_prob(u, mean, std)
def act(self, obs, rng):
with torch.no_grad():
mean, std = self(torch.as_tensor(obs))
u = mean.numpy() + std.numpy() * rng.standard_normal(
mean.shape, dtype=np.float32)
return c * np.tanh(u)
def act_greedy(self, obs, rng=None):
with torch.no_grad():
return c * np.tanh(self(torch.as_tensor(obs))[0].numpy())