forked from mila-iqia/babyai
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathtest_mission_gen.py
executable file
·102 lines (80 loc) · 2.23 KB
/
test_mission_gen.py
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
102
#!/usr/bin/env python3
import random
import time
from optparse import OptionParser
from babyai.levels import level_dict
from PyQt5.QtWidgets import QApplication
from gym_minigrid.rendering import Window
def test():
parser = OptionParser()
parser.add_option(
"--level-name",
default='OpenRedDoor'
)
parser.add_option(
"--seed",
type="int",
default=-1
)
parser.add_option(
"--partial-obs",
action='store_true'
)
(options, args) = parser.parse_args()
rng = random.Random()
seed = options.seed
level = level_dict[options.level_name]
mission = None
app = QApplication([])
window = Window()
def reset():
nonlocal seed
nonlocal mission
if options.seed == -1:
seed = rng.randint(0, 0xFFFFFF)
mission = level(seed=seed)
print('seed=%d' % seed)
print(mission.instrs)
print(mission.surface)
pixmap = mission.render('pixmap')
window.setPixmap(pixmap)
window.setKeyDownCb(keyDownCb)
def keyDownCb(keyName):
if keyName == 'ESCAPE':
window.close()
return
if keyName == 'BACKSPACE':
reset()
return
action = 0
if keyName == 'LEFT':
action = mission.actions.left
elif keyName == 'RIGHT':
action = mission.actions.right
elif keyName == 'UP':
action = mission.actions.forward
elif keyName == 'SPACE':
action = mission.actions.toggle
elif keyName == 'PAGE_UP':
action = mission.actions.pickup
elif keyName == 'PAGE_DOWN':
action = mission.actions.drop
else:
return
obs, reward, done, info = mission.step(action)
print("is done:", done)
if done == True:
reset()
reset()
while True:
time.sleep(0.01)
if options.partial_obs:
obs = mission.gen_obs()
pixmap = mission.unwrapped.get_obs_render(obs['image'], 32)
else:
pixmap = mission.render('pixmap')
window.setPixmap(pixmap)
app.processEvents()
if window.closed:
break
test()