Fix action indexing

reinforcement_learning/actor_critic.py

@@ -46,7 +46,7 @@ def forward(self, x):
         x = F.relu(self.affine1(x))
         action_scores = self.action_head(x)
         state_values = self.value_head(x)
-        return F.softmax(action_scores), state_values
+        return F.softmax(action_scores, dim=1), state_values
 
 
 model = Policy()
@@ -59,7 +59,7 @@ def select_action(state):
     m = Multinomial(probs)
     action = m.sample()
     model.saved_actions.append(SavedAction(m.log_prob(action), state_value))
-    return action.data
+    return action.data[0]
 
 
 def finish_episode():
@@ -88,7 +88,7 @@ def finish_episode():
     state = env.reset()
     for t in range(10000):  # Don't infinite loop while learning
         action = select_action(state)
-        state, reward, done, _ = env.step(action[0, 0])
+        state, reward, done, _ = env.step(action)
         if args.render:
             env.render()
         model.rewards.append(reward)

reinforcement_learning/reinforce.py

@@ -40,7 +40,7 @@ def __init__(self):
     def forward(self, x):
         x = F.relu(self.affine1(x))
         action_scores = self.affine2(x)
-        return F.softmax(action_scores)
+        return F.softmax(action_scores, dim=1)
 
 
 policy = Policy()
@@ -53,7 +53,7 @@ def select_action(state):
     m = Multinomial(probs)
     action = m.sample()
     policy.saved_actions.append(m.log_prob(action))
-    return action.data
+    return action.data[0]
 
 
 def finish_episode():
@@ -79,7 +79,7 @@ def finish_episode():
     state = env.reset()
     for t in range(10000):  # Don't infinite loop while learning
         action = select_action(state)
-        state, reward, done, _ = env.step(action[0, 0])
+        state, reward, done, _ = env.step(action)
         if args.render:
             env.render()
         policy.rewards.append(reward)