From ed943a635623d77fede8ea80599d03e06e8c147e Mon Sep 17 00:00:00 2001 From: Mike Clark Date: Fri, 15 Jun 2018 08:24:14 +0800 Subject: [PATCH] Use logits in CategoricalActorCriticNet A really minor change. Should we input logits instead of probs into `Categorical`? It will do less converting from logits and back, leading to increased stability during training (theoretically). --- deep_rl/network/network_heads.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/deep_rl/network/network_heads.py b/deep_rl/network/network_heads.py index 974f507..a9a0c4a 100644 --- a/deep_rl/network/network_heads.py +++ b/deep_rl/network/network_heads.py @@ -187,9 +187,9 @@ class CategoricalActorCriticNet(nn.Module, BaseNet): phi = self.network.phi_body(obs) phi_a = self.network.actor_body(phi) phi_v = self.network.critic_body(phi) - prob = F.softmax(self.network.fc_action(phi_a), dim=-1) + logits = self.network.fc_action(phi_a) v = self.network.fc_critic(phi_v) - dist = torch.distributions.Categorical(probs=prob) + dist = torch.distributions.Categorical(logits=logits) if action is None: action = dist.sample() log_prob = dist.log_prob(action).unsqueeze(-1)