-
Notifications
You must be signed in to change notification settings - Fork 7
/
rocket_play.py
executable file
·52 lines (42 loc) · 1.61 KB
/
rocket_play.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
#!/usr/bin/env python3
import argparse
import gym
from gym import wrappers
from lib import model, kfac
from PIL import Image
import numpy as np
import torch
ENV_ID = "RocketLander-v0"
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("-m", "--model", required=True, help="Model file to load")
parser.add_argument("-e", "--env", default=ENV_ID, help="Environment name to use, default=" + ENV_ID)
parser.add_argument("-r", "--record", help="If specified, sets the recording dir, default=Disabled")
parser.add_argument("-s", "--save", type=int, help="If specified, save every N-th step as an image")
args = parser.parse_args()
env = gym.make(args.env)
if args.record:
env = wrappers.Monitor(env, args.record)
net = model.ModelActor(env.observation_space.shape[0], env.action_space.shape[0])
net.load_state_dict(torch.load(args.model))
obs = env.reset()
total_reward = 0.0
total_steps = 0
while True:
env.render()
obs_v = torch.FloatTensor(obs)
mu_v = net(obs_v)
action = mu_v.squeeze(dim=0).data.numpy()
action = np.clip(action, -1, 1)
if np.isscalar(action):
action = [action]
obs, reward, done, _ = env.step(action)
total_reward += reward
total_steps += 1
if done:
break
if args.save is not None and total_steps % args.save == 0:
o = env.render('rgb_array')
img = Image.fromarray(o)
img.save("img_%05d.png" % total_steps)
print("In %d steps we got %.3f reward" % (total_steps, total_reward))