1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216 | """State Dependent Exploration (gSDE) distribution."""
from typing import Optional, TypeVar
import torch
from torch import nn
from torch.distributions import Normal
from lit_sb3.distributions.base import Distribution, sum_independent_dims
from lit_sb3.distributions.tanh_bijector import TanhBijector
SelfStateDependentNoiseDistribution = TypeVar(
"SelfStateDependentNoiseDistribution", bound="StateDependentNoiseDistribution"
)
class StateDependentNoiseDistribution(Distribution):
"""Distribution class for using generalized State Dependent Exploration (gSDE).
Paper: https://arxiv.org/abs/2005.05719.
It is used to create the noise exploration matrix and
compute the log probability of an action with that noise.
Args:
action_dim: Dimension of the action space.
full_std: Whether to use (n_features x n_actions) parameters
for the std instead of only (n_features,)
use_expln: Use ``expln()`` function instead of ``exp()`` to ensure
a positive standard deviation (cf paper). It allows to keep variance
above zero and prevent it from growing too fast. In practice, ``exp()`` is usually enough.
squash_output: Whether to squash the output using a tanh function,
this ensures bounds are satisfied.
learn_features: Whether to learn features for gSDE or not.
This will enable gradients to be backpropagated through the features
``latent_sde`` in the code.
epsilon: small value to avoid NaN due to numerical imprecision.
"""
bijector: TanhBijector | None
latent_sde_dim: int | None
weights_dist: Normal
_latent_sde: torch.Tensor
exploration_mat: torch.Tensor
exploration_matrices: torch.Tensor
def __init__(
self,
action_dim: int,
full_std: bool = True,
use_expln: bool = False,
squash_output: bool = False,
learn_features: bool = False,
epsilon: float = 1e-6,
) -> None:
super().__init__()
self.action_dim = action_dim
self.latent_sde_dim = None
self.mean_actions = None
self.log_std = None
self.use_expln = use_expln
self.full_std = full_std
self.epsilon = epsilon
self.learn_features = learn_features
if squash_output:
self.bijector = TanhBijector(epsilon)
else:
self.bijector = None
def get_std(self, log_std: torch.Tensor) -> torch.Tensor:
"""Get the standard deviation from the learned parameter.
(log of it by default). This ensures that the std is positive.
Args:
log_std: Log standard deviation of the action distribution.
"""
if self.use_expln:
# From gSDE paper, it allows to keep variance
# above zero and prevent it from growing too fast
below_threshold = torch.exp(log_std) * (log_std <= 0)
# Avoid NaN: zeros values that are below zero
safe_log_std = log_std * (log_std > 0) + self.epsilon
above_threshold = (torch.log1p(safe_log_std) + 1.0) * (log_std > 0)
std = below_threshold + above_threshold
else:
# Use normal exponential
std = torch.exp(log_std)
if self.full_std:
return std
assert self.latent_sde_dim is not None
# Reduce the number of parameters:
return torch.ones(self.latent_sde_dim, self.action_dim).to(log_std.device) * std
def sample_weights(self, log_std: torch.Tensor, batch_size: int = 1) -> None:
"""Sample weights for the noise exploration matrix,.
using a centered Gaussian distribution.
Args:
log_std: Log standard deviation of the action distribution.
batch_size: Number of elements to sample.
"""
std = self.get_std(log_std)
self.weights_dist = Normal(torch.zeros_like(std), std)
# Reparametrization trick to pass gradients
self.exploration_mat = self.weights_dist.rsample()
# Pre-compute matrices in case of parallel exploration
self.exploration_matrices = self.weights_dist.rsample((batch_size,))
def proba_distribution_net(
self, latent_dim: int, log_std_init: float = -2.0, latent_sde_dim: int | None = None
) -> tuple[nn.Module, nn.Parameter]:
"""Create the layers and parameter that represent the distribution.
one output will be the deterministic action, the other parameter will be the
standard deviation of the distribution that control the weights of the noise matrix.
Args:
latent_dim: Dimension of the last layer of the policy (before the action layer)
log_std_init: Initial value for the log standard deviation
latent_sde_dim: Dimension of the last layer of the features extractor
for gSDE. By default, it is shared with the policy network.
"""
# Network for the deterministic action, it represents the mean of the distribution
mean_actions_net = nn.Linear(latent_dim, self.action_dim)
# When we learn features for the noise, the feature dimension
# can be different between the policy and the noise network
self.latent_sde_dim = latent_dim if latent_sde_dim is None else latent_sde_dim
# Reduce the number of parameters if needed
log_std = (
torch.ones(self.latent_sde_dim, self.action_dim) if self.full_std else torch.ones(self.latent_sde_dim, 1)
)
# Transform it to a parameter so it can be optimized
log_std = nn.Parameter(log_std * log_std_init, requires_grad=True)
# Sample an exploration matrix
self.sample_weights(log_std)
return mean_actions_net, log_std
def proba_distribution(
self: SelfStateDependentNoiseDistribution,
mean_actions: torch.Tensor,
log_std: torch.Tensor,
latent_sde: torch.Tensor,
) -> SelfStateDependentNoiseDistribution:
"""Create the distribution given its parameters (mean, std).
Args:
mean_actions: Mean of the action distribution.
log_std: Log standard deviation of the action distribution.
latent_sde: Latent variable for the SDE exploration noise.
"""
# Stop gradient if we don't want to influence the features
self._latent_sde = latent_sde if self.learn_features else latent_sde.detach()
variance = torch.mm(self._latent_sde**2, self.get_std(log_std) ** 2)
self.distribution = Normal(mean_actions, torch.sqrt(variance + self.epsilon))
return self
def log_prob(self, actions: torch.Tensor) -> torch.Tensor: # ty:ignore[invalid-method-override]
gaussian_actions = self.bijector.inverse(actions) if self.bijector is not None else actions
# log likelihood for a gaussian
log_prob = self.distribution.log_prob(gaussian_actions) # ty:ignore[unresolved-attribute]
# Sum along action dim
log_prob = sum_independent_dims(log_prob)
if self.bijector is not None:
# Squash correction (from original SAC implementation)
log_prob -= torch.sum(self.bijector.log_prob_correction(gaussian_actions), dim=1)
return log_prob
def entropy(self) -> torch.Tensor | None:
if self.bijector is not None:
# No analytical form,
# entropy needs to be estimated using -log_prob.mean()
return None
return sum_independent_dims(self.distribution.entropy()) # ty:ignore[unresolved-attribute]
def sample(self) -> torch.Tensor:
noise = self.get_noise(self._latent_sde)
actions = self.distribution.mean + noise # ty:ignore[unresolved-attribute]
if self.bijector is not None:
return self.bijector.forward(actions)
return actions
def mode(self) -> torch.Tensor:
actions = self.distribution.mean # ty:ignore[unresolved-attribute]
if self.bijector is not None:
return self.bijector.forward(actions)
return actions
def get_noise(self, latent_sde: torch.Tensor) -> torch.Tensor:
latent_sde = latent_sde if self.learn_features else latent_sde.detach()
# Default case: only one exploration matrix
if len(latent_sde) == 1 or len(latent_sde) != len(self.exploration_matrices):
return torch.mm(latent_sde, self.exploration_mat)
# Use batch matrix multiplication for efficient computation
# (batch_size, n_features) -> (batch_size, 1, n_features)
latent_sde = latent_sde.unsqueeze(dim=1)
# (batch_size, 1, n_actions)
noise = torch.bmm(latent_sde, self.exploration_matrices)
return noise.squeeze(dim=1)
def actions_from_params(
self, mean_actions: torch.Tensor, log_std: torch.Tensor, latent_sde: torch.Tensor, deterministic: bool = False
) -> torch.Tensor:
# Update the proba distribution
self.proba_distribution(mean_actions, log_std, latent_sde)
return self.get_actions(deterministic=deterministic)
def log_prob_from_params(
self, mean_actions: torch.Tensor, log_std: torch.Tensor, latent_sde: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
actions = self.actions_from_params(mean_actions, log_std, latent_sde)
log_prob = self.log_prob(actions)
return actions, log_prob
|