class_docstr: Abstract base class for Python Policies.
  
  The `action(time_step, policy_state)` method returns a PolicyStep named tuple
  containing the following nested arrays:
    `action`: The action to be applied on the environment.
    `state`: The state of the policy (E.g. RNN state) to be fed into the next
      call to action.
    `info`: Optional side information such as action log probabilities.
  
  For stateful policies, e.g. those containing RNNs, an initial policy state can
  be obtained through a call to `get_initial_state()`.
  
  Example of simple use in Python:
  
    py_env = PyEnvironment()
    policy = PyPolicy()
  
    time_step = py_env.reset()
    policy_state = policy.get_initial_state()
  
    acc_reward = 0
    while not time_step.is_last():
      action_step = policy.action(time_step, policy_state)
      policy_state = action_step.state
      time_step = py_env.step(action_step.action)
      acc_reward += time_step.reward
class_docstr_tok: ["Abstract","base","class","for","Python","Policies",".","The","`","action","(","time_step",",","policy_state",")","`","method","returns","a","PolicyStep","named","tuple","containing","the","following","nested","arrays",":","`","action","`",":","The","action","to","be","applied","on","the","environment",".","`","state","`",":","The","state","of","the","policy","(","E.g",".","RNN","state",")","to","be","fed","into","the","next","call","to","action",".","`","info","`",":","Optional","side","information","such","as","action","log","probabilities",".","For","stateful","policies",",","e.g",".","those","containing","RNNs",",","an","initial","policy","state","can","be","obtained","through","a","call","to","`","get_initial_state","(",")","`",".","Example","of","simple","use","in","Python",":","py_env","=","PyEnvironment","(",")","policy","=","PyPolicy","(",")","time_step","=","py_env.reset","(",")","policy_state","=","policy.get_initial_state","(",")","acc_reward","=","0","while","not","time_step.is_last","(",")",":","action_step","=","policy.action","(","time_step",",","policy_state",")","policy_state","=","action_step.state","time_step","=","py_env.step","(","action_step.action",")","acc_reward","+=","time_step.reward"]
docstr: Initialization of PyPolicy class.
  
  Args:
    time_step_spec: A `TimeStep` ArraySpec of the expected time_steps. Usually
      provided by the user to the subclass.
    action_spec: A nest of BoundedArraySpec representing the actions. Usually
      provided by the user to the subclass.
    policy_state_spec: A nest of ArraySpec representing the policy state.
      Provided by the subclass, not directly by the user.
    info_spec: A nest of ArraySpec representing the policy info. Provided by
      the subclass, not directly by the user.
    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.
docstr_tok: ["Initialization","of","PyPolicy","class",".","Args",":","time_step_spec",":","A","`","TimeStep","`","ArraySpec","of","the","expected","time_steps",".","Usually","provided","by","the","user","to","the","subclass",".","action_spec",":","A","nest","of","BoundedArraySpec","representing","the","actions",".","Usually","provided","by","the","user","to","the","subclass",".","policy_state_spec",":","A","nest","of","ArraySpec","representing","the","policy","state",".","Provided","by","the","subclass",",","not","directly","by","the","user",".","info_spec",":","A","nest","of","ArraySpec","representing","the","policy","info",".","Provided","by","the","subclass",",","not","directly","by","the","user",".","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","."]
code: def __init__(
      self,
      time_step_spec: ts.TimeStep,
      action_spec: types.NestedArraySpec,
      policy_state_spec: types.NestedArraySpec = (),
      info_spec: types.NestedArraySpec = (),
      observation_and_action_constraint_splitter: Optional[
          types.Splitter
      ] = None,
  ):
      """Initialization of PyPolicy class.
  
      Args:
        time_step_spec: A `TimeStep` ArraySpec of the expected time_steps. Usually
          provided by the user to the subclass.
        action_spec: A nest of BoundedArraySpec representing the actions. Usually
          provided by the user to the subclass.
        policy_state_spec: A nest of ArraySpec representing the policy state.
          Provided by the subclass, not directly by the user.
        info_spec: A nest of ArraySpec representing the policy info. Provided by
          the subclass, not directly by the user.
        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.
      """
      common.tf_agents_gauge.get_cell("TFAPolicy").set(True)
      common.assert_members_are_not_overridden(
          base_cls=PyPolicy, instance=self
      )
      self._time_step_spec = tensor_spec.to_array_spec(time_step_spec)
      self._action_spec = tensor_spec.to_array_spec(action_spec)
      # TODO(kbanoop): rename policy_state to state.
      self._policy_state_spec = tensor_spec.to_array_spec(
          policy_state_spec
      )
      self._info_spec = tensor_spec.to_array_spec(info_spec)
      self._setup_specs()
      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
  import numpy
  import six
  import tensorflow
  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
label_desc: Use image node_id 1 to create a new PyPolicy object from inherited base classes: object with example: obj = PyPolicy(time_step_spec, action_spec, policy_state_spec, info_spec, observation_and_action_constraint_splitter)