corrected output neurons in reinforce

algorithms/reinforce.py

@@ -10,7 +10,6 @@ class PolicyNetwork(nn.Module):
     def __init__(self, n_state, n_hidden, n_action, lr):
         '''
         Initialize the policy neural network:
-            Use one hidden layer
             Input: a state, followed by a hidden layer
             Output: the probability of taking possible individual actions
                     use softmax function as the activation for the output layer
@@ -22,7 +21,7 @@ class PolicyNetwork(nn.Module):
             nn.ReLU(),
             nn.Linear(n_hidden, n_hidden),
             nn.ReLU(),
-            nn.Linear(n_hidden, n_hidden),
+            nn.Linear(n_hidden, n_action),
             nn.Softmax(dim=-1)
         )