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

Track down autograd error

Browse source 
"Trying to backward through the graph a second time" -> added a comment
This commit is contained in:
Antonin Raffin 2019-11-27 17:29:47 +01:00
parent fe67a98711
commit fbe29a7298
3 changed files with 17 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
torchy_baselines/common/distributions.py
View file

@ -288,7 +288,7 @@ class StateDependentNoiseDistribution(Distribution):
if self.full_std:
return std
# 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):
"""

2
torchy_baselines/sac/policies.py
View file

@ -33,7 +33,7 @@ class Actor(BaseNetwork):
if self.use_sde:
# TODO: check for the learn_features
self.action_dist = StateDependentNoiseDistribution(action_dim, full_std=full_std, use_expln=False,
learn_features=False, squash_output=True)
learn_features=True, squash_output=True)
self.mu, self.log_std = self.action_dist.proba_distribution_net(latent_dim=net_arch[-1],
log_std_init=log_std_init)
else:

23
torchy_baselines/sac/sac.py
View file

@ -170,11 +170,13 @@ class SAC(BaseRLModel):
obs, action_batch, next_obs, done, reward = replay_data
# TODO: check if there is another way to fix pytorch complain
# if we don't sample the weights again
# (Trying to backward through the graph a second time)
if self.use_sde:
self.actor.reset_noise()
# Two options: retain_graph=True in the actor_loss.backward()
# or sample again the noise matrix
# otherwise the intermediate step `std = th.exp(log_std)`
# is lost and we cannot backpropagate through again
# if self.use_sde:
# self.actor.reset_noise()
# Action by the current actor for the sampled state
action_pi, log_prob = self.actor.action_log_prob(obs)
log_prob = log_prob.reshape(-1, 1)
@ -196,10 +198,10 @@ class SAC(BaseRLModel):
ent_coef_loss.backward()
self.ent_coef_optimizer.step()
# Select action according to policy
next_action, next_log_prob = self.actor.action_log_prob(next_obs)
with th.no_grad():
# Select action according to policy
next_action, next_log_prob = self.actor.action_log_prob(next_obs)
# Compute the target Q value
target_q1, target_q2 = self.critic_target(next_obs, next_action)
target_q = th.min(target_q1, target_q2)
@ -227,7 +229,10 @@ class SAC(BaseRLModel):
# Optimize the actor
self.actor.optimizer.zero_grad()
actor_loss.backward()
# Cf comment above, otherwise pytorch raises an error
# ("Trying to backward through the graph a second time")
retain_graph = True if self.use_sde and gradient_steps > 1 else False
actor_loss.backward(retain_graph=retain_graph)
self.actor.optimizer.step()
# Update target networks
@ -239,6 +244,8 @@ class SAC(BaseRLModel):
logger.logkv("ent_coef", ent_coef.item())
logger.logkv("actor_loss", actor_loss.item())
logger.logkv("critic_loss", critic_loss.item())
if ent_coef_loss is not None:
logger.logkv("ent_coef_loss", ent_coef_loss.item())
def learn(self, total_timesteps, callback=None, log_interval=4,
eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="SAC",