play_atari.py

# -*- coding: utf-8 -*-

import gym
import numpy as np
import cv2
import six

class Controller(object):
    def __init__(self):
        self.keymap = list(range(49, 59))

    def __call__(self, pause=30):
        key = cv2.waitKey(pause)
        if key == -1:
            return 0
        return 1 + self.keymap.index(key)

env = gym.make('BattleZone-v0')
ctl = Controller()

for i in six.moves.range(1):
    obs = env.reset()
    done = False
    R = 0
    t = 0

    while not done:
        env.render()
        #frame = env.render(mode='rgb_array')
        #frame = cv2.resize(frame, (frame.shape[1] * 4, frame.shape[0] * 4))
        #cv2.imshow('game', frame)
        #act = ctl(pause=10)

        act = env.action_space.sample()
        obs, r, done, _ = env.step(act)

        R += r
        t += 1