class_docstr: Abstract base class for TF Policies.
  
  The Policy represents a mapping from `time_steps` recieved from the
  environment to `actions` that can be applied to the environment.
  
  Agents expose two policies. A `policy` meant for deployment and evaluation,
  and a `collect_policy` for collecting data from the environment. The
  `collect_policy` is usually stochastic for exploring the environment better
  and may log auxilliary information such as log probabilities required for
  training as well. `Policy` objects can also be created directly by the users
  without using an `Agent`.
  
  The main methods of TFPolicy are:
  
  * `action`: Maps a `time_step` from the environment to an action.
  * `distribution`: Maps a `time_step` to a distribution over actions.
  * `get_initial_state`: Generates the initial state for stateful policies, e.g.
      RNN/LSTM policies.
  
  Example usage:
  
  ```
  env = SomeTFEnvironment()
  policy = TFRandomPolicy(env.time_step_spec(), env.action_spec())
  # Or policy = agent.policy or agent.collect_policy
  
  policy_state = policy.get_initial_state(env.batch_size)
  time_step = env.reset()
  
  while not time_step.is_last():
    policy_step = policy.action(time_step, policy_state)
    time_step = env.step(policy_step.action)
  
    policy_state = policy_step.state
    # policy_step.info may contain side info for logging, such as action log
    # probabilities.
  ```
  
  Policies can be saved to disk as SavedModels (see policy_saver.py and
  policy_loader.py) or as TF Checkpoints.
  
  A `PyTFEagerPolicy` can be used to wrap a `TFPolicy` so that it works with
  `PyEnvironment`s.
  
  
  **NOTE**: For API consistency, subclasses are not allowed to override public
  methods of `TFPolicy` class. Instead, they may implement the protected methods
  including `_get_initial_state`, `_action`, and `_distribution`. This
  public-calls-private convention allowed this base class to do things like
  properly add `spec` and shape checks, which provide users an easier experience
  when debugging their environments and networks.
  
  For researchers, and those developing new Policies, the `TFPolicy` base class
  constructor also accept a `validate_args` parameter.  If `False`, this
  disables all spec structure, dtype, and shape checks in the public methods of
  these classes.  It allows algorithm developers to iterate and try different
  input and output structures without worrying about overly restrictive
  requirements, or input and output states being in a certain format.  However,
  *disabling argument validation* can make it very hard to identify structural
  input or algorithmic errors; and should not be done for final, or
  production-ready, Policies.  In addition to having implementations that may
  disagree with specs, this mean that the resulting Policy may no longer
  interact well with other parts of TF-Agents.  Examples include impedance
  mismatches with Actor/Learner APIs, replay buffers, and the model export
  functionality in `PolicySaver.
docstr: Initialization of TFPolicy class.
  
  Args:
    time_step_spec: A `TimeStep` spec of the expected time_steps. Usually
      provided by the user to the subclass.
    action_spec: A nest of BoundedTensorSpec representing the actions. Usually
      provided by the user to the subclass.
    policy_state_spec: A nest of TensorSpec representing the policy_state.
      Provided by the subclass, not directly by the user.
    info_spec: A nest of TensorSpec representing the policy info. Provided by
      the subclass, not directly by the user.
    clip: Whether to clip actions to spec before returning them.  Default
      True. Most policy-based algorithms (PCL, PPO, REINFORCE) use unclipped
      continuous actions for training.
    emit_log_probability: Emit log-probabilities of actions, if supported. If
      True, policy_step.info will have CommonFields.LOG_PROBABILITY set.
      Please consult utility methods provided in policy_step for setting and
      retrieving these. When working with custom policies, either provide a
      dictionary info_spec or a namedtuple with the field 'log_probability'.
    automatic_state_reset:  If `True`, then `get_initial_policy_state` is used
      to clear state in `action()` and `distribution()` for for time steps
      where `time_step.is_first()`.
    observation_and_action_constraint_splitter: A function used to process
      observations with action constraints. These constraints can indicate,
      for example, a mask of valid/invalid actions for a given state of the
      environment. The function takes in a full observation and returns a
      tuple consisting of 1) the part of the observation intended as input to
      the network and 2) the constraint. An example
      `observation_and_action_constraint_splitter` could be as simple as: ```
      def observation_and_action_constraint_splitter(observation): return
      observation['network_input'], observation['constraint'] ``` *Note*: when
      using `observation_and_action_constraint_splitter`, make sure the
      provided `q_network` is compatible with the network-specific half of the
      output of the `observation_and_action_constraint_splitter`. In
      particular, `observation_and_action_constraint_splitter` will be called
      on the observation before passing to the network. If
      `observation_and_action_constraint_splitter` is None, action constraints
      are not applied.
    validate_args: Python bool.  Whether to verify inputs to, and outputs of,
      functions like `action` and `distribution` against spec structures,
      dtypes, and shapes.  Research code may prefer to set this value to
      `False` to allow iterating on input and output structures without being
      hamstrung by overly rigid checking (at the cost of harder-to-debug
      errors).  See also `TFAgent.validate_args`.
    name: A name for this module. Defaults to the class name.
code: def __init__(
      self,
      time_step_spec: ts.TimeStep,
      action_spec: types.NestedTensorSpec,
      policy_state_spec: types.NestedTensorSpec = (),
      info_spec: types.NestedTensorSpec = (),
      clip: bool = True,
      emit_log_probability: bool = False,
      automatic_state_reset: bool = True,
      observation_and_action_constraint_splitter: Optional[
          types.Splitter
      ] = None,
      validate_args: bool = True,
      name: Optional[Text] = None,
  ):
      """Initialization of TFPolicy class.
  
      Args:
        time_step_spec: A `TimeStep` spec of the expected time_steps. Usually
          provided by the user to the subclass.
        action_spec: A nest of BoundedTensorSpec representing the actions. Usually
          provided by the user to the subclass.
        policy_state_spec: A nest of TensorSpec representing the policy_state.
          Provided by the subclass, not directly by the user.
        info_spec: A nest of TensorSpec representing the policy info. Provided by
          the subclass, not directly by the user.
        clip: Whether to clip actions to spec before returning them.  Default
          True. Most policy-based algorithms (PCL, PPO, REINFORCE) use unclipped
          continuous actions for training.
        emit_log_probability: Emit log-probabilities of actions, if supported. If
          True, policy_step.info will have CommonFields.LOG_PROBABILITY set.
          Please consult utility methods provided in policy_step for setting and
          retrieving these. When working with custom policies, either provide a
          dictionary info_spec or a namedtuple with the field 'log_probability'.
        automatic_state_reset:  If `True`, then `get_initial_policy_state` is used
          to clear state in `action()` and `distribution()` for for time steps
          where `time_step.is_first()`.
        observation_and_action_constraint_splitter: A function used to process
          observations with action constraints. These constraints can indicate,
          for example, a mask of valid/invalid actions for a given state of the
          environment. The function takes in a full observation and returns a
          tuple consisting of 1) the part of the observation intended as input to
          the network and 2) the constraint. An example
          `observation_and_action_constraint_splitter` could be as simple as: ```
          def observation_and_action_constraint_splitter(observation): return
          observation['network_input'], observation['constraint'] ``` *Note*: when
          using `observation_and_action_constraint_splitter`, make sure the
          provided `q_network` is compatible with the network-specific half of the
          output of the `observation_and_action_constraint_splitter`. In
          particular, `observation_and_action_constraint_splitter` will be called
          on the observation before passing to the network. If
          `observation_and_action_constraint_splitter` is None, action constraints
          are not applied.
        validate_args: Python bool.  Whether to verify inputs to, and outputs of,
          functions like `action` and `distribution` against spec structures,
          dtypes, and shapes.  Research code may prefer to set this value to
          `False` to allow iterating on input and output structures without being
          hamstrung by overly rigid checking (at the cost of harder-to-debug
          errors).  See also `TFAgent.validate_args`.
        name: A name for this module. Defaults to the class name.
      """
      super(TFPolicy, self).__init__(name=name)
      common.check_tf1_allowed()
      common.tf_agents_gauge.get_cell("TFAPolicy").set(True)
      common.assert_members_are_not_overridden(
          base_cls=TFPolicy, instance=self
      )
      if not isinstance(time_step_spec, ts.TimeStep):
          raise ValueError(
              "The `time_step_spec` must be an instance of `TimeStep`, but is `{}`.".format(
                  type(time_step_spec)
              )
          )
  
      self._time_step_spec = tensor_spec.from_spec(time_step_spec)
      self._action_spec = tensor_spec.from_spec(action_spec)
      self._policy_state_spec = tensor_spec.from_spec(policy_state_spec)
      self._emit_log_probability = emit_log_probability
      self._validate_args = validate_args
  
      if emit_log_probability:
          log_probability_spec = tensor_spec.BoundedTensorSpec(
              shape=(),
              dtype=tf.float32,
              maximum=0,
              minimum=-float("inf"),
              name="log_probability",
          )
          log_probability_spec = tf.nest.map_structure(
              lambda _: log_probability_spec, action_spec
          )
          info_spec = policy_step.set_log_probability(
              info_spec, log_probability_spec
          )  # pytype: disable=wrong-arg-types
  
      self._info_spec = tensor_spec.from_spec(info_spec)
      self._setup_specs()
      self._clip = clip
      self._action_fn = common.function_in_tf1(reduce_retracing=False)(
          self._action
      )
      self._automatic_state_reset = automatic_state_reset
      self._observation_and_action_constraint_splitter = (
          observation_and_action_constraint_splitter
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import abc
  from typing import Optional, Sequence, Text
  import six
  import tensorflow
  import tensorflow_probability
  from tf_agents.distributions import reparameterized_sampling
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import policy_step
  from tf_agents.trajectories import time_step
  from tf_agents.trajectories import trajectory
  from tf_agents.typing import types
  from tf_agents.utils import common
  from tf_agents.utils import nest_utils
label_desc: Use image node_id 1 to create a new TFPolicy object from inherited base classes: tf with example: obj = TFPolicy(time_step_spec, action_spec, policy_state_spec, info_spec, clip, emit_log_probability, automatic_state_reset, observation_and_action_constraint_splitter, validate_args, name)