diff --git a/rllib/offline/off_policy_estimator.py b/rllib/offline/off_policy_estimator.py index fff235f82..ea163c1cd 100644 --- a/rllib/offline/off_policy_estimator.py +++ b/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):