fix categorical action in cql loss

pytorch · Nov 8, 2023 · 53ef411 · 53ef411
1 parent 9af9f99
commit 53ef411
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/torchrl/objectives/cql.py b/torchrl/objectives/cql.py
@@ -1152,6 +1152,9 @@ def cql_loss(self, tensordict):
 
         logsumexp = torch.logsumexp(qvalues, dim=-1, keepdim=True)
         if self.action_space == "categorical":
+            if current_action.shape != qvalues.shape:
+                # unsqueeze the action if it lacks on trailing singleton dim
+                current_action = current_action.unsqueeze(-1)
             q_a = qvalues.gather(-1, current_action)
         else:
             q_a = (qvalues * current_action).sum(dim=-1, keepdim=True)