wrapper_env.py

import deepmind_lab

class WrapperEnv(object):
  """A gym-like wrapper environment for DeepMind Lab.
  (Work in Progress)

  Attributes:
      env: The corresponding DeepMind Lab environment.
      length: Maximum number of frames

  Args:
      env (deepmind_lab.Lab): DeepMind Lab environment.

  Written by:
    Kevin COSTA (kcosta42)
    Michael TRAZZI (mtrazzi)
  """
  def __init__(self, env, length):
    self.env = env
    self.length = length
    self.reset()

  def step(self, action):
    done = not self.env.is_running()
    if (done):
      self.reset()
    obs = self.env.observations()
    reward = self.env.step(action, num_steps=1)

    return obs['RGB_INTERLEAVED'], reward, done, self.env.num_steps()

  def reset(self):
    self.env.reset()
    obs = self.env.observations()
    return obs['RGB_INTERLEAVED']