make action probabilities a numpy array (#10122)

rllib/offline/off_policy_estimator.py

@@ -1,10 +1,13 @@
 from collections import namedtuple
 import logging
 
+import numpy as np
+
 from ray.rllib.policy.sample_batch import MultiAgentBatch, SampleBatch
 from ray.rllib.policy import Policy
 from ray.rllib.utils.annotations import DeveloperAPI
 from ray.rllib.offline.io_context import IOContext
+from ray.rllib.utils.numpy import convert_to_numpy
 from ray.rllib.utils.typing import TensorType, SampleBatchType
 from typing import List
 
@@ -53,7 +56,7 @@ class OffPolicyEstimator:
         raise NotImplementedError
 
     @DeveloperAPI
-    def action_prob(self, batch: SampleBatchType) -> TensorType:
+    def action_prob(self, batch: SampleBatchType) -> np.ndarray:
         """Returns the probs for the batch actions for the current policy."""
 
         num_state_inputs = 0
@@ -61,13 +64,13 @@ class OffPolicyEstimator:
             if k.startswith("state_in_"):
                 num_state_inputs += 1
         state_keys = ["state_in_{}".format(i) for i in range(num_state_inputs)]
-        log_likelihoods = self.policy.compute_log_likelihoods(
+        log_likelihoods: TensorType = self.policy.compute_log_likelihoods(
             actions=batch[SampleBatch.ACTIONS],
             obs_batch=batch[SampleBatch.CUR_OBS],
             state_batches=[batch[k] for k in state_keys],
             prev_action_batch=batch.data.get(SampleBatch.PREV_ACTIONS),
             prev_reward_batch=batch.data.get(SampleBatch.PREV_REWARDS))
-        return log_likelihoods
+        return convert_to_numpy(log_likelihoods)
 
     @DeveloperAPI
     def process(self, batch: SampleBatchType):