class_docstr: Exposes a Python environment as an in-graph TF environment.
  
  This class supports Python environments that return nests of arrays as
  observations and accept nests of arrays as actions. The nest structure is
  reflected in the in-graph environment's observation and action structure.
  
  Implementation notes:
  
  * Since `tf.py_func` deals in lists of tensors, this class has some additional
    `tf.nest.flatten` and `tf.nest.pack_structure_as` calls.
  
  * This class currently cast rewards and discount to float32.
class_docstr_tok: ["Exposes","a","Python","environment","as","an","in-graph","TF","environment",".","This","class","supports","Python","environments","that","return","nests","of","arrays","as","observations","and","accept","nests","of","arrays","as","actions",".","The","nest","structure","is","reflected","in","the","in-graph","environment","'s","observation","and","action","structure",".","Implementation","notes",":","*","Since","`","tf.py_func","`","deals","in","lists","of","tensors",",","this","class","has","some","additional","`","tf.nest.flatten","`","and","`","tf.nest.pack_structure_as","`","calls",".","*","This","class","currently","cast","rewards","and","discount","to","float32","."]
docstr: Returns a TensorFlow op to step the environment.
  
  Args:
    actions: A Tensor, or a nested dict, list or tuple of Tensors
      corresponding to `action_spec()`.
  
  Returns:
    A `TimeStep` tuple of:
      step_type: A scalar int32 tensor representing the `StepType` value.
      reward: A float32 tensor representing the reward at this
        time_step.
      discount: A scalar float32 tensor representing the discount [0, 1].
      observation: A Tensor, or a nested dict, list or tuple of Tensors
        corresponding to `observation_spec()`.
  
  Raises:
    ValueError: If any of the actions are scalars or their major axis is known
    and is not equal to `self.batch_size`.
docstr_tok: ["Returns","a","TensorFlow","op","to","step","the","environment",".","Args",":","actions",":","A","Tensor",",","or","a","nested","dict",",","list","or","tuple","of","Tensors","corresponding","to","`","action_spec","(",")","`",".","Returns",":","A","`","TimeStep","`","tuple","of",":","step_type",":","A","scalar","int32","tensor","representing","the","`","StepType","`","value",".","reward",":","A","float32","tensor","representing","the","reward","at","this","time_step",".","discount",":","A","scalar","float32","tensor","representing","the","discount","[","0",",","1","]",".","observation",":","A","Tensor",",","or","a","nested","dict",",","list","or","tuple","of","Tensors","corresponding","to","`","observation_spec","(",")","`",".","Raises",":","ValueError",":","If","any","of","the","actions","are","scalars","or","their","major","axis","is","known","and","is","not","equal","to","`","self.batch_size","`","."]
code: def _step(self, actions):
      """Returns a TensorFlow op to step the environment.
  
      Args:
        actions: A Tensor, or a nested dict, list or tuple of Tensors
          corresponding to `action_spec()`.
  
      Returns:
        A `TimeStep` tuple of:
          step_type: A scalar int32 tensor representing the `StepType` value.
          reward: A float32 tensor representing the reward at this
            time_step.
          discount: A scalar float32 tensor representing the discount [0, 1].
          observation: A Tensor, or a nested dict, list or tuple of Tensors
            corresponding to `observation_spec()`.
  
      Raises:
        ValueError: If any of the actions are scalars or their major axis is known
        and is not equal to `self.batch_size`.
      """
  
      def _step_py(*flattened_actions):
          with _check_not_called_concurrently(self._lock):
              packed = tf.nest.pack_sequence_as(
                  structure=self.action_spec(),
                  flat_sequence=flattened_actions,
              )
              self._time_step = self._env.step(packed)
              return tf.nest.flatten(self._time_step)
  
      def _isolated_step_py(*flattened_actions):
          return self._execute(_step_py, *flattened_actions)
  
      with tf.name_scope("step"):
          flat_actions = [
              tf.identity(x) for x in tf.nest.flatten(actions)
          ]
          if self._check_dims:
              for action in flat_actions:
                  dim_value = tensor_shape.dimension_value(
                      action.shape[0]
                  )
                  if action.shape.rank == 0 or (
                      dim_value is not None
                      and dim_value != self.batch_size
                  ):
                      raise ValueError(
                          "Expected actions whose major dimension is batch_size (%d), "
                          "but saw action with shape %s:\n   %s"
                          % (self.batch_size, action.shape, action)
                      )
          outputs = tf.numpy_function(
              _isolated_step_py,
              flat_actions,
              self._time_step_dtypes,
              name="step_py_func",
          )
          return self._time_step_from_numpy_function_outputs(outputs)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import contextlib
  from multiprocessing import pool
  import threading
  from typing import Any, Optional, Text
  from absl import logging
  import gin
  import numpy
  import tensorflow
  from tf_agents.environments import batched_py_environment
  from tf_agents.environments import py_environment
  from tf_agents.environments import tf_environment
  from tf_agents.specs import tensor_spec
  from tf_agents.typing import types
  from tensorflow.python.framework import tensor_shape
label_desc: Use image node_id 8 for calling the TFPyEnvironment obj's underlying member method code with example usage: obj._step(actions) and returns: self, self, tf