Cartpole test
@@ -194,6 +194,23 @@ def calc_loss(batch, net, cuda=False):
return nn.MSELoss()(q_v, y_v)
+def play_episode(env, net):
+ state = env.reset()
+ total_reward = 0.0
+
+ while True:
+ state_v = Variable(torch.FloatTensor([state]))
+ q_vals_v = net(state_v)
+ _, act_v = torch.max(q_vals_v, dim=1)
+ action = act_v.data.numpy()[0]
+ new_state, reward, is_done, _ = env.step(action)
+ total_reward += reward
+ if is_done:
+ break
+ state = new_state
+ return total_reward
+
+
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--cuda", default=False, action='store_true', help="Enable cuda mode")
@@ -202,6 +219,8 @@ if __name__ == "__main__":
writer = SummaryWriter(comment='-cart')
env = BufferWrapper(gym.make("CartPole-v0"), n_steps=4)
+ test_env = BufferWrapper(ImageWrapper(gym.make("Pong-v4")), n_steps=4)
+
# net = DQN(env.observation_space.shape, env.action_space.n)
net = SimpleDQN(env.observation_space.shape, env.action_space.n)
tgt_net = TargetNet(net)
@@ -227,7 +246,6 @@ if __name__ == "__main__":
print("%d: reward %f" % (frame_idx, reward))
writer.add_scalar("reward", reward, frame_idx)
- if len(exp_buffer) >= BATCH_SIZE:
batch = exp_buffer.sample(BATCH_SIZE)
optimizer.zero_grad()
loss_v = calc_loss(batch, net, cuda=args.cuda)
@@ -236,9 +254,14 @@ if __name__ == "__main__":
epsilon = max(0.1, 1.0 - frame_idx / 10**5)
if frame_idx % SUMMARY_EVERY_FRAME == 0:
writer.add_scalar("epsilon", epsilon, frame_idx)
- print("%d: epsilon %f" % (frame_idx, epsilon))
- #writer.add_scalar("loss", np.mean(losses), iter_idx)
+ loss = loss_v.data.cpu().numpy()[0]
+ writer.add_scalar("loss", loss, frame_idx)
+ print("%d: epsilon %f, loss %f" % (frame_idx, epsilon, loss))
if frame_idx % SYNC_TARGET_FRAMES == 0:
tgt_net.sync()
+ reward = play_episode(test_env, tgt_net.target_model)
+ writer.add_scalar("reward_test", reward, frame_idx)
+ print("%d: synced, test episode reward=%f" % (frame_idx, reward))
+
frame_idx += 1
