Learning Reinforcement Learning: REINFORCE with PyTorch! | by Christian Hubbs | Towards Data Science

The REINFORCE algorithm is one of the first policy gradient algorithms in reinforcement learning and a great jumping off point to get into more advanced approaches. Policy gradients are different than Q-value algorithms because PG’s try to learn a parameterized policy instead of estimating Q-values of state-action pairs. So the policy output is represented as a probability distribution over actions rather than a set of Q-value estimates. If any of this is confusing or unclear, don’t worry, we’ll break it down step-by-step!
In this post, we’ll look at the REINFORCE algorithm and test it using OpenAI’s CartPole environment with PyTorch. We assume a basic understanding of reinforcement learning, so if you don’t know what states, actions, environments and the like mean, check out some of the links to other articles here or the simple primer on the topic here.
We can distinguish policy gradient algorithms from Q-value approaches (e.g. Deep Q-Networks) in that policy gradients make action selection without reference to the action values. Some policy gradients learn an estimate of values to help find a better policy, but this value estimate isn’t required to select an action. The output of a DQN is going to be a vector of value estimates while the output of the policy gradient is going to be a probability distribution over actions.
For example, consider we have two networks, a policy network and a DQN network that have learned the CartPole task with two actions (left and right). If we pass a state s to each, we might get the following from the DQN:
And this from the policy gradient:
The DQN gives us estimates of the discounted future rewards of the state and we make our selection based on these values (typically taking the maximum value according to some ε-greedy rule). The policy gradient, on the other hand, gives us probabilities of our actions. The way we make our selection, in this case, is by choosing action 0 28% of the time and action 1 72% of the time. These probabilities will change as the network gains more experience.
To get these probabilities, we use a simple function called softmax at the output layer. The function is given below:
This squashes all of our values to be between 0 and 1, and ensures that all of the outputs sum to 1 (Σ σ(x) = 1). Because we’re using the exp(x) function to scale our values, the largest ones tend to dominate and get more of the probability assigned to them.
Now for the algorithm itself.
If you’ve followed along with some previous posts, this shouldn’t look too daunting. However, we’ll walk through it anyway for clarity.
The requirements are rather straightforward, we need a differentiable policy, for which we can use a neural network, and a few hyperparameters like our step size (α), discount rate (γ), batch size (K) and max episodes (N). From there, we initialize our network, and run our episodes. After each episode, we discount our rewards, which is the sum of all of the discounted rewards from that reward onward. We’ll update our policy after each batch (e.g. K episodes) is complete. This tends to help provide stability for training.
The policy loss (L(θ)) looks a bit complicated at first but isn’t that difficult to understand if you look at it closely. Recall that the output of the policy network is a probability distribution. What we’re doing with the π(a | s, θ), is just getting the probability estimate of our network at each state. We then multiply that by the sum of the discounted rewards (G) to get the network’s expected value.
For example, say we’re at a state s the network is split between two actions, so the probability of choosing a=0 is 50% and a=1 is also 50%. The network randomly selects a=0, we get a reward of 1 and the episode ends (let’s assume discount factor is 1). When we go back and update our network, this state-action pair gives us (1)(0.5)=0.5, which translates into the network’s expected value of that action taken at that state. From here, we take the log of the probability and sum over all of the steps in our batch of episodes. Finally, we average this out and take the gradient of this value to make our updates.
Just for a quick refresher, the goal of Cart-Pole is to keep the pole in the air for as long as possible. Your agent needs to determine whether to push the cart to the left or the right to keep it balanced while not going over the edges on the left and right. If you don’t have OpenAI’s library installed yet, just run pip install gym and you should be set. Also, grab the latest off of pytorch.org if you haven’t already.
Go ahead and import some packages:
import numpy as npimport matplotlib.pyplot as pltimport gymimport sysimport torchfrom torch import nnfrom torch import optimprint("PyTorch:\t{}".format(torch.__version__))PyTorch:	1.1.0
With our packages imported, we’re going to set up a simple class called policy_estimator that will contain our neural network. It’s going to have two hidden layers with a ReLU activation function and softmax output. We’ll also give it a method called predict that enables us to do a forward pass through the network.
class policy_estimator():    def __init__(self, env):        self.n_inputs = env.observation_space.shape[0]        self.n_outputs = env.action_space.n                # Define network        self.network = nn.Sequential(            nn.Linear(self.n_inputs, 16),             nn.ReLU(),             nn.Linear(16, self.n_outputs),            nn.Softmax(dim=-1))        def predict(self, state):        action_probs = self.network(torch.FloatTensor(state))        return action_probs
Note that calling the predict method requires us to convert our state into a FloatTensor for PyTorch to work with it. Actually, the predict method itself is somewhat superfluous in PyTorch as a tensor could be passed directly to our network to get the results, but I include it here just for clarity.
The other thing we need is our discounting function to discount future rewards based on the discount factor γ we use.
def discount_rewards(rewards, gamma=0.99):    r = np.array([gamma**i * rewards[i]         for i in range(len(rewards))])    # Reverse the array direction for cumsum and then    # revert back to the original order    r = r[::-1].cumsum()[::-1]    return r — r.mean()
One thing I’ve done here that’s a bit non-standard is subtract the mean of the rewards at the end. This helps to stabilize the learning, particularly in cases such as this one where all the rewards are positive because the gradients change more with negative or below-average rewards than they would if the rewards weren’t normalized like this.
Now for the REINFORCE algorithm itself.
def reinforce(env, policy_estimator, num_episodes=2000,              batch_size=10, gamma=0.99):    # Set up lists to hold results    total_rewards = []    batch_rewards = []    batch_actions = []    batch_states = []    batch_counter = 1        # Define optimizer    optimizer = optim.Adam(policy_estimator.network.parameters(),                            lr=0.01)        action_space = np.arange(env.action_space.n)    ep = 0    while ep < num_episodes:        s_0 = env.reset()        states = []        rewards = []        actions = []        done = False        while done == False:            # Get actions and convert to numpy array            action_probs = policy_estimator.predict(                s_0).detach().numpy()            action = np.random.choice(action_space,                 p=action_probs)            s_1, r, done, _ = env.step(action)                        states.append(s_0)            rewards.append(r)            actions.append(action)            s_0 = s_1                        # If done, batch data            if done:                batch_rewards.extend(discount_rewards(                    rewards, gamma))                batch_states.extend(states)                batch_actions.extend(actions)                batch_counter += 1                total_rewards.append(sum(rewards))                                # If batch is complete, update network                if batch_counter == batch_size:                    optimizer.zero_grad()                    state_tensor = torch.FloatTensor(batch_states)                    reward_tensor = torch.FloatTensor(                        batch_rewards)                    # Actions are used as indices, must be                     # LongTensor                    action_tensor = torch.LongTensor(                       batch_actions)                                        # Calculate loss                    logprob = torch.log(                        policy_estimator.predict(state_tensor))                    selected_logprobs = reward_tensor * \                          torch.gather(logprob, 1,                         action_tensor).squeeze()                    loss = -selected_logprobs.mean()                                        # Calculate gradients                    loss.backward()                    # Apply gradients                    optimizer.step()                                        batch_rewards = []                    batch_actions = []                    batch_states = []                    batch_counter = 1                                    avg_rewards = np.mean(total_rewards[-100:])                # Print running average                print("\rEp: {} Average of last 100:" +                        "{:.2f}".format(                     ep + 1, avg_rewards), end="")                ep += 1                    return total_rewards
For the algorithm, we pass our policy_estimator and env objects, set a few hyperparameters and we’re off.
A few points on the implementation, always be certain to ensure your outputs from PyTorch are converted back to NumPy arrays before you pass the values to env.step() or functions like np.random.choice() to avoid errors. Also, we use torch.gather() to separate the actual actions taken from the action probabilities to ensure we’re calculating the loss function properly as discussed above. Finally, you can change the ending so that the algorithm stops running once the environment is “solved” instead of running for a preset number of steps (CartPole is solved after an average score of 195 or more for 100 consecutive episodes).
To run this, we just need a few lines of code to put it all together.
env = gym.make('CartPole-v0')policy_est = policy_estimator(env)rewards = reinforce(env, policy_est)
Plotting the results, we can see that it works quite well!
Even simple policy gradient algorithms can work quite nicely and they have less baggage than DQN’s which often employ additional features like memory replay to learn effectively.
See what you can do with this algorithm on more challenging environments!