Deep-Learning-Atari/ALEWrapper.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
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
from ale_python_interface import ALEInterface

import BaseROM


class ALEWrapper:
    ale: ALEInterface = None

    def __init__(self):
        self.ale = ALEInterface()

    def set_seed(self, seed):
        self.ale.setInt(b'random_seed', seed)

    def set_display(self, display):
        self.ale.setBool(b'display_screen', display)

    def set_recording(self, sub_dir):
        dir = f"recording/{sub_dir}"
        self.ale.setString(b'record_screen_dir', str.encode(dir))
        self.ale.setInt(b'fragsize', 64)

    def set_repeat_action(self, probability=0.):
        self.ale.setFloat(b'repeat_action_probability', probability)

    def load_rom(self, rom: BaseROM):
        self.ale.loadROM(str.encode("roms/"+rom.name))

    def game_over(self):
        return self.ale.game_over()

    def get_minimal_actions(self):
        return list(self.ale.getMinimalActionSet())

    def act(self, action):
        return self.ale.act(action)

    def restore_state(self, state):
        self.ale.restoreState(state)

    def restore_system_state(self, state):
        self.ale.restoreSystemState(state)

    def copy_state(self):
        return self.ale.cloneState()

    def copy_system_state(self):
        return self.ale.cloneSystemState()

    def get_screen_dim(self):
        return self.ale.getScreenDims()

    def get_frame(self):
        return self.ale.getScreenRGB()

    def reset(self):
        self.ale.reset_game()

    def get_hashable_state(self, state):
        state.flags.writeable = False
        return state.ravel().data

    def map_action(self, action):
        if action == 0:
            return "noop"
        elif action == 1:
            return "fire"
        elif action == 2:
            return "up"
        elif action == 3:
            return "right"
        elif action == 4:
            return "left"
        elif action == 5:
            return "down"
        elif action == 6:
            return "up-right"
        elif action == 7:
            return "up-left"
        elif action == 8:
            return "down-right"
        elif action == 9:
            return "down-left"
        elif action == 10:
            return "up-fire"
        elif action == 11:
            return "right-fire"
        elif action == 12:
            return "left-fire"
        elif action == 13:
            return "down-fire"
        elif action == 14:
            return "up-right-fire"
        elif action == 15:
            return "up-left-fire"
        elif action == 16:
            return "down-right-fire"
        elif action == 17:
            return "down-left-fire"
        elif action == 40:
            return "reset"