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: Initializes a new `TFPyEnvironment`.
  
  Args:
    environment: Environment to interact with, implementing
      `py_environment.PyEnvironment`.  Or a `callable` that returns an
      environment of this form.  If a `callable` is provided and
      `thread_isolation` is provided, the callable is executed in the
      dedicated thread.
    check_dims: Whether should check batch dimensions of actions in `step`.
    isolation: If this value is `False` (default), interactions with the
      environment will occur within whatever thread the methods of the
      `TFPyEnvironment` are run from.  For example, in TF graph mode, methods
      like `step` are called from multiple threads created by the TensorFlow
      engine; calls to step the environment are guaranteed to be sequential,
      but not from the same thread.  This creates problems for environments
      that are not thread-safe.  Using isolation ensures not only that a
      dedicated thread (or thread-pool) is used to interact with the
      environment, but also that interaction with the environment happens in a
      serialized manner.  If `isolation == True`, a dedicated thread is
      created for interactions with the environment.  If `isolation` is an
      instance of `multiprocessing.pool.Pool` (this includes instances of
      `multiprocessing.pool.ThreadPool`, nee `multiprocessing.dummy.Pool` and
      `multiprocessing.Pool`, then this pool is used to interact with the
      environment.  **NOTE** If using `isolation` with a
      `BatchedPyEnvironment`, ensure you create the `BatchedPyEnvironment`
      with `multithreading=False`, since otherwise the multithreading in that
      wrapper reverses the effects of this one.
  
  Raises:
    TypeError: If `environment` is not an instance of
      `py_environment.PyEnvironment` or subclasses, or is a callable that does
      not return an instance of `PyEnvironment`.
    TypeError: If `isolation` is not `True`, `False`, or an instance of
      `multiprocessing.pool.Pool`.
docstr_tok: ["Initializes","a","new","`","TFPyEnvironment","`",".","Args",":","environment",":","Environment","to","interact","with",",","implementing","`","py_environment.PyEnvironment","`",".","Or","a","`","callable","`","that","returns","an","environment","of","this","form",".","If","a","`","callable","`","is","provided","and","`","thread_isolation","`","is","provided",",","the","callable","is","executed","in","the","dedicated","thread",".","check_dims",":","Whether","should","check","batch","dimensions","of","actions","in","`","step","`",".","isolation",":","If","this","value","is","`","False","`","(","default",")",",","interactions","with","the","environment","will","occur","within","whatever","thread","the","methods","of","the","`","TFPyEnvironment","`","are","run","from",".","For","example",",","in","TF","graph","mode",",","methods","like","`","step","`","are","called","from","multiple","threads","created","by","the","TensorFlow","engine",";","calls","to","step","the","environment","are","guaranteed","to","be","sequential",",","but","not","from","the","same","thread",".","This","creates","problems","for","environments","that","are","not","thread-safe",".","Using","isolation","ensures","not","only","that","a","dedicated","thread","(","or","thread-pool",")","is","used","to","interact","with","the","environment",",","but","also","that","interaction","with","the","environment","happens","in","a","serialized","manner",".","If","`","isolation","==","True","`",",","a","dedicated","thread","is","created","for","interactions","with","the","environment",".","If","`","isolation","`","is","an","instance","of","`","multiprocessing.pool.Pool","`","(","this","includes","instances","of","`","multiprocessing.pool.ThreadPool","`",",","nee","`","multiprocessing.dummy.Pool","`","and","`","multiprocessing.Pool","`",",","then","this","pool","is","used","to","interact","with","the","environment",".","*","*","NOTE","*","*","If","using","`","isolation","`","with","a","`","BatchedPyEnvironment","`",",","ensure","you","create","the","`","BatchedPyEnvironment","`","with","`","multithreading=False","`",",","since","otherwise","the","multithreading","in","that","wrapper","reverses","the","effects","of","this","one",".","Raises",":","TypeError",":","If","`","environment","`","is","not","an","instance","of","`","py_environment.PyEnvironment","`","or","subclasses",",","or","is","a","callable","that","does","not","return","an","instance","of","`","PyEnvironment","`",".","TypeError",":","If","`","isolation","`","is","not","`","True","`",",","`","False","`",",","or","an","instance","of","`","multiprocessing.pool.Pool","`","."]
code: def __init__(
      self,
      environment: py_environment.PyEnvironment,
      check_dims: bool = False,
      isolation: bool = False,
  ):
      """Initializes a new `TFPyEnvironment`.
  
      Args:
        environment: Environment to interact with, implementing
          `py_environment.PyEnvironment`.  Or a `callable` that returns an
          environment of this form.  If a `callable` is provided and
          `thread_isolation` is provided, the callable is executed in the
          dedicated thread.
        check_dims: Whether should check batch dimensions of actions in `step`.
        isolation: If this value is `False` (default), interactions with the
          environment will occur within whatever thread the methods of the
          `TFPyEnvironment` are run from.  For example, in TF graph mode, methods
          like `step` are called from multiple threads created by the TensorFlow
          engine; calls to step the environment are guaranteed to be sequential,
          but not from the same thread.  This creates problems for environments
          that are not thread-safe.  Using isolation ensures not only that a
          dedicated thread (or thread-pool) is used to interact with the
          environment, but also that interaction with the environment happens in a
          serialized manner.  If `isolation == True`, a dedicated thread is
          created for interactions with the environment.  If `isolation` is an
          instance of `multiprocessing.pool.Pool` (this includes instances of
          `multiprocessing.pool.ThreadPool`, nee `multiprocessing.dummy.Pool` and
          `multiprocessing.Pool`, then this pool is used to interact with the
          environment.  **NOTE** If using `isolation` with a
          `BatchedPyEnvironment`, ensure you create the `BatchedPyEnvironment`
          with `multithreading=False`, since otherwise the multithreading in that
          wrapper reverses the effects of this one.
  
      Raises:
        TypeError: If `environment` is not an instance of
          `py_environment.PyEnvironment` or subclasses, or is a callable that does
          not return an instance of `PyEnvironment`.
        TypeError: If `isolation` is not `True`, `False`, or an instance of
          `multiprocessing.pool.Pool`.
      """
      if not isolation:
          self._pool = None
      elif isinstance(isolation, pool.Pool):
          self._pool = isolation
      elif isolation:
          self._pool = pool.ThreadPool(1)
      else:
          raise TypeError(
              "isolation should be True, False, or an instance of "
              "a multiprocessing Pool or ThreadPool.  Saw: {}".format(
                  isolation
              )
          )
  
      if callable(environment):
          environment = self._execute(environment)
      if not isinstance(environment, py_environment.PyEnvironment):
          raise TypeError(
              "Environment should implement py_environment.PyEnvironment"
          )
  
      if not environment.batched:
          # If executing in an isolated thread, do not enable multiprocessing for
          # this environment.
          environment = batched_py_environment.BatchedPyEnvironment(
              [environment], multithreading=not self._pool
          )
      self._env = environment
      self._check_dims = check_dims
  
      if isolation and getattr(self._env, "_parallel_execution", None):
          logging.warning(
              "Wrapped environment is executing in parallel.  "
              "Perhaps it is a BatchedPyEnvironment with multithreading=True, "
              "or it is a ParallelPyEnvironment.  This conflicts with the "
              "`isolation` arg passed to TFPyEnvironment: interactions with the "
              "wrapped environment are no longer guaranteed to happen in a common "
              "thread.  Environment: %s",
              (self._env,),
          )
  
      action_spec = tensor_spec.from_spec(self._env.action_spec())
      time_step_spec = tensor_spec.from_spec(self._env.time_step_spec())
      batch_size = self._env.batch_size if self._env.batch_size else 1
      self._render_shape = None
  
      super(TFPyEnvironment, self).__init__(
          time_step_spec, action_spec, batch_size
      )
  
      # Gather all the dtypes and shapes of the elements in time_step.
      self._time_step_dtypes = [
          s.dtype for s in tf.nest.flatten(self.time_step_spec())
      ]
  
      self._time_step = None
      self._lock = threading.Lock()
  
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 1 to create a new TFPyEnvironment object from inherited base classes: tf_environment with example: obj = TFPyEnvironment(environment, check_dims, isolation)