class_docstr: Base class to build policies based on reward predictions.
class_docstr_tok: ["Base","class","to","build","policies","based","on","reward","predictions","."]
docstr: Returns an action distribution based on predicted rewards.
  
  Sub-classes are expected to implement this method.
  
  Args:
    mask: A 2-D tensor of binary masks shaped as [batch_size, num_of_arms].
    predicted_rewards: A 2-D tensor of predicted rewards shaped as
      [batch_size, num_of_arms].
  
  Returns:
    A tuple of a `tfp.distributions.Distribution` and a tensor of policy types
    shaped as [batch_size].
docstr_tok: ["Returns","an","action","distribution","based","on","predicted","rewards",".","Sub-classes","are","expected","to","implement","this","method",".","Args",":","mask",":","A","2-D","tensor","of","binary","masks","shaped","as","[","batch_size",",","num_of_arms","]",".","predicted_rewards",":","A","2-D","tensor","of","predicted","rewards","shaped","as","[","batch_size",",","num_of_arms","]",".","Returns",":","A","tuple","of","a","`","tfp.distributions.Distribution","`","and","a","tensor","of","policy","types","shaped","as","[","batch_size","]","."]
code: def _action_distribution(
      self,
      mask: Optional[types.Tensor],
      predicted_rewards: types.Tensor,
  ) -> Tuple[tfp.distributions.Distribution, types.Tensor]:
      """Returns an action distribution based on predicted rewards.
  
      Sub-classes are expected to implement this method.
  
      Args:
        mask: A 2-D tensor of binary masks shaped as [batch_size, num_of_arms].
        predicted_rewards: A 2-D tensor of predicted rewards shaped as
          [batch_size, num_of_arms].
  
      Returns:
        A tuple of a `tfp.distributions.Distribution` and a tensor of policy types
        shaped as [batch_size].
      """
      pass
  
imports: import abc
  from typing import Iterable, Optional, Text, Tuple
  import gin
  import six
  import tensorflow
  import tensorflow_probability
  from tf_agents.bandits.networks import heteroscedastic_q_network
  from tf_agents.bandits.policies import constraints
  from tf_agents.bandits.specs import utils
  from tf_agents.policies import tf_policy
  from tf_agents.policies import utils
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import policy_step
  from tf_agents.trajectories import time_step
  from tf_agents.typing import types
label_desc: Use image node_id 4 for calling the RewardPredictionBasePolicy obj's underlying member method code with example usage: obj._action_distribution(mask, predicted_rewards) without return types