fix bug on prediction

takuseno · Jul 15, 2024 · 2520211 · 2520211
1 parent 3b15aa5
commit 2520211
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/d3rlpy/algos/qlearning/torch/ddpg_impl.py b/d3rlpy/algos/qlearning/torch/ddpg_impl.py
@@ -242,7 +242,7 @@ def compute_target(self, batch: TorchMiniBatch) -> torch.Tensor:
         pass
 
     def inner_predict_best_action(self, x: TorchObservation) -> torch.Tensor:
-        return torch.argmax(self._modules.policy(x).mean).unsqueeze(0)
+        return torch.argmax(self._modules.policy(x).probs).unsqueeze(0)
 
     @abstractmethod
     def inner_sample_action(self, x: TorchObservation) -> torch.Tensor: