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: Renders the environment.
  
  Note for compatibility this will convert the image to uint8.
  
  Args:
    mode: One of ['rgb_array', 'human']. Renders to an numpy array, or brings
      up a window where the environment can be visualized.
  
  Returns:
    A Tensor of shape [width, height, 3] denoting an RGB image if mode is
    `rgb_array`. Otherwise return nothing and render directly to a display
    window.
  Raises:
    NotImplementedError: If the environment does not support rendering.
docstr_tok: ["Renders","the","environment",".","Note","for","compatibility","this","will","convert","the","image","to","uint8",".","Args",":","mode",":","One","of","[","'rgb_array","'",",","'human","'","]",".","Renders","to","an","numpy","array",",","or","brings","up","a","window","where","the","environment","can","be","visualized",".","Returns",":","A","Tensor","of","shape","[","width",",","height",",","3","]","denoting","an","RGB","image","if","mode","is","`","rgb_array","`",".","Otherwise","return","nothing","and","render","directly","to","a","display","window",".","Raises",":","NotImplementedError",":","If","the","environment","does","not","support","rendering","."]
code: def render(
      self, mode: Text = "rgb_array"
  ) -> Optional[types.NestedTensor]:
      """Renders the environment.
  
      Note for compatibility this will convert the image to uint8.
  
      Args:
        mode: One of ['rgb_array', 'human']. Renders to an numpy array, or brings
          up a window where the environment can be visualized.
  
      Returns:
        A Tensor of shape [width, height, 3] denoting an RGB image if mode is
        `rgb_array`. Otherwise return nothing and render directly to a display
        window.
      Raises:
        NotImplementedError: If the environment does not support rendering.
      """
  
      if not self._render_shape:
          # Make sure the environment has been initialized.
          self.current_time_step()
          img = self._env.render("rgb_array")
          self._render_shape = img.shape
  
      def _render(mode):
          """Pywrapper fn to the environments render."""
          # Mode might be passed down as bytes or ndarray.
          # If so, convert to a str first.
          if isinstance(mode, np.ndarray):
              mode = str(mode)
          if isinstance(mode, bytes):
              mode = mode.decode("utf-8")
          if mode == "rgb_array":
              img = self._env.render(mode)
              img = img.astype(np.uint8, copy=False)
              return img
          elif mode == "human":
              # Generate mock img to keep outputs the same.
              self._env.render(mode)
              return np.zeros(self._render_shape, dtype=np.uint8)
  
      img = tf.numpy_function(
          lambda mode: self._execute(_render, mode),
          [mode],
          [tf.uint8],
          name="render_py_func",
      )
  
      if not tf.executing_eagerly():
          # Extract from list returned from np_function.
          img = img[0]
          img.set_shape(tf.TensorShape(self._render_shape))
      return img
  
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 9 for calling the TFPyEnvironment obj's underlying member method code with example usage: obj.render(mode) and returns: img, img, np