fix(log_prob): fixed bugs when calculating log_prob while using funct…

…ion `squash_action`. (#34) thanks to @BlueFisher
StepNeverStop · Aug 30, 2021 · 078b523 · 078b523
1 parent 23f910c
commit 078b523
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 12 deletions.
diff --git a/rls/algorithms/single/sac.py b/rls/algorithms/single/sac.py
@@ -137,10 +137,11 @@ def _train_continuous(self, BATCH):
         if self.is_continuous:
             target_mu, target_log_std = self.actor(
                 BATCH.obs_, begin_mask=BATCH.begin_mask)   # [T, B, A]
-            dist = td.Normal(target_mu, target_log_std.exp())
+            dist = td.Independent(
+                td.Normal(target_mu, target_log_std.exp()), 1)
             target_pi = dist.sample()   # [T, B, A]
             target_pi, target_log_pi = squash_action(
-                target_pi, dist.log_prob(target_pi))   # [T, B, A], [T, B, 1]
+                target_pi, dist.log_prob(target_pi).unsqueeze(-1))   # [T, B, A], [T, B, 1]
         else:
             target_logits = self.actor(
                 BATCH.obs_, begin_mask=BATCH.begin_mask)  # [T, B, A]
@@ -170,10 +171,10 @@ def _train_continuous(self, BATCH):
         if self.is_continuous:
             mu, log_std = self.actor(
                 BATCH.obs, begin_mask=BATCH.begin_mask)  # [T, B, A]
-            dist = td.Normal(mu, log_std.exp())
+            dist = td.Independent(td.Normal(mu, log_std.exp()), 1)
             pi = dist.rsample()  # [T, B, A]
             pi, log_pi = squash_action(
-                pi, dist.log_prob(pi))   # [T, B, A], [T, B, 1]
+                pi, dist.log_prob(pi).unsqueeze(-1))   # [T, B, A], [T, B, 1]
             entropy = dist.entropy().mean()  # 1
         else:
             logits = self.actor(

diff --git a/rls/algorithms/single/sac_v.py b/rls/algorithms/single/sac_v.py
@@ -145,10 +145,10 @@ def _train_continuous(self, BATCH):
         if self.is_continuous:
             mu, log_std = self.actor(
                 BATCH.obs, begin_mask=BATCH.begin_mask)  # [T, B, A]
-            dist = td.Normal(mu, log_std.exp())
+            dist = td.Independent(td.Normal(mu, log_std.exp()), 1)
             pi = dist.rsample()  # [T, B, A]
             pi, log_pi = squash_action(
-                pi, dist.log_prob(pi))   # [T, B, A], [T, B, 1]
+                pi, dist.log_prob(pi).unsqueeze(-1))   # [T, B, A], [T, B, 1]
         else:
             logits = self.actor(
                 BATCH.obs, begin_mask=BATCH.begin_mask)  # [T, B, A]
@@ -189,10 +189,10 @@ def _train_continuous(self, BATCH):
         if self.is_continuous:
             mu, log_std = self.actor(
                 BATCH.obs, begin_mask=BATCH.begin_mask)  # [T, B, A]
-            dist = td.Normal(mu, log_std.exp())
+            dist = td.Independent(td.Normal(mu, log_std.exp()), 1)
             pi = dist.rsample()  # [T, B, A]
             pi, log_pi = squash_action(
-                pi, dist.log_prob(pi))   # [T, B, A], [T, B, 1]
+                pi, dist.log_prob(pi).unsqueeze(-1))   # [T, B, A], [T, B, 1]
             entropy = dist.entropy().mean()  # 1
         else:
             logits = self.actor(

diff --git a/rls/algorithms/single/tac.py b/rls/algorithms/single/tac.py
@@ -119,10 +119,11 @@ def _train(self, BATCH):
         if self.is_continuous:
             target_mu, target_log_std = self.actor(
                 BATCH.obs_, begin_mask=BATCH.begin_mask)  # [T, B, A]
-            dist = td.Normal(target_mu, target_log_std.exp())
+            dist = td.Independent(
+                td.Normal(target_mu, target_log_std.exp()), 1)
             target_pi = dist.sample()   # [T, B, A]
             target_pi, target_log_pi = squash_action(
-                target_pi, dist.log_prob(target_pi), is_independent=False)  # [T, B, A]
+                target_pi, dist.log_prob(target_pi).unsqueeze(-1), is_independent=False)  # [T, B, A]
             target_log_pi = tsallis_entropy_log_q(
                 target_log_pi, self.entropic_index)   # [T, B, 1]
         else:
@@ -160,10 +161,10 @@ def _train(self, BATCH):
         if self.is_continuous:
             mu, log_std = self.actor(
                 BATCH.obs, begin_mask=BATCH.begin_mask)  # [T, B, A]
-            dist = td.Normal(mu, log_std.exp())
+            dist = td.Independent(td.Normal(mu, log_std.exp()), 1)
             pi = dist.rsample()  # [T, B, A]
             pi, log_pi = squash_action(pi, dist.log_prob(
-                pi), is_independent=False)  # [T, B, A]
+                pi).unsqueeze(-1), is_independent=False)  # [T, B, A]
             log_pi = tsallis_entropy_log_q(
                 log_pi, self.entropic_index)  # [T, B, 1]
             entropy = dist.entropy().mean()  # 1