saymrwulf/stable-baselines3
1
0
Fork 0
mirror of https://github.com/saymrwulf/stable-baselines3.git synced 2026-09-16 22:20:26 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Bug fix in SAC with constant ent coeff + try batch sde matrices

Browse source
This commit is contained in:
Antonin RAFFIN 2019-12-01 13:11:13 +01:00
parent fbe29a7298
commit 879191b26a
3 changed files with 22 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
torchy_baselines/common/distributions.py
View file

@ -290,7 +290,7 @@ class StateDependentNoiseDistribution(Distribution):
# Reduce the number of parameters: # Reduce the number of parameters:
return th.ones(self.latent_sde_dim, self.action_dim).to(log_std.device) * std return th.ones(self.latent_sde_dim, self.action_dim).to(log_std.device) * std
def sample_weights(self, log_std): def sample_weights(self, log_std, batch_size=1):
""" """
Sample weights for the noise exploration matrix, Sample weights for the noise exploration matrix,
using a centered gaussian distribution. using a centered gaussian distribution.
@ -300,6 +300,7 @@ class StateDependentNoiseDistribution(Distribution):
std = self.get_std(log_std) std = self.get_std(log_std)
self.weights_dist = Normal(th.zeros_like(std), std) self.weights_dist = Normal(th.zeros_like(std), std)
self.exploration_mat = self.weights_dist.rsample() self.exploration_mat = self.weights_dist.rsample()
self.exploration_matrices = self.weights_dist.rsample((batch_size,))
def proba_distribution_net(self, latent_dim, log_std_init=-2.0, latent_sde_dim=None): def proba_distribution_net(self, latent_dim, log_std_init=-2.0, latent_sde_dim=None):
""" """
@ -354,7 +355,13 @@ class StateDependentNoiseDistribution(Distribution):
def get_noise(self, latent_sde): def get_noise(self, latent_sde):
latent_sde = latent_sde if self.learn_features else latent_sde.detach() latent_sde = latent_sde if self.learn_features else latent_sde.detach()
if len(latent_sde) != len(self.exploration_matrices):
return th.mm(latent_sde, self.exploration_mat) return th.mm(latent_sde, self.exploration_mat)
# (batch_size, n_features) -> (batch_size, 1, n_features)
latent_sde = latent_sde.unsqueeze(1)
# (batch_size, 1, n_actions)
noise = th.bmm(latent_sde, self.exploration_matrices)
return noise.squeeze(1)
def sample(self, latent_sde): def sample(self, latent_sde):
noise = self.get_noise(latent_sde) noise = self.get_noise(latent_sde)

6
torchy_baselines/sac/policies.py
View file

@ -53,11 +53,13 @@ class Actor(BaseNetwork):
""" """
return self.action_dist.get_std(self.log_std) return self.action_dist.get_std(self.log_std)
def reset_noise(self): def reset_noise(self, batch_size=1):
""" """
Sample new weights for the exploration matrix, when using SDE. Sample new weights for the exploration matrix, when using SDE.
:param batch_size: (int)
""" """
self.action_dist.sample_weights(self.log_std) self.action_dist.sample_weights(self.log_std, batch_size=batch_size)
def get_action_dist_params(self, obs): def get_action_dist_params(self, obs):
latent = self.actor_net(obs) latent = self.actor_net(obs)

15
torchy_baselines/sac/sac.py
View file

@ -87,6 +87,7 @@ class SAC(BaseRLModel):
self.n_episodes_rollout = n_episodes_rollout self.n_episodes_rollout = n_episodes_rollout
self.action_noise = action_noise self.action_noise = action_noise
self.gamma = gamma self.gamma = gamma
self.ent_coef_optimizer = None
if _init_setup_model: if _init_setup_model:
self._setup_model() self._setup_model()
@ -124,7 +125,7 @@ class SAC(BaseRLModel):
# Force conversion to float # Force conversion to float
# this will throw an error if a malformed string (different from 'auto') # this will throw an error if a malformed string (different from 'auto')
# is passed # is passed
self.ent_coef = float(self.ent_coef) self.ent_coef = th.tensor(float(self.ent_coef)).to(self.device)
self.replay_buffer = ReplayBuffer(self.buffer_size, obs_dim, action_dim, self.device) self.replay_buffer = ReplayBuffer(self.buffer_size, obs_dim, action_dim, self.device)
self.policy = self.policy(self.observation_space, self.action_space, learning_rate=self.learning_rate, self.policy = self.policy(self.observation_space, self.action_space, learning_rate=self.learning_rate,
@ -174,15 +175,15 @@ class SAC(BaseRLModel):
# or sample again the noise matrix # or sample again the noise matrix
# otherwise the intermediate step `std = th.exp(log_std)` # otherwise the intermediate step `std = th.exp(log_std)`
# is lost and we cannot backpropagate through again # is lost and we cannot backpropagate through again
# if self.use_sde: if self.use_sde:
# self.actor.reset_noise() self.actor.reset_noise(batch_size=batch_size)
# Action by the current actor for the sampled state # Action by the current actor for the sampled state
action_pi, log_prob = self.actor.action_log_prob(obs) action_pi, log_prob = self.actor.action_log_prob(obs)
log_prob = log_prob.reshape(-1, 1) log_prob = log_prob.reshape(-1, 1)
ent_coef_loss = None ent_coef_loss = None
if not isinstance(self.ent_coef, float): if self.ent_coef_optimizer is not None:
# Important: detach the variable from the graph # Important: detach the variable from the graph
# so we don't change it with other losses # so we don't change it with other losses
# see https://github.com/rail-berkeley/softlearning/issues/60 # see https://github.com/rail-berkeley/softlearning/issues/60
@ -200,6 +201,8 @@ class SAC(BaseRLModel):
with th.no_grad(): with th.no_grad():
if self.use_sde:
self.actor.reset_noise(batch_size=batch_size)
# Select action according to policy # Select action according to policy
next_action, next_log_prob = self.actor.action_log_prob(next_obs) next_action, next_log_prob = self.actor.action_log_prob(next_obs)
# Compute the target Q value # Compute the target Q value
@ -231,8 +234,8 @@ class SAC(BaseRLModel):
self.actor.optimizer.zero_grad() self.actor.optimizer.zero_grad()
# Cf comment above, otherwise pytorch raises an error # Cf comment above, otherwise pytorch raises an error
# ("Trying to backward through the graph a second time") # ("Trying to backward through the graph a second time")
retain_graph = True if self.use_sde and gradient_steps > 1 else False # retain_graph = True if self.use_sde and gradient_steps > 1 else False
actor_loss.backward(retain_graph=retain_graph) actor_loss.backward(retain_graph=False)
self.actor.optimizer.step() self.actor.optimizer.step()
# Update target networks # Update target networks