Experience processing
#!/usr/bin/env python3
import gym
+import time
import universe
import argparse
import numpy as np
@@ -14,6 +15,9 @@ from torch.autograd import Variable
ENV_NAME = "wob.mini.BisectAngle-v0"
REMOTE_ADDR = 'vnc://gpu:5900+15900'
+GAMMA = 0.99
+REWARD_STEPS = 4
+
def step_env(env, action):
while True:
@@ -27,6 +31,7 @@ def step_env(env, action):
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("-n", "--name", required=True, help="Name of the run")
+ parser.add_argument("--cuda", default=False, action='store_true', help="CUDA mode")
args = parser.parse_args()
env = gym.make(ENV_NAME)
@@ -34,16 +39,25 @@ if __name__ == "__main__":
env = wob_vnc.MiniWoBCropper(env)
env.configure(remotes=REMOTE_ADDR)
- obs = env.reset()
net = model_vnc.Model(input_shape=(3, wob_vnc.HEIGHT, wob_vnc.WIDTH),
n_actions=env.action_space.n)
print(net)
- obs, reward, done, info = step_env(env, env.action_space.sample())
- obs_v = Variable(torch.from_numpy(np.array(obs)))
- r = net(obs_v)
- print(r[0].size(), r[1].size())
+ agent = ptan.agent.PolicyAgent(lambda x: net(x)[0], cuda=args.cuda,
+ apply_softmax=True)
+ exp_source = ptan.experience.ExperienceSourceFirstLast(
+ [env], agent, gamma=GAMMA, steps_count=REWARD_STEPS, vectorized=True)
+
+ # obs, reward, done, info = step_env(env, env.action_space.sample())
+ # obs_v = Variable(torch.from_numpy(np.array(obs)))
+ # r = net(obs_v)
+ # print(r[0].size(), r[1].size())
+ for idx, exp in enumerate(exp_source):
+ print(exp)
+ if idx > 100:
+ break
+ time.sleep(0.5)
pass
