Fix returns_to_go shape in CalQL

takuseno · takuseno · commit a02a9d4f5639 · 2025-03-07T18:11:13.000+09:00
diff --git a/d3rlpy/algos/qlearning/torch/cal_ql_impl.py b/d3rlpy/algos/qlearning/torch/cal_ql_impl.py
@@ -18,4 +18,4 @@ def _compute_policy_is_values(
             value_obs=value_obs,
             returns_to_go=returns_to_go,
         )
-        return torch.maximum(values, returns_to_go), log_probs
+        return torch.maximum(values, returns_to_go.view(1, -1, 1)), log_probs

Original file line number	Diff line number	Diff line change
`@@ -18,4 +18,4 @@ def _compute_policy_is_values(`
`18`	`18`	`value_obs=value_obs,`
`19`	`19`	`returns_to_go=returns_to_go,`
`20`	`20`	`)`
`21`		`- return torch.maximum(values, returns_to_go), log_probs`
	`21`	`+ return torch.maximum(values, returns_to_go.view(1, -1, 1)), log_probs`