atari_utils.py

from collections import deque
from gym import Wrapper, ObservationWrapper, ActionWrapper
from gym.spaces.box import Box
import numpy as np
import cv2


def _process_frame42(frame):
    reshaped_screen = np.reshape(frame, [210, 160, 3]).astype(np.float32).mean(2)
    resized_screen = cv2.resize(reshaped_screen, (84, 110))
    x_t = resized_screen[18:102, :]
    x_t = cv2.resize(x_t, (42, 42))
    x_t *= (1.0 / 255.0)
    x_t = np.reshape(x_t, [42, 42, 1])
    return x_t


class AtariRescale42x42Env(ObservationWrapper):
    def __init__(self, env=None):
        super(AtariRescale42x42Env, self).__init__(env)
        self.observation_space = Box(0, 255, [42, 42, 1])

    def _observation(self, observation):
        return _process_frame42(observation)


def _process_frame84(frame):
    img = np.reshape(frame, [210, 160, 3]).astype(np.float32)
    img = img[:, :, 0] * 0.2126 + img[:, :, 1] * 0.0722 + img[:, :, 2] * 0.7152
    resized_screen = cv2.resize(img, (84, 110),  interpolation=cv2.INTER_LINEAR)
    x_t = resized_screen[18:102, :]
    x_t /= 255.0
#    x_t -= 0.5
    x_t = np.reshape(x_t, [84, 84, 1])
    return x_t


class AtariRescale84x84Env(ObservationWrapper):
    def __init__(self, env=None):
        super(AtariRescale84x84Env, self).__init__(env)
        self.observation_space = Box(0, 255, [84, 84, 1])

    def _observation(self, observation):
        return _process_frame84(observation)


class RandomizedResetEnv(Wrapper):
    def __init__(self, env, no_op_max=7):
        super(RandomizedResetEnv, self).__init__(env)
        self._no_op_max = no_op_max

    def _reset(self):
        ob = self.env.reset()
        action = 0

        # randomize initial state
        if self._no_op_max > 0:
            no_op = np.random.randint(0, self._no_op_max + 1)
            for _ in range(no_op):
                ob, _, _, _ = self.env.step(action)
        return ob


class OneLiveResetEnv(Wrapper):
    def _step(self, action):
        lives = self.env.unwrapped.ale.lives()
        observation, reward, done, info = self.env.step(action)
        if lives != self.env.unwrapped.ale.lives():
            done = True
        return observation, reward, done, info


class UnstuckPolicyEnv(ActionWrapper):
    actions = deque(maxlen=30)

    def _action(self, action):
        if self.actions.count(action) == 30:
            action = 1
        self.actions.append(action)
        return action

    def _reverse_action(self, action):
        return action


class ObservationBuffer(Wrapper):
    def __init__(self, env, buffer_size=4):
        super(ObservationBuffer, self).__init__(env)
        self.buffer_size = buffer_size
        self.buffer = deque(maxlen=self.buffer_size)
        assert len(self.env.observation_space.shape) == 3
        self._shape = list(self.env.observation_space.shape)
        self._num_channels = self._shape[2]
        self._shape[2] *= self.buffer_size
        self.observation_space = Box(-0.5, 0.5, self._shape)

    def _step(self, action):
        observation, reward, done, info = self.env.step(action)
        self.buffer.append(observation)
        return np.concatenate(self.buffer, axis=2), reward, done, info

    def _reset(self):
        obs = self.env.reset()
        for _ in range(self.buffer_size):
            self.buffer.append(obs)
        return np.concatenate(self.buffer, axis=2)

#    def _render(self, mode='human', close=False):
#        if mode == "rgb_array":
#            return self.buffer[-1]
#        return self.env.render(mode, close)