# Path: ai_safety_gridworlds/environments/shared/rl/array_spec.py
# class ArraySpec(object):
# class BoundedArraySpec(ArraySpec):
#   def __init__(self, shape, dtype, name=None):
#   def shape(self):
#   def dtype(self):
#   def name(self):
#   def __repr__(self):
#   def __eq__(self, other):
#   def __ne__(self, other):
#   def _fail_validation(self, message, *args):
#   def validate(self, value):
#   def generate_value(self):
#   def __init__(self, shape, dtype, minimum, maximum, name=None):
#   def minimum(self):
#   def maximum(self):
#   def __repr__(self):
#   def __eq__(self, other):
#   def validate(self, value):
#   def generate_value(self):
# 
# Path: ai_safety_gridworlds/environments/shared/rl/environment.py
# class TimeStep(collections.namedtuple(
#     'TimeStep', ['step_type', 'reward', 'discount', 'observation'])):
# class StepType(enum.IntEnum):
# class Base(object):
#   def first(self):
#   def mid(self):
#   def last(self):
#   def first(self):
#   def mid(self):
#   def last(self):
#   def reset(self):
#   def step(self, action):
#   def observation_spec(self):
#   def action_spec(self):
#   def close(self):
#   def __enter__(self):
#   def __exit__(self, unused_exception_type, unused_exc_value, unused_traceback):
# def restart(observation):
# def transition(reward, observation, discount=1.0):
# def termination(reward, observation):
# def truncation(reward, observation, discount=1.0):
#   FIRST = 0
#   MID = 1
#   LAST = 2


from ai_safety_gridworlds.environments.shared.rl import array_spec as specs
from ai_safety_gridworlds.environments.shared.rl import environment
from six.moves import zip
import numpy as np
import six
          or empty lists.
    """
    # Save important constructor arguments.
    self._game_factory = game_factory
    self._default_reward = default_reward
    self._observation_distiller = observation_distiller
    self._max_iterations = max_iterations

    # These slots comprise an Environment's internal state. They are:
    self._state = None              # Current Environment game step state.
    self._current_game = None       # Current pycolab game instance.
    self._game_over = None          # Whether the instance's game has ended.
    self._last_observations = None  # Last observation received from the game.
    self._last_reward = None        # Last reward, if any, or default reward.
    self._last_discount = None      # Last discount factor from the game.

    # Attempt to distill our action spec.
    self._valid_actions, self._action_size = self._compute_action_spec(
        discrete_actions, continuous_actions)

    # With this, we're ready to compute our own observation spec. This is done
    # by starting a new episode, inspecting the observations returned in the
    # first step, then closing the episode and resetting internal variables
    # to a default value.
    self._observation_spec = self._compute_observation_spec()

  def reset(self):
    """Start a new episode."""
    # Build a new game and retrieve its first set of state/reward/discount.
    self._current_game = self._game_factory()

# Path: ai_safety_gridworlds/environments/shared/rl/array_spec.py
# class ArraySpec(object):
# class BoundedArraySpec(ArraySpec):
#   def __init__(self, shape, dtype, name=None):
#   def shape(self):
#   def dtype(self):
#   def name(self):
#   def __repr__(self):
#   def __eq__(self, other):
#   def __ne__(self, other):
#   def _fail_validation(self, message, *args):
#   def validate(self, value):
#   def generate_value(self):
#   def __init__(self, shape, dtype, minimum, maximum, name=None):
#   def minimum(self):
#   def maximum(self):
#   def __repr__(self):
#   def __eq__(self, other):
#   def validate(self, value):
#   def generate_value(self):
# 
# Path: ai_safety_gridworlds/environments/shared/rl/environment.py
# class TimeStep(collections.namedtuple(
#     'TimeStep', ['step_type', 'reward', 'discount', 'observation'])):
# class StepType(enum.IntEnum):
# class Base(object):
#   def first(self):
#   def mid(self):
#   def last(self):
#   def first(self):
#   def mid(self):
#   def last(self):
#   def reset(self):
#   def step(self, action):
#   def observation_spec(self):
#   def action_spec(self):
#   def close(self):
#   def __enter__(self):
#   def __exit__(self, unused_exception_type, unused_exc_value, unused_traceback):
# def restart(observation):
# def transition(reward, observation, discount=1.0):
# def termination(reward, observation):
# def truncation(reward, observation, discount=1.0):
#   FIRST = 0
#   MID = 1
#   LAST = 2


# Path: ai_safety_gridworlds/environments/shared/rl/pycolab_interface.py
from ai_safety_gridworlds.environments.shared.rl import array_spec as specs
from ai_safety_gridworlds.environments.shared.rl import environment
from six.moves import zip
import numpy as np
import six

          or empty lists.
    """
    # Save important constructor arguments.
    self._game_factory = game_factory
    self._default_reward = default_reward
    self._observation_distiller = observation_distiller
    self._max_iterations = max_iterations

    # These slots comprise an Environment's internal state. They are:
    self._state = None              # Current Environment game step state.
    self._current_game = None       # Current pycolab game instance.
    self._game_over = None          # Whether the instance's game has ended.
    self._last_observations = None  # Last observation received from the game.
    self._last_reward = None        # Last reward, if any, or default reward.
    self._last_discount = None      # Last discount factor from the game.

    # Attempt to distill our action spec.
    self._valid_actions, self._action_size = self._compute_action_spec(
        discrete_actions, continuous_actions)

    # With this, we're ready to compute our own observation spec. This is done
    # by starting a new episode, inspecting the observations returned in the
    # first step, then closing the episode and resetting internal variables
    # to a default value.
    self._observation_spec = self._compute_observation_spec()

  def reset(self):
    """Start a new episode."""
    # Build a new game and retrieve its first set of state/reward/discount.
    self._current_game = self._game_factory()

    self._state = environment.StepType.FIRST