fix action shape

2017-10-27 09:06:10 +11:00
parent ce4a8286fd
commit a7b14b8091
1 changed files with 1 additions and 1 deletions
--- a/tutorial-contents/405_DQN_Reinforcement_learning.py
+++ b/tutorial-contents/405_DQN_Reinforcement_learning.py
@ -58,7 +58,7 @@ class DQN(object):
        # input only one sample
        if np.random.uniform() < EPSILON:   # greedy
            actions_value = self.eval_net.forward(x)
-            action = torch.max(actions_value, 1)[1].data.numpy()[0]     # return the argmax
+            action = torch.max(actions_value, 1)[1].data.numpy()[0, 0]     # return the argmax
        else:   # random
            action = np.random.randint(0, N_ACTIONS)
        return action