Proximal Policy Optimization (PPO) - How to train Large Language Models
Luis Serrano AcademyDisclaimer: The transcript on this page is for the YouTube video titled "Proximal Policy Optimization (PPO) - How to train Large Language Models" from "Luis Serrano Academy". All rights to the original content belong to their respective owners. This transcript is provided for educational, research, and informational purposes only. This website is not affiliated with or endorsed by the original content creators or platforms.
Watch the original video here: https://www.youtube.com/watch?v=TjHH_--7l8g
Hello, my name is Luis Serrano and this is Serrano Academy. This video is about proximal policy optimization, or PPO. Proximal policy optimization, or PPO, is tremendously important in reinforcement learning and in machine learning in general, especially nowadays because PPO is the method that is used to train large language models.
See, large language models are trained using something called reinforcement learning with human feedback, or RLHF. RLHF is the topic of the next video, and it pretty much consists of training a large language model by making it play a game. In this game, the model scores highly if the humans rate it well, and the technique used in RLHF is proximal policy optimization.
So in this video, I'm going to tell you everything about proximal policy optimization. In short, PPO is very similar to other reinforcement learning methods, like the ones you've seen in the previous video, where you have two neural networks: you train a value network and a policy network. Now, what's special about PPO is that the value network and the policy network get trained at the same time. Let me show you how.
In this video, we're going to learn three things: first, the value neural network; then, the policy neural network; and then, how to train them. But the trick here is that we're going to train them simultaneously.
So let's introduce a grid world, which is very similar to the one in the previous video. Here we have a grid and an agent that's going to walk around this grid. Now, in this grid, some cells are special. For example, here there's a lot of money, here there's also a lot of money, and here there is a dragon. So if you land in the money, you're good, and if you land in the dragon, then that's bad because the dragon will eat you.
We can put some numbers: if you land in the money over here, you get $+5$ points; over here, you get $+4$ points; and if you get eaten by the dragon, you get $-1$ point. There are also some barriers that you are not allowed to cross, like the one here and the one here.
So the game consists of the agent walking around the grid, moving up, down, left, or right, and when it hits one of the special boxes, then it gets that many points. In this run, it won four points. If it plays again, let's see what's going to happen. Now it got eaten by the dragon, so it got $-1$ point. It doesn't have to start in the bottom left; it can start anywhere. In this game, it gained five points. The idea is to get the agent, which is that orange ball there, to learn how to play the game in an optimal way.
That grid world is pretty much the quintessential example of reinforcement learning. But you can think of many problems being solved using reinforcement learning, like playing chess, playing Go, or playing many games like Atari or Nintendo games. Also, robots learn to move using reinforcement learning, and most importantly, large language models, like the ones we use all the time, are trained using reinforcement learning.
So let's talk about states. A state is very simple: it's any place where the agent can be. It cannot be on the barriers, but it can be on pretty much any cell on the grid. In other words, a state is any position where the game can be.
We also have actions. Let's look at some square, like the one on the right, and let's say it has four neighbors. It can go up, down, left, and right, and those are the four actions. So the actions are basically the arrows where the agent could potentially move. Now, we can tell the agent not to walk into the walls, so we can remove some of these arrows, and those are the available actions that one can have for each of the states.
Now let's talk about values. Let's look, for example, at this cell in the middle. If the agent is here, then what do I know about the game? Well, let's say that if the game is played well, then I can reach a value of five from here. So that means that if I'm in this state, I can assign a value of five because it's sort of the best I can reach from that place. Now it's clear that pretty much any square other than the four and the $-1$ have a value of five, because from any place in this grid, I can reach that five on the top right.
Now, that game is not very interesting, so let's make it a little more interesting. Let's say that it costs one point to move one spot in this grid. So what happens if I'm in this spot over here? Well, I can move to the five and I earn five points, but I lost one from moving there, so the value here is four. That means that four points is the best score that I can achieve if I am standing at that cell.
What if I'm standing over here? Well, I can walk around in an optimal way in this direction. Notice that if I'm here, the value is four, and if I'm here, the value is three, because whatever I do, I need to do two steps to get to the five. And if I'm here, the value is two. That means if I'm in that grid, the best I can do is two points, which is if I walk as fast as I can to the $+5$.
Now, what happens if I'm here? Well, I could try to walk all the way to the $+5$, and if I continue playing this game, the best I can do is $-4$ because I lost a lot of points by walking, but then I gained five. However, I can also just walk up and get that four, and it's going to be much better because now I make two points. So the value here is not $-4$, it is $2$. In other words, the value is the maximum score that I can achieve if I play as intelligently as possible.
Now you may see a rule happening here. Let's say I have a box, its four neighbors, and the four actions I can make. Let's assume that I can make four actions and that I have calculated the values of all the neighbors. If the neighbors have values of $4$, $1$, $-3$, and $2$, then obviously I want to walk towards the $4$, which is the highest, and then I subtract $1$ because it costs me $1$ to get there, choosing this action over here.
Therefore, we have a rule, which is that the value of a cell is the maximum of the values of the neighbors minus $1$. Using that, we can start actually propagating the cells and calculating all the values. All the neighbors of a $5$ have a value of $4$. Now, among the stuff that I haven't filled, any neighbor of a $4$ has a $3$. Any neighbor of a $3$ that I haven't yet scored has a value of $2$, and any neighbor of a $2$ that I haven't yet scored has a value of $1$. Then we do the same thing with $0$ and $-1$.
Notice here that any neighbor of $-1$ that I haven't yet labeled has a value of $-2$. Notice something peculiar: if I'm here, the best I can do is walk down to the $-1$, so I have a value of $-2$. But if I'm here, I shouldn't walk to the $-1$; I should walk to this $4$ over here and get a value of $-2$. So those are the values.
Now that you've learned what values are, let me tell you what a policy is. Let's go back to the first game, where it was free to move around. If it's free to move around, then I always want to get to this $5$ over here, and I never want to hit the other two numbers. So if I don't want to hit this $4$, that means I should forbid these two arrows that lead to the $4$. And if I never want to hit my $-1$, then I should forbid these two arrows over here.
Now notice that whatever the agent does, as long as it follows those arrows, it's never going to hit the $4$ or the $-1$, and it's always going to get to the $5$ eventually, with probability $1$. It may take forever, it may loop around, but it'll win the game. So these arrows form what's called a policy.
Now let's look at a policy when you actually charge one point for moving to another cell. Let's remember what the values were: these are the values for each cell. So what do you think a good policy would be? Basically, if I put an arrow pointing to the highest value among the neighbors, that's a good policy.
So if I'm here at this $1$, then a good policy would be, for example, to move to this $2$. I could also move to the $2$ upstairs, but I'm definitely not going to move left. And if I'm here, a good policy is to move up, maybe right, and then back up.
If I have, for example, four neighbors and the four possible directions I can go, and I look at the values, and they're again $4$, $1$, $-3$, and $2$, we saw the value is $3$, but the policy is also the arrows that point to the highest value. I could have more than one arrow coming out if that highest value is repeated. So a good policy for this game would be this one. As long as I follow those arrows, I will always end up with the maximum number of points I can achieve at that particular cell, and that is the policy.
It seems like we solved the problem, because we have calculated the values and we have calculated the policy, and with either one of those separately, we can win this game. So what's the catch? As you may have noticed, in order to calculate the values and the policy, we had to visit every square in the grid. Not only that, but we had to visit them several times because we needed to iterate and actually recalculate these values several times until we got the perfect values and policy.
Visiting every square on this grid is not a problem when the grid is small, but imagine a grid that is really, really large—let's say thousands of columns and thousands of rows. For such a grid, it's hard to visit every square, and even harder to visit them several times. So what could we do?
As a matter of fact, this problem can get even harder because you could have infinitely many states. For example, if you're training a robot arm at a factory, every position of the arm is a state, and there are infinitely many, so there's no way you can visit all of them. We somehow need to approximate the value function, and for that, we're going to use neural networks.
So how is it going to work? The idea of a neural network is that it sort of gives us a continuous function. For the value network, what do I mean by that? Let's say we have discovered that the value on this square is pretty high; that means that the values on the squares around are probably high as well. Likewise, if a value is really low, then the values for the squares or states around them are also pretty low.
So we don't need to visit all the squares. If we manage to visit a bunch of them, we are able to calculate the function for the rest, or at least approximate it, and for that, we use the value neural network. Now, if you don't know what a neural network is, I recommend checking out the video on my channel where I explain neural networks in detail and how to train them.
But all you need to know is that in this neural network, we take the coordinates of the state or of the box, which are $X$ and $Y$ (the horizontal and vertical coordinates), and that's the input. The output is what the neural network thinks the value is, so in this case, for example, $3$. Of course, for another problem, the input of the neural network can be larger than two; it can be basically a bunch of variables that describe the state, and the output is the value.
We're going to draw this neural network like this. In here, it only has one layer, but it can have many, many layers—it could be really large—but the input is the state and the output is the value.
Now, just like we had a value neural network, we also have a policy neural network, and it's based on the same premise. Let's say that we've managed to calculate a good policy for this state over here. Notice that this policy has four arrows with different sizes; that means that the policy recommends us to go up with more probability than to go down. I'll elaborate on this policy in a little bit.
But in other words, if the policy for the state is this one over here, then it makes sense to think that the policies for its neighboring states are also similar. So just like we did with the value, we can approximate the policy with a neural network called a policy neural network. The input again is anything describing the state, so in this case, it is the two coordinates, and the output is the policy that's going to go in that state.
Now, I'm going to draw the policy neural network like this. It can also have many, many layers, but I'm going to draw just one. The input is a state, and the output is four things: the probability of going up, the probability of going right, of going down, and of going left. In this case, let's say that the output is $0.6$ for the probability of going up, $0.25$ for going right, $0.05$ for going down, and $0.1$ for going left.
This is called a stochastic policy, and I'll tell you more about it in a little bit, but the idea is that at every state, we have these probabilities of going in any of the four directions, and this policy hints that the value is going to be higher in the state to the top because the probability of going up is higher.
So in summary, we have the values over here, but since we can't calculate all the values because it's too much work, we use the value neural network to approximate them pretty well. It's going to take a lot fewer iterations to train a pretty good value neural network. It may not be exactly the values we want, but they're going to be close enough to give us a pretty good idea of how to do well at this game. Then we have the policy over here, and since we also don't have the bandwidth to calculate all this policy, we have the policy neural network that's going to approximate the policy really well.
Now, I mentioned before that this policy on the right is different because it's probabilistic, and that's actually a good thing. You see, the policy on the left is fully deterministic, so it tells you exactly what to do. Sometimes it gives you options, but it pretty much tells you what path to follow to get to the perfect place.
However, as we play this game, we don't want a deterministic policy because a deterministic policy won't explore the space very much. On the other hand, a stochastic policy like the one on the right allows us to really explore the space well because it gives us probabilities. It may not take us to the best score every time, but it allows us to walk around the grid and be able to improve our values and our policy. So a stochastic policy is always better than a deterministic one when you're training a model.
You may have heard the terms "explore" and "exploit" in machine learning. The policy on the left exploits, because once it finds a good path, it keeps taking that path and getting the best possible score it can from what it knows. However, the stochastic policy on the right actually balances explore and exploit, because it tries really hard to get high scores, but it also allows itself, with a small probability, to explore the entire space and be able to improve the values and the policy.
Now that we have the value and the policy neural network, let me show you how to train them. Okay, so recall that we have two neural networks: the value neural network, whose job is to approximate the values, and the policy neural network, whose job is to approximate the policy. I'm going to tell you how to train both of them, but first, I'm going to tell you how to train the value neural network.
Here's the main idea for training values, and I'm going to be a little vague, but don't worry, I'll get into the details slowly. Let's say that we have a state whose value is two, meaning the value neural network thinks that this value is two. However, let's say that the actual value is five. So this neural network underestimated the value, and we need to tell the neural network to increase this value of two.
So we introduce something called the gain. The gain here is three, which is the actual value minus the value the neural network thinks. In order to tell the neural network to increase this two, that gain is going to become a loss function that gets fed into the neural network. Again, I'll be more specific about this loss function in a little bit, but for now, think that we're going to use the gain to tell the neural network to increase this value of two, and then maybe it'll make it higher, like a $2.1$.
In the opposite scenario, let's say that the neural network thinks that this value is five and the actual value is three. So now the neural network overestimated, and the gain is $3 - 5$, which is $-2$. So that's actually a negative gain. Well, we're going to take this negative gain and tell the neural network to decrease this five in the training, and the neural network will output something like a $4.9$. In other words, we are training the neural network to get closer and closer to the actual values.
Now, I promised that I was going to be more specific, so let me be more specific. On the left, we have the values that the neural network gives. This neural network is not very well trained; it's in its early stages. As you can see, the values don't look very much like the actual values we calculated before, but we're going to do one iteration to make it better.
Let's say we're standing at this spot in the middle, and what we're going to do is walk around until we reach a terminal state. How do we know how to walk around? We can use the policy neural network that we have so far. It may not be great, but it's going to give us some path. Now, here I was lucky and we reached the terminal state of five, but we could have reached the terminal state of $-1$ or $+4$. It doesn't matter. The fact is that from this path, we can calculate some values, and we calculate them like we did at the very beginning.
Let's mimic the path over here. The values for this path are a $4$, because it's one square away from the $5$, then a $3$, then a $2$, $1$, $0$, and $-1$. So those are the actual values. Notice that they're not necessarily the exact optimal values; the policy could have taken me elsewhere, and I could have calculated different values. That doesn't matter; we're going to trust that if we do this many times, we're going to get close to the actual values.
Now, how do we tell the neural network how far its values on the gray path on the left are from what they should be on the colored path on the right? We take the difference and then we square it. If you've seen loss functions in machine learning, this is a textbook loss function: it's the squared difference, because the smaller the squared difference is, the closer the neural network is to calculating the actual values.
So let's take the difference of these numbers and square them, and we get these values over here. Now, it's a custom in loss functions to actually take the average instead of the sum, so we take the average. That means we divide by six, since there are six numbers. And this sum divided by six is $3.36$.
So $3.36$ is the loss. It is a quantification of how bad this neural network is, because the higher the loss, the farther the neural network is from approximating the values, and the lower the loss, the better it is. In an ideal situation where the loss was zero, that means the value neural network managed to approximate the actual values perfectly.
And just like we do with neural networks, we're going to take this loss and use it to train the value neural network. After one iteration of this training, the neural network will be able to output slightly better values for those cells and for everything around. As you can imagine, if we continue the training by taking different paths and always training the neural network, at some point we're going to get to something pretty good—a neural network that approximates the values pretty well. And that's the value neural network.
Now, if you like formulas, then here is the formula for the loss function of the value neural network:
Now, what is this? Well, $V$ is the value neural network, and $S_t$ represents the states. So those are the values that the neural network thinks those particular states in that path are. And what is $R_t$? Those are the actual values that the neural network wants to get to. They may depend on the path, but let's pretend for a moment that this is the goal of the neural network—the values that it wants to get to. Finally, this squared is the squared exponent over here.
So this is a squared loss function, and we use this loss function to train the neural network. If we can get the loss function to be as small as possible, that means that the values that the neural network is outputting are as close as possible to the actual ones. And that's how you train the value neural network.
Now let me show you how to train the policy neural network. Recall that we have two neural networks: the value neural network and the policy neural network. I already told you how to train the value neural network, so now I'm going to tell you how to train the policy neural network. It's a similar idea, but a little more complicated. Let's get through it.
Let's say that we have two states and a policy that tells us the probability of going from the left state to the right state, and it says it's $0.2$. This probability is given by the policy neural network. Our goal is to fine-tune and to train this policy neural net to come up with a better estimate for this probability.
Since we're training them at the same time, let's look at the value neural network, and let's say that the value neural network thinks that the value for the state at the right is two. And now let's say that the actual value is five. By actual value, I mean the same as before: we take a path, we follow it until it terminates, and then we calculate the values from there, and the path was given using the policy neural network. Again, it could be the path that gives us the best values, or it could be a path that gives us some other values—it doesn't matter for now.
Because the value neural network thought the value was two and the actual value was five, that means there was a gain of three. In other words, the value neural network underestimated the value by three. So we need to do two things: the one we did before, which is we tell the neural network to increase the value for the state because it estimated it as too low, and we also want to tell the policy neural network to increase the probability to get to that state.
Both neural networks are trained at the same time, so they both underestimate this box, and we need to tell them both to increase the value of the box and the probability to get to the box. So the policy neural network is told to increase it, and let's say it increases it to $0.3$. That's the idea.
Let's say we have the opposite scenario, where now the value neural network thinks the value is high—six, for example—and the actual value is not that high, let's say it's one. Then the gain is $-5$, which means it's a loss of five. We need to do two things: tell the value neural network to decrease the value on that state to get closer to one, and also tell the policy neural network to decrease the probability to get to that state because both neural networks overestimated this state. Therefore, the policy neural network would decrease this probability to, let's say, $0.1$. And that's the idea of how to train the policy neural network. Notice that they both get trained at the same time.
However, there's a special ingredient in training the policy neural network, which is impulse or momentum. Let me show you how momentum works with a simple example.
Let's say that there's a kid running—that's scenario one—and the kid is running towards the right with a velocity of $10$. We decide to help the kid out, so we push them. That means we apply a force in the same direction that the kid is running, which increases the speed, let's say from $10$ to $20$.
Now let's look at another scenario. In scenario two, the kid is running, but now they're running towards the left, and we still push them towards the right. That means that we change their direction, but the velocity obviously doesn't grow; it shrinks. Why? Because on the left, the velocity increases a lot, from $10$ to $20$, because we push the kid in the same direction that they were running. On the other hand, on the right, the velocity went down because we pushed the kid in a different direction than they were running. In other words, the impulse or the momentum was in the same direction as the velocity on the left, and on the right, it was in the opposite direction.
Obviously, you know how physics works, so this is nothing new. But I'm going to show you how this is used for the training.
Let's look at scenario one and make a little plot. In the horizontal axis, we have the iterations of the neural network. As the neural network is trained, there are many iterations, and in each one of them, the probability of a certain state changes.
Let's say that as we train the model, the probability of a certain state has been increasing, and then we notice that there's a gain, which means we need to increase the probability for the state. Since it's been growing, it increases by a lot, in the same way that when the kid gets pushed in the same direction that they are running, the velocity increases.
In the other scenario, let's say the probability of a certain state has been decreasing through the iterations of the neural network, but we discover that this is a pretty good state, so we have a gain—a positive gain—that means we need to increase the probability for this action. Therefore, it increases, but not very much, because now we're pushing the kid in the opposite direction than they were running.
In other words, on the left, we had an impulse. As the probability is increasing, we have a gain, and therefore the probability increases more. On the right, the impulse is going down, and the gain is a force that pushes it in the other direction, so it doesn't increase that much.
We also have two other scenarios, which I'm going to tell you quickly because I know you can imagine what's happening. If the probability is increasing and the gain is negative, then the probability now decreases. On the other hand, if the probability is decreasing and the gain is negative, then the probability will decrease much faster. This is because on the left, the impulse goes up and the gain goes down, so it doesn't decrease that much, whereas in scenario four, the impulse goes down and the gain goes in the same direction, so it decreases a lot more. I'm going to skip scenarios three and four numerically, but I encourage you to pause the video and actually work them out if you like.
So, what do I mean exactly by increasing or decreasing the probability at a certain rate? Let's say that we're training the neural network, and we are at iteration $100$. At iteration $100$, the policy tells us that going between this state on the left and the state on the right has a probability of $0.4$. That also happens in scenario two.
Now, let's say that we find a gain of $10$. This means that the neural network underestimated the value of this state by $10$, and that also happens in the scenario on the right. Clearly, we have to increase this probability. This is a pretty good state, so we want to increase that $P = 0.4$. In iteration $101$, the probability is going to be much higher than $0.4$. However, how high is it going to be? That depends on the impulse that this probability is carrying.
We go back to iteration $99$ and we look at what the probability was, or what the policy was, for going between these two states. Let's say in scenario one, iteration $99$ gave a probability of $0.2$. On the other hand, in scenario two, we had a probability of $0.8$.
What happens in scenario one? The probability increased from $0.2$ to $0.4$. That means it carries an upward impulse, and it's going to increase a lot more—let's say to $0.8$—because this impulse tells us that the probability is growing, we have a positive gain, and the probability grows a lot more.
In scenario two, however, the probability will grow from $0.4$, but not that much; it's going to grow to $0.5$. Why? Because it has been decreasing. The probability has an impulse that tells it that it's been decreasing, and then we apply a gain, which means it's not going to grow that much. It's the case when the kid is running in one direction and we push them in the other direction. Here, the probability is decreasing, and we tell it to increase. That's what I meant by momentum.
Let's look at the actual formula. Just like we used the gain to train the value neural network, we're going to use something called the surrogate objective function to train the policy neural network.
The way we're going to encode the impulse here is by dividing. We take the probability at iteration $100$, which is $0.4$, and divide it by the probability at iteration $99$, which is $0.2$, then multiply by the gain, and we get $20$. On the right, we do the same thing: we take the $0.4$ divided by the $0.8$, multiplied by the gain, and we get $5$.
The reason this $20$ is big is because we divide the $0.4$ by a small number, $0.2$. The smaller the probability is at iteration $99$, the higher the surrogate objective function is. On the right, we get $5$, which is a small number because we divided by a larger number, $0.8$.
In that way, iteration $99$ tells us exactly what to do. If it came from small to big, we're dividing by a small number and we get a large number like the $20$ on the left. If it's coming from big to small, we're dividing by a large number, so we have a smaller surrogate objective function. The surrogate objective function here has that momentum ingrained inside it.
If you like formulas, here is the formula. It looks complex, but it actually encodes exactly this quotient times the gain:
The idea is that this entire value is the surrogate objective function used to update the policy neural network. We use this function to train the policy neural network, and if we can optimize this, we reach a point where the policy neural network is approximating the actual policy really well. And that's how you train the policy neural network.
Now, it turns out that I haven't told you the whole story. We have a small problem. Let's say that given a policy, the probability of a certain action is $0.4$, but let's say that that increased drastically because in the previous iteration it was $0.001$. That ratio is $4,000$, and if we multiply it by the gain, we still have a very large number.
Putting a very large number into a neural network is not very good; it's going to make the training a little chaotic. So we actually want this surrogate objective function to not be drastically large or drastically small; we want it to be somewhere in the middle. What do we do? We're going to clip it. Let me show you how.
Recall that the numerator is the current probability and the denominator is the previous probability—that means the probability in the previous iteration of the training of the neural network. Since probabilities are between $0$ and $1$, let's look at where the ratio of these two could fall.
What we want is for this ratio to be close to $1$, because in principle, these probabilities shouldn't change too drastically in a single step. Let's say we make a small radius around $1$ of $0.3$. So the smallest possible value that we want it to be is $0.7$, and the highest possible value that we want it to be is $1.3$.
That means that if this ratio happens to fall between $0.7$ and $1.3$—for example, $0.9$—then we keep it as it is. Now let's say that this ratio doesn't fall in our desired interval. What if it's smaller than $0.7$? Well, we don't accept that, and we move it all the way to $0.7$. In other words, if the ratio is too small, we make it $0.7$.
What happens if it's too big? If it's anything bigger than $1.3$, we don't accept it, and we move it all the way to $1.3$. So if it's too large, we make it $1.3$. This is what clipping means, and we can select that parameter—in this case, it was $0.3$—to be anything we want.
This is how the surrogate objective function looks when it's clipped. We're going to call it $L^{\text{CLIP}}(\theta)$, and we're going to multiply this ratio by the gain, and then compare it with the clipped ratio times the gain, taking the minimum of these two:
Now I've told you the full story, because this clipped surrogate objective function is the one we use to train the policy neural network.
And that's it! That's proximal policy optimization. As a summary, you have the value neural network, which approximates the values, and the way to train it with the gain given by the value loss function. Then you have the policy neural network, which approximates the policy, and it's trained using the clipped surrogate objective function.
That's it for proximal policy optimization in this video. Stay tuned for the next one, where I'll show you how to use PPO in order to do reinforcement learning with human feedback, which is the way we train large language models.
As before, if you want to review the basics of deep reinforcement learning in detail, check out the previous video over here. They're all linked in the comments. And that's all, folks! Thank you very much for your attention.
As usual, if you enjoyed this video, please hit like, comment, or share it among your friends. I always enjoy reading the comments, especially if you have other topics you'd like me to teach. I get many ideas from there.
If you want to tweet at me, it's @SerranoAcademy, or check out my page at serrano.academy, where I have all these videos plus blog posts and a lot more material. Finally, if you want to check out my book, *Grokking Machine Learning*, there's a link below where you can get it, and there's a $40\%$ discount code only for my YouTube subscribers. Thank you very much, and see you in the next video!