class_docstr: Base class to build policies based on reward predictions.
class_docstr_tok: ["Base","class","to","build","policies","based","on","reward","predictions","."]
docstr: Implementation of `action`.
  
  Args:
    time_step: A `TimeStep` tuple corresponding to `time_step_spec()`.
    policy_state: A Tensor, or a nested dict, list or tuple of Tensors
      representing the previous policy_state.
    seed: Seed to use if action performs sampling (optional).
  
  Returns:
    A `PolicyStep` named tuple containing:
      `action`: An action Tensor matching the `action_spec`.
      `state`: A policy state tensor to be fed into the next call to action.
      `info`: Optional side information such as action log probabilities.
docstr_tok: ["Implementation","of","`","action","`",".","Args",":","time_step",":","A","`","TimeStep","`","tuple","corresponding","to","`","time_step_spec","(",")","`",".","policy_state",":","A","Tensor",",","or","a","nested","dict",",","list","or","tuple","of","Tensors","representing","the","previous","policy_state",".","seed",":","Seed","to","use","if","action","performs","sampling","(","optional",")",".","Returns",":","A","`","PolicyStep","`","named","tuple","containing",":","`","action","`",":","An","action","Tensor","matching","the","`","action_spec","`",".","`","state","`",":","A","policy","state","tensor","to","be","fed","into","the","next","call","to","action",".","`","info","`",":","Optional","side","information","such","as","action","log","probabilities","."]
code: def _action(
      self,
      time_step: ts.TimeStep,
      policy_state: types.NestedTensor,
      seed: Optional[types.Seed] = None,
  ) -> policy_step.PolicyStep:
      """Implementation of `action`.
  
      Args:
        time_step: A `TimeStep` tuple corresponding to `time_step_spec()`.
        policy_state: A Tensor, or a nested dict, list or tuple of Tensors
          representing the previous policy_state.
        seed: Seed to use if action performs sampling (optional).
  
      Returns:
        A `PolicyStep` named tuple containing:
          `action`: An action Tensor matching the `action_spec`.
          `state`: A policy state tensor to be fed into the next call to action.
          `info`: Optional side information such as action log probabilities.
      """
      step = super(RewardPredictionBasePolicy, self)._action(
          time_step, policy_state, seed
      )
      return self._maybe_save_chosen_arm_features(
          time_step, step.action, step
      )
  
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 6 for calling the RewardPredictionBasePolicy obj's underlying member method code with example usage: obj._action(time_step, policy_state, seed) and returns: self