Deep-Learning-Atari/BaseROM.py at python · shamrickus/Deep-Learning-Atari · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
import numpy
from skimage.transform import rescale


class BaseROM:
    name = None

    def process_image(self, image):
        return self.rescale(self.gray_scale(self.crop(image)))

    def rescale(self, image):
        result_size = 84
        assert image.shape[0] == image.shape[1]
        return rescale(image, result_size / image.shape[0])

    def crop(self, image):
        return image

    def gray_scale(self, image):
        return numpy.dot(image[..., :3], [0.299, 0.587, 0.114])

    def actions(self):
        return range(0, 17)


class Breakout(BaseROM):
    def __init__(self):
        self.name = "Breakout.a26"

    def crop(self, image):
        return image[35:-15, ...]

    def actions(self):
        return range(0, 5)


class Beamrider(BaseROM):
    def __init__(self):
        self.name = "Beamrider.a26"


class Enduro(BaseROM):
    def __init__(self):
        self.name = "Enduro.a26"


class Pong(BaseROM):
    def __init__(self):
        self.name = "Pong.a26"


class Qbert(BaseROM):
    def __init__(self):
        self.name = "Qbert.a26"


class Seaquest(BaseROM):
    def __init__(self):
        self.name = "Seaquest.a26"


class SpaceInvaders(BaseROM):
    def __init__(self):
        self.name = "SpaceInvaders.a26"