Fix bc_loss calculation in ReBRAC

takuseno · Aug 25, 2024 · 87751da · 87751da
1 parent 3433de5
commit 87751da
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/d3rlpy/algos/qlearning/torch/rebrac_impl.py b/d3rlpy/algos/qlearning/torch/rebrac_impl.py
@@ -56,10 +56,12 @@ def compute_actor_loss(
             reduction="min",
         )
         lam = 1 / (q_t.abs().mean()).detach()
-        bc_loss = ((batch.actions - action.squashed_mu) ** 2).mean()
+        bc_loss = ((batch.actions - action.squashed_mu) ** 2).sum(
+            dim=1, keepdim=True
+        )
         return TD3PlusBCActorLoss(
-            actor_loss=lam * -q_t.mean() + self._actor_beta * bc_loss,
-            bc_loss=bc_loss,
+            actor_loss=(lam * -q_t + self._actor_beta * bc_loss).mean(),
+            bc_loss=bc_loss.mean(),
         )
 
     def compute_target(self, batch: TorchMiniBatch) -> torch.Tensor: