CMU AI Agents 2026: 9. Reinforcement Learning Basics

0:00 Okay, so let's go ahead and get started. So today is going to be about an intro to reinforcement learning which will build on the SFT for agents that we had last week. This is going to be the first of three lectures that we'll have on reinforcement learning. And so in this first one, we're going to be talking about basic poly gradient methods, which are ways where agents interact with the environment and get feedback and then update their weights to make things that were successful that got high reward more likely and things that got low reward less likely. We'll then have a second lecture which will actually be next week on more advanced RL algorithms talking about off policy methods about some of the things that you need to do to stabilize RL and get it to work. And then our third lecture will be on RL systems and going into some of the practical frameworks and you know things that you need to think about and pay attention to in hyperparameters when you're training these agents. And all of these should set you up pretty well for assignment three, which is going to be focused on training models with reinforcement learning in two settings. So the first setting is going to be a simple kind of sandboxed setting, which we'll actually use as a running example in this lecture. It's going to be called number search. And the second one is going to be like a simple kind of crafting environment that's a little bit like Minecraft. So for these, they're definitely going to require more compute than your previous assignments have.

1:36 so you'll be using the compute credits more heavily, but you should have the foundations to get started on those assignments after these first two lectures. And yeah, and we're looking forward to to seeing what you all do for those. We'll also have a guest lecture on Thursday from Akari Asai who's a faculty here at CMU who's an expert on deep research and she'll be talking about that as another domain for agents that complements the the ones that we've talked about so far. Cool.

2:08 Any any quick logistic questions before we get started? Okay. Yeah. So today we'll have a bridge from supervised fine-tuning to policy gradient methods which are kind of the most the most basic but also one of the most elegant forms of reinforcement learning. So the running example that we'll use is this task that we'll call number search which is basically like training an agent to guess a number based on feedback about whether you're too high or too low. So this is a pretty simple task but you know easy to understand and easy to come up with a reasonable policy for it. So the setup is there's this number which is part of the environment for this task but isn't known to the agent and it's an integer between 1 and 16. The agent is going to get four attempts to guess the number. And after it makes a guess, the if it got it correct, the environment will say correct. But otherwise, the environment will say, you should go higher or you should go lower. And then the agent can make another guess up to four guesses in total. If at any point the agent gets it correct, the task ends and it gets a reward of one. But if you reach the end of your four guesses and you haven't gotten it correct, then it's a reward of zero.

3:30 Okay. So, a lot of the trajectories that we'll walk through are from an example of this where the hidden number is 11. So, the agent is going to try to guess that. And here's four possible attempts at this where in the first one the agent says 8, gets feedback, you should go higher. It says 12, gets feedback, you should go lower. 10 higher and then 11 so it gets it correct on the last guess. In another attempt we might start with a lower number. So this is actually what's the optimal policy here? What's the best starting number?

4:07 The middle, right? So it's kind of like 20 questions. You want to like split the space, right? So the best first guess is eight. But and actually how many of you folks have played Wordle? you know the guessing game from the New York New York Times. It's a fun game. It's kind of in a similar space and there's like an optimal first guess that you can make which reduces the the the space in an information theoretic way.

4:33 Similar concept here but simpler. So you could start off with a you know suboptimal first guess of four but you could still succeed in the task right by taking this feedback into account going higher then going higher then going you know lower and eventually getting it correct but you could also go wrong right so we could start with a very bad guess of 16 and unless we got that right the feedback's going to be lower and in this case we might time out so we don't get it within four guesses and similar Similarly, we could have another pretty bad policy of starting with one and then going two and three and four and and timing out and failing at the task.

5:14 Everybody clear on the the setup? Okay, great. So, we'll also have rewards for trajectories. And a reward is just a scalar value which tells you how good or bad that trajectory is. So in agentive settings, this might correspond to like did the code that you write pass all of the unit tests in like Swedbench for example and you get a reward of one if it passes all the tests a reward of zero if it doesn't pass all the tests or you could have you know like a value that's not zero or one that's not a binary reward too like what's the fraction of tests that you passed but in a lot of the domains that people are training on these days in what's called RLVR like reinforcement learning with verified rewards. The rewards will be either zero or one. but the methods that we'll be talking about today are more general.

6:06 they also apply to scaler valued rewards. So just a quick reminder of how SFT works. So in supervised fine-tuning, you have a demonstration of a task and you're training the agent to copy that demonstration by maximizing the probability of the demonstrated actions. Right? So if we break down one trajectory into these steps, we have a prompt that's given to the agent. we have a guess feedback from the environment higher. We have another guess more feedback from the environment and so on. And our task is going to be given these demonstrated actions which we'll write a star. We want to maximize the probability that the agent assigns to those actions conditioned on all the past history up to that point. Right? And so this will be just you know the supervised fine-tuning loss is maximizing the log of the probability of that sequence which is the sum of the log of the probabilities of the individual actions right and there's a lot of ways we could get these target actions that we're trying to imitate. So they could be from a person who gives a demonstration from the task. They could be from like a stronger model that we're trying to copy to distill from today. They could also be from trajectories from the agent itself that we're trying to reinforce in the model.

7:32 But this should hopefully be a review from you know your previous classes on LLM training and and from the SFT lecture. Okay. But SFT has some problems which reinforcement learning is going to address. So there's roughly like three of these that we'll focus on. The first one of these is what we might call task mismatch. So supervised fine-tuning maximizes the probability of demonstrated actions at each point in time. But what we're really interested in is maximizing the probability that the agent carries out the task successfully.

8:07 And these are similar to each other, but they're not quite the same. And there's a couple reasons for this. So the first reason is that there are multiple ways to carry out the task successfully, right? So, starting with eight is kind of the best thing to do, as a first guess, but, it would also be reasonable to start with, you know, four as long as you're able to, take the feedback into account. And for this, you know, here after this first step of feedback, we have various reasonable actions that we could take as well. and a demonstration, whether that's, you know, from a person or another model, only shows one way of solving the task, but we want the agent to be able to potentially use other ways as well.

8:56 And a second problem is that some actions that you can take are worse than others, but to SFT, they all look the same. So it would be very bad for us to in response to the feedback you should go higher it would be very bad for us to like output a number below eight right but SFT doesn't make a distinction between the actions that weren't demonstrated and so it would get the same loss if you know the agent had highest probability on a lower alternative action rather than a higher alternative action that wasn't the one that was demonstrated.

9:36 So that's one problem of the mismatch between this SFT loss and the objective that we're actually interested in of maximizing you know expected reward or the probability of success on the task. The second problem is data mismatch. So we'd like to be able to learn from suboptimal data too. Even if we assume that our demonstrations are the best way to achieve a task, there's, you know, if we had access to other demonstrations that were less optimal or maybe even weren't successful, those could provide useful signal for our models too. So say that we know that this is a bad trajectory like it times out. We wouldn't want to just do SFT on this, right? We wouldn't want to maximize the probability of those actions, but we might want to have some learning method that could have the model not take trajectories like this. And reinforcement learning is going to be one way of doing that.

10:42 The third problem is what's known as exposure bias. So in SFT, you're training a model to make predictions of the next action given a fixed sequence of past actions, right? And that fixed sequence of actions comes from somewhere. But then at evaluation time, the previous actions are going to be generated by the model itself, right? Because you have to produce actions one at a time. You have the actions you produced in the past. It's this auto reggressive process of interacting with the environment.

11:16 So if the agent produces an action that is very out of distribution from the histories that it was trained on, it might just not have the the response that it produces might just not be calibrated because it never saw data like that. So say that we only trained on demonstrations where the next integer was, you know, did match the feedback that we got. So it was always higher if the feedback was higher. Now if at generation time the model makes a mistake and samples a lower number, it might just totally throw it off because it didn't see data like that in training time if we only trained on these optimal policies.

11:58 Does that make sense? So one way of viewing reinforcement learning is it's a way of letting the model generate its own training data and also learn how to recover from the mistakes that it makes so that it's prepared to handle its own mistakes when you generate from it. And this property of having the model generate its own training data is called being on policy. It's generally a good thing to have because it prepares the model to you know deal with its own mistakes and the supervision that you get is a good fit for what the model's able to currently predict. but it can have some efficiency trade-offs and in Graham's lecture we'll talk a little bit we'll talk a lot actually about methods that relax these assumptions that go more off policy.

12:43 Cool. So these are three problems with supervised fine-tuning for agents and for sequential prediction generally. Any questions about them? All right. So, reinforcement learning the general idea is instead of having like a trajectory a fixed trajectory which you learn to imitate by you know maximizing the probability of the actions. We're going to have the agent or you know the model generate its own actions and then we'll evaluate the reward of those actions in the trajectory and then we'll adjust the parameters of the agent so that it makes high reward trajectories more likely and lowreward trajectories less likely. So formally and we'll see you know the formulation of this in a little bit we're going to be training the agent to maximize its expected reward you know expectation in probability under the agent and under the environment.

13:46 So for just schematically if we have you know a task like this number search game with the underlying goal in the environment B to guess number 11 we would have the agent generate multiple trajectories just by producing likely actions at each step from the agent. You know just in the way that we would typically evaluate the agent in the environment for each of those trajectories. we'll get a reward by by executing these actions in the environment and then we'll use that reward as feedback for the reinforcement learning method which is going to change the parameters of the agent to make these highreward trajectories more likely.

14:33 So this is going to address the three problems that we laid out with SFT. So task mismatch is addressed because we're now training the agent directly to maximize expected reward. I haven't shown you how that's done yet. A lot of the first part of the lecture is going to be about how we do that. We'll formulate expected reward. We'll take derivatives of it. We'll optimize it directly. It will also fix data mismatch because as we'll see especially towards the middle and end of the lecture, the agent will now be able to learn from all trajectories. So it will get some signal from these bad trajectories and will avoid taking actions that lead to trajectories like that in the future.

15:14 And it will also help deal with exposure bias because the agent is now going to be generating these trajectories that are used in training. It's going to be on policy. So if it makes mistakes and sees that those mistakes lead to a bad reward, it will get penalized for that and then learn to avoid those in the future. So it's sort of being prepared to deal with its own errors. Yeah. And yeah, like a lot of the issues that people had with like language models up to maybe like three or four years ago where they were being like they had kind of repetition biases. They would get stuck in loops. A lot of those issues got fixed once people started doing a final post-training step with reinforcement learning where the models were now able to recover from their own mistakes.

16:05 Okay. So we'll talk about how we you know formulate interactions with the environment as trajectories which will then let us set up the objectives that we optimize. Okay. So we won't spend too much time on this because it's should just be a recap from the first part of the course. But we have an agent which is going to be represented as a policy which we'll term pi theta. Theta are the parameters of the language model that you know that drives the agent and the agent's going to be producing actions by sampling from the language model and you know then applying those those samples in the harness that it's operating in that will produce actions which you then take in the environment and the environment gives you back observations and also rewards.

17:01 A lot of the settings that we'll look at only give you a reward at the end of the task, like when you actually execute the test cases. but pretty much everything that we talk about will also apply to settings that have intermediate rewards where for individual actions you take, you can also get some reward about the value of that action. But just some examples of how you can apply this general framework for text generation. the task instance might be the prompt that you're trying to produce a completion to. The observation might be the tokens you have so far. The action might be the next token. And the reward might be the score that you get from a judge model.

17:42 For example, if you're doing reinforcement learning from human feedback for a web agent, like the setting we talked about a week or two ago, the task instance might be the task and the website. And here there's a difference between the state of the environment, sort of everything that determines the environment and what the agent gets to observe. There might be some hidden state like what's the content of the backing databases, but the agent only gets to observe the page that it's currently on. And it can take actions like clicking, typing or scrolling. And it might get reward like intermediate reward like what are the number of subtasks from the the task description that it has completed so far. And for number search we've seen the task instance is the number that the agent is trying to guess which is going to be hidden from the agent. So in number search every task looks the same to the agent at the start. the observation would be high or lower or correct.

18:41 Action would be the next guess. And the reward is given if the agent gets it correct that's a one or if it times out then it gets a zero. So given these these observations and actions and tasks we can start to formulate some notation which will let us derive the reward terms that we're expected reward terms that we're trying to maximize. So a trajectory is a sequence of interled actions and observations and we'll represent the sequence of observations and actions up to a given time step as the history h.

19:20 So the history ht is going to have all of the actions and observations up to that point in time. And our agent policy pi theta remember theta are the parameters of the language model inside the agent. that's going to give us a distribution over the next actions that I would take g at a particular history. Right? So you take the history, you plug it into your agent framework, you get a distribution over next tokens which correspond to actions in the environment. Any questions about that?

19:54 Okay. So we also have rewards and in the most general setting you could be getting rewards along the way after each action or observation that you take. Those formulations whether it's after actions or observations can be equivalent to each other. in the number search setting, we only get a reward after you either time out or get the feedback that it was correct. So in the number search setting, your rewards will be binary and they'll just come at the end of the trajectory.

20:24 And we'll write the reward of the trajectory as the sum over the rewards at all intermediate time steps. So toao is going to represent trajectory which is you know the sequence of these observations actions and rewards. Okay. So that's the notation and given this we can also you know have probabilities over trajectories and the probability of a trajectory is going to depend on two things. So the first thing is the probabilities assigned by the agent. So we can decompose this into a product of probabilities over time steps where h sorry I dropped a p here but there should be a p before this term which is the probability from the environment of getting this observation and this reward conditioned on this history the action taken by the agent at that time step and the task and that probability from the environment we're going to multiply by the probability ility from the agent policy of producing this next particular action. So this is just the chain rule decomposing this probability of producing this entire trajectory into the product of its component parts.

21:43 And a key thing in RL is we don't actually know we we don't have to know anything about the environment. We don't have to know, you know, what's the probability of getting correct as our next observation or what's the probability of getting a next observation as I'm interacting with Amazon.com or something. The key thing about RL is we're going to interact with the environment which will by taking actions which will be a way of sampling from this distribution over observations and rewards. And that's actually going to be enough information for us to be able to learn how to differentiate an objective that depends on this with respect to our parameters and update the parameters. And that's kind of magic in some sense, but it also makes reinforcement learning algorithms a bit slow because you kind of just have to like try stuff and see what happens.

22:39 and that's why reinforcement learning takes a lot of compute to run because you have to do these experiments experiments with the environment, see what happens, and use that as a somewhat indirect source of feedback about how your reward depends on your parameters. Any questions about this? >> Yeah. Yeah. >> Yeah. So I've rewritten this a bit from so the question was in classical RL there's a marov assumption which says that the transition to the next state depends just on the previous state that you're in and the actions. So we're actually working with a palm dp here a partially observed marov decision process where there's a difference between the underlying state which you can't see and the observations that you can. So in a palm dp your probability of the next observation has to depend on all of the previous observations and actions because you don't have access to the state. but for a pom dp these formulations should be equivalent and this one's just a little bit easier to work with for our purposes.

23:54 Good question. Okay. Yeah, sorry about that P missing. Okay, so our objective is going to be to maximize the expected reward. and so what does this mean? So what we want to do is we want to take samples from our agent like we want to have the agent generate trajectories through the environment. and we're typically going to be doing that by sampling actions from the language model and then feeding those actions through the agent framework, right? so we can write that as sampling a trajectory from this distribution which is parameterized by the language model itself. And for each of those trajectories that we get by sampling from that distribution, it's a very complex distribution that depends on the agent, on the language model, and on the environment, we're going to get a trajectory and we're going to evaluate the reward of that trajectory that gives us a number. And we want to maximize our chance of getting a high reward here. So that's going to be maximizing this expectation. And we'll just write that objective as J of theta because it's a function of the parameters of our language model.

25:12 So in our very simple setting let's you know we have this number search we have say there's only four trajectories which are actually producible by the language model. So it's learned kind of a very sharp policy. say that it puts probability of 0.3 on this first trajectory 0.2 on the second and 0.25 25 on the third and fourth. generally blue is going to mean success, a reward of one and red is going to mean failure, a reward of zero.

25:42 So the expected reward if these are the only four trajectories which the agent can produce is going to be 0.5. Right? So just you know taking a weighted sum of these trajectory rewards weighted by the probability under the agent policy and the objective of training is going to be to optimize this. So we want to adjust the parameters theta to maximize this expected reward. So you can think of this as like this is kind of a bad current agent because it has a pretty high probability 0.5 of producing a trajectory with zero reward. We want to shift some of that probability. We want to shift ideally all of that probability to predicting trajectories that have reward of one. So shifting the probability mass from red to blue.

26:32 This is obviously simplified because you know really the agent could produce an inf not infinite but like combinatorally exponentially large number of possible trajectories but we're just kind of simplifying the visualization a bit here. So one way that you can do this is a pretty simple modification to SFT that's really well motivated. It's called expert iteration. And a particular type of expert iteration is called reinforced self-training or rest for short. So the idea in rest is going to be that we will just generate data from our current policy from our current parameters of the language model and then we will choose the trajectories that are successful and use those as fine-tuning data and update the model on those. And so that's where the reinforced term comes from. It's going to sort of reinforce the correct behaviors that the agent is already producing. so it has two steps which are going to be iterated. rest em is actually named like rest I think expectation maximization because it looks really similar to the em algorithm which some of you might have seen in an ML class. But the idea is we will alternate between these two steps. The first step is you sample trajectories from the current policy. You get the rewards of them. The second step is you will have the model make those trajectories very likely by just doing a bunch of SFT on those collected trajectories. and then you repeat you will get more trajectories from the current policy and you add those trajectories in to the existing set that you have of good trajectories and then you'll train your model on all of those. So you have this growing data set which you are having the model iteratively get very good at producing.

28:32 yeah so that's kind of a high level description and the way this works mechanically is you will generate a bunch of trajectories. You will evaluate the reward of each one. You'll select using like a threshold and this is easy if your reward is binary. if it's binary, you just keep the trajectories that have a reward of one and throw out the ones that have a reward of zero. And then you combine all of these kept trajectories across your instances. And those are the those are the next iteration of your training set.

29:09 And so then the objective is just the SFT objective. We want to maximize the probability of these actions in these cap trajectories. the one where the reward is for binary rewards one. Does that make sense? Any questions about that? So, schematically, this is actually going to look quite a bit like the reinforcement learning methods that we will talk about. this isn't like technically referred to as reinforcement learning. but it really has a lot of similarities to it because you're generating data from the model.

29:44 you're updating on that data in a way that is determined by the reward of the trajectories. And we'll see how reinforcement learning does this in a little bit more of a fine grained way later. But this is a very simple and pretty effective baseline that you can start with which reuses all the same machinery of supervised fine-tuning and is a good fit if you have binarized rewards. >> Yeah. Why is that self?

30:16 >> So the difference from self distillation is just that you're using the reward to filter. Self-distillation could also refer to training on all of the trajectories produced by the model. >> Yeah. So self it's called selfraining because it looks a lot like you know like that but it's reinforced because you're using the reward and so you're reinforcing the good trajectories. Yeah. I mean this is a more basic question but how can you guarantee trajectories that do have a positive reward?

30:49 >> What if like for example in the worst case you just have a bunch of bad trajectories? >> Yeah. So that's a great question like how do you guarantee that some of your trajectories have good reward? if your model can't produce trajectories that have good reward then reinforcement learning and this method are probably not a great fit. So there's a few things you can do. one is you can have some demonstrations of good trajectories which you train the model on. So you know like the R1 paper and a lot of reasoning models will have an initial what they call cold start stage where they will train the model on reasoning traces before they start doing reinforcement learning. So that just makes it easier for RL or something like this to be able to produce trajectories that have high reward. You can also like look at curriculum. So you can select problems that do have enough trajectories which produce reward. Train on those first and then as the model gets better you're able to tackle harder problems. You can also have like intermediate rewards like if you have some knowledge of the problem and like know it's useful to explore the code repo first. You could give the model reward for doing that and then hopefully like after it learns to do that it'll kind of unlock solving the whole task. But this is a really big kind of issue that RL algorithms have to grapple with and those are a few solutions to it.

32:14 Yeah. Cool. Okay. So now let's actually get into the policy gradient methods. So we're interested to maximize this quantity J of theta the expectation of reward. And we can't just use back propagation to do this for two reasons. The first is that we are sampling actions from our policy. And the second reason is that we are then executing these actions in this like blackbox environment which is potentially which is almost always like non-ifferiable.

32:53 and so this breaks our ability to do back propagation to you know directly you know like use the chain rule like you saw in a machine learning class to get the derivative of your reward with respect to the parameters. So actions are discrete and they're sampled. So the derivative with respect to the parameters of a sample is zero almost everywhere because you know the parameters change the probabilities but don't until you get to kind of a step point don't change the action sampled and the environment is usually not differentiable. So the general solution to this is actually a pretty classic method called reinforce which was from Williams in 1992 and it's going to use this definition of the expected reward and then basically just some algebra tricks in order to compute an estimate of this that uses samples from the model and the same back propagation procedures that we've looked at before to get par you know to get partial derivatives with respect to these parameters theta and we'll show you how to do that on the next slide.

34:13 So this is the objective that we want to optimize. It's this expectation over trajectories from our environment and our policy and the expectation of the reward of those trajectories. So this is a if you recall what an expectation is it's a sum over these things these trajectories of the term in here the reward but weighted by the probability of that trajectory. So this is a sum over a really large space like all possible trajectories from the language model.

34:47 We're almost never going to be able to compute that exactly, but this is what it would come out to be if we wanted to compute it. Taking the gradient of this. So we just, you know, take the partial derivative with respect to the parameters theta. we can move that inside the sum and this reward term depends only on the trajectory. So once we know the trajectory, we know what the reward is. it doesn't depend on the parameters anymore. The parameters produce the trajectory. but then the reward is just a function of the trajectory. So this reward term is going to be multiplied by this partial derivative of this probability.

35:32 But this is still difficult because this is a sum over like a potentially very large number of trajectories. So we're going to do an algebra trick here. And this is called the score ratio trick. we're going to just multiply and divide by this term p of theta. So we have that over here and we also have that in the denominator. Those cancel. So this is the same as that. And then we can use this trick which is the partial derivative of log p with respect to theta. We can just use the chain rule and it's get it's the partial derivative of p / p. So we can use that identity here and we get this term. And this is going to be kind of magical because we now have this which is going to be an expectation over trajectories times a reward term times this term here.

36:27 And this term is going to be this is the same computation that we do in SFT. Right? In SFT, we're also taking a partial derivative of the log of a bunch of unit tokens with respect to the parameters of the language model. And we're going to be able to estimate this quantity by sampling trajectories and then computing this term here. And this will give us an estimate of the of the gradient of this expected reward. Any questions about that?

37:09 >> So this is really general doesn't make any assumptions about the language model or about the environment. >> Yeah. >> Are you going the series of states and actions but the actions are generated by large are also sequences of >> yeah normal way to reconcile that >> yeah we'll we'll break that down in a couple slides but yeah it's a good question yeah so we'll just copy that term here and then we'll break this into the component pieces so these trajectories are going to be sampled by having the language model and the agent on top of it interact with the environment and that's you know the same procedure that you would use when you're just running this agent. That gives you a trajectory. and we this should be a sum over this should be an average over trajectories because we can estimate an expectation by drawing some number of samples from this distribution and then averaging the quantity inside the expectation. So if we average over a bunch of trajectories drawn from this distribution this quantity here then we will get an estimate of this gradient.

38:33 So what are these terms inside here? Well it's since this is a trajectory we break this up into steps. So it'll be a sum over time steps of the gradient that's due to that particular action. And we'll see how to compute that in a slide. So yeah, so the three components here are the trajectories. We get those by sampling. the task level evaluation which we get from the environment. And then this gradient here which is going to be produced using the same way that we get gradients in in supervised fine-tuning. But the key thing is these actions here are going to be the ones sampled from the agent rather than ones that are demonstrated by you know some ground truth.

39:25 Okay. So yeah, how do we get log probabilities of trajectories? I'll go a little bit quickly through this. The the key thing to keep in mind is that the probability of a trajectory depends both on the probability from the language model plus agent but it also depends on the probability of getting the particular observations and rewards that you did from the environment. But the key thing is is that this probability doesn't depend on the parameters of the language model in any way. So if we differentiate this log probability of you know the trajectory with respect to the parameters this entire term just goes away because it's a constant with respect to the parameters. So we can just ignore it and this gradient for the trajectory is the sum of the gradients for the actions within the trajectory.

40:25 So this is mostly just kind of like bookkeeping for rigor. but one thing that is important to think about is these actions that are being produced by the agent might correspond to multiple tokens from your language model, right? So like maybe you need to say you need to say you're you're making a complex tool call and that splits up into multiple tokens. of course, the log probabilities of those tokens are just going to be summed up to produce the log probability of the entire action.

41:02 And when you have this gradient of a trajectory, the gradient for a trajectory, you can just multiply by the reward, which is a scalar value from the environment to get the contribution of this trajectory that you sampled to the expectation. And how does this like actually look if you are running this on a real example. So you would have your trajectory which is sampled from the agent. And these blue tokens are the ones that are actually generated by the model. All the gray tokens are ones that are either like part of the prompt or part of the environment or added by the scaffold. But these are the ones that are actually produced by generating from the language model. and those are going to be the ones which contribute to the loss for policy gradient.

42:02 So we'll have all of these generated tokens. We'll get log probabilities for those of that generated action selected from the distribution of the language model over particular tokens at that point. And then we will just scale all of those token log props by the reward that corresponds to this trajectory. Does that make sense? And then if you take the gradient of this, the gradient passes inside and gets applied to this log probability here. So that your gradient is going to be the reward times the gradient. So when I first looked at the code that implemented reinforce on like a pretty complex task, it was like using RL to optimize dialogue models from a paper called deal or no deal. It was kind of like really surprising that the implementation is so simple like you just do basically an SFT gradient computation and then multiply that gradient by a reward or multiply the loss by a reward and take the gradient of that. But it it works and the math works out.

43:13 Yeah. So we'll have a mask here which is going to be you know you mask only to select the tokens that you generate and then you scale all of those by the reward for the trajectory. And that's how you implement the reinforce policy gradient algorithm. And so the training loop is just you sample task instances you evaluate reward on them. get the loss using that slide that we just showed you and then do an update on all of the instances in your batch of examples here.

43:51 Later on we'll when we cover gpo we'll talk about how you could generate multiple possible trajectories for a given task example I and then use all of those in the update. But if you're doing just like standard reinforce, you would just be using one trajectory for each example and then computing the update from that. So your batch size would be the number of task instances that you're training on. but when we get to gpo you will also have multiple samples per task instance. Any questions about this?

44:29 Okay. Yeah. So just very quickly with binary rewards this looks a lot like the rest procedure that we talked about. Both of these are just scaling log probabilities of actions by a reward. The difference is is that this reinforce method is like fully online. You're like sampling these trajectories and then updating on them directly. Whereas with REST, you build up a very large data set and then train on all of those and maybe do like, you know, a full epoch or multiple epochs on all of those instances. But the intuition behind them is basically the same.

45:08 Okay, but this has a lot of issues too. And so it's very rare for people to just run standard reinforce policy gradient out of the box. So the biggest I think the the the way of kind of the strongest motivation for the next methods that we'll talk about is thinking about how rewards should really be the update that you do from rewards should really be relative. So let's look at these two different trajectories. this one is successful and so all the actions are going to be upweighted by a reward of one if we use the update rules we've talked about.

45:44 This one is unsuccessful. So we had a bad first guess, we eventually time out. And if we're just using basic reinforce or arrest, we are not going to do any update based on this example at all. Right? Because it's scaled by the reward. The reward is zero. So but we actually want to be able to learn not to do this, right? so the intuition for the next part is going to be about shifting our rewards so that we still get a correct estimate of the gradient but we're able to learn from rewards at all different scales or shifts. and another way of thinking about this is that the update that you do should depend on how well the policy would be expected to do on this example. So say that this is just like an easy example and it and the model like 95% of the time gets a reward of one on it. If 5% of the time you make a really dumb mistake and get a reward of zero, you should update very heavily against those costly but rare mistakes. And baselines and advantages that we'll talk about next will be a way of doing that which will compare rewards to the reward that you would expect to get from the policy running on this example.

47:08 And the math for that looks like this. We'll define an advantage which is going to be a function of the history up to a particular time step and an action. And that's going to be say that you take this action from you take this particular action and then you sample from the language model for the rest of the trajectory that will tell you and you contrast that against if you just sample from the language model including that action and then continue on. This will tell you how good is this particular action relative to how good is the what's the value of this particular action taking into account that the language model is going to take all the rest of the actions in the future.

47:53 So maybe a concrete example will will help with this. Let's say that we are partway through this guessing game and we've guessed eight and gotten feedback higher, guessed 12, gotten feedback lower. So we know that the numbers have to be between 8 and 12. So they can either be 9, 10 or 11. So we don't know exactly what this value function is. It depends on what our policy is. But we can think about it's probably something like 2/3 or one, right? So like if the policy is is does the best thing possible always you should be able to just always win from this situation. But if it's a more random policy within the set of things that are valid it might be more like 2/3.

48:43 But we can think about the advantage functions of particular guesses relative to this baseline. So like a guess of 11 would be bad because it's at one end of this range and if 11 is not correct then we have two candidates and one guess left so this advantage should be low whereas 10 would be good because whether or you know even if this is wrong then we have one guess left and we could go either higher or lower.

49:16 So a lot of reinforcement learning algorithms are derived from different ways of estimating these advantage and value th this advantage as a function of this individual Q value for actions and values and we'll talk about a few of those for the rest of this lecture. But a key idea is that if we don't know qpi or vpi, we can estimate them by we can estimate the expected reward of this particular action just by acting and we can estimate this value of how well we would expect to do using what we'll call a baseline function.

50:00 We'll see some examples of baseline functions for the next few slides. So a baseline is just a way of rescaling your rewards. So this is the reinforce estimator of the expected return here. which you know scales the reward by the log probabilities of which uses the reward to rescale this gradient of the log probabilities. And these advantage-based methods will generally have this form here of subtracting off a baseline from this reward.

50:42 And these are different these are not numerically equivalent to each other but they are both different valid ways to estimate the policy gradient. We'll have a slide next that shows that this is also a valid unbiased estimator of this same quantity. But for different like baselines that we can define we'll get different reinforcement learning algorithms. So we'll have an example coming up which kind of motivates why subtracting a baseline can be a good thing to do.

51:15 so and in that example the baseline's just going to be a constant. The baseline could also be a running average of all the rewards that you've gotten so far from all the examples that you've seen. It could also be a trained like neural model and this is the case in actor critic methods or a method called PO which we may cover in the next lecture. but you need a separate model from your policy for that that can be expensive. And in the final part of this lecture we'll actually not have a separate model but we will just estimate this expected reward under the policy by sampling from the policy itself. and that'll be the gpo algorithm.

52:00 Okay. So, any questions about this? We'll see a little bit more motivation for baselines in the next few slides, too. So, just a quick note about about baselines and bias. So both of these terms are estimates of the underlying expected reward.

52:33 And an estimate could have a bias. Estimates have bias and variance, right? A bias is like how much is the mean of this estimate potentially shifted from the true value. And the variance is like how much spread is there in the values that I'm drawing for that estimate. So you might be worried about bias in your estimator. And there's this this lema that shows that subtracting a baseline which depends only on the history up to this point which all of these baselines that we've defined here do does not bias the gradient. So the way to see this is this is our advantage our baseline based estimate of the expected return and we can decompose this the expectation of this into these terms here. So using just the associative property putting this here that's the definition of the partial derivative that's the that's the the gradient for J minus this expectation here of B * this and this is going to be equal to this if this expectation of this term here is zero And we can see that by just taking this term here and writing out the expectation like this. it's the expectation over actions sampled from this policy. and bringing out this baseline term which does not depend on anything other than the history up to this point. We just get the sum over actions of the probability of those actions.

54:26 And that's going to we can use this same log derivative trick and rewrite that as this. And then this will just become one and it'll be a bias term times the partial derivative of one which is going to be zero. So this tells you that you can subtract a baseline which depends only on the history of actions up to that point, actions and observations up to that point without biasing your estimate of the policy gradient.

55:04 But let's look at an example of how these like baselines can actually reduce the noise in your estimate of the gradient. And this is going to actually be really important because a huge amount of the work in reinforcement learning is reducing the variance in your estimates of the gradient so that learning is stable and robust to all the variation that you have from sampling from the policy and from the environment. But to motivate this we'll look at just a very simple example where you just have two actions.

55:39 there's no multiple time steps. There's just a single decision that you make and two actions. You're either guessing a or b. And we'll have just a very simple policy which only has one parameter theta. And we'll get a probabilistic policy from this by applying the sigmoid function to theta. So sigmoid is this squashing function, right? Which maps a real value from negative infinity to infinity into a value between 0 and one. and we'll take that value as the probability of choosing a.

56:15 So say that for our current policy we have this theta is negative 1.1 which will give us these very convenient probabilities of 0.25 for guessing a which has reward of one and probability 0.75 of guessing B which has probability zero. Is this a good policy or a bad policy? Think about how the probabilities correspond to the rewards.

56:48 This is bad, right? Because it puts higher probability on the option that has no reward. So we want our policy gradient to update this parameter theta. Which way do we want theta to get updated? Yeah, we want it to become more positive, right? Because if theta is larger, we'll be farther up on this curve, we'll get a higher probability for guessing A, which is the option that has higher reward. So let's compute our gradients without using a baseline and see if policy gradient tells us the right thing. So we can compute the expectation we can compute the values that we we would get as estimates of the policy gradient from each of the actions that we could take under the policy that we could sample.

57:43 So if we sample guess A, what's our estimate of the gradient? Well, we get this by multiplying the reward by the partial derivative, which using properties of sigmoid, the partial derivative here is 75. So the gradient is just the product of those, it's 75. We can do the same thing. what's our estimate of the gradient if instead of sampling A, we sample B? That because the reward is zero actually gives us a zero gradient back. So this is kind of looks like the trajectories that we saw before where we didn't get any gradient from ones with reward zero. Here we're not getting any gradient if we sample this action which has reward zero. But if you take kind of the exact definition of the policy gradient which is the expectation over all the actions in this setting where there's only two actions we can compute that exactly. we can just multiply each of these gradients by the probability of sampling that action and that'll be like we'll see the exact number in a couple slides but I think it's like8 so the actual policy gradient the expectation over all the actions is8 and but we will only we'll have kind of a very noisy estimate of that where 0.25% 25% of the time we will get a sample of 75 which is much higher than 0.18 and then 75% of the time we'll get zero as an estimate for it. And what we really want is to have both of those estimates that we could sample be much closer to the true value.18.

59:26 So a baseline actually does this if so the modification from the past thing is we're just subtracting a baseline from our reward here and otherwise the computation is the same and it works out to be that if we sample guess A we now have an estimate for the gradient of 375 and if we sample B we have an estimate for the gradient of.125. So the expectation multiplied by the probability of these actions is still going to be like.18.

60:00 So it's going to be the same as what we had before. But these individual samples are much closer to that true value because the baseline actually reduced the variance in the samples that we get. We'll have a visualization of this in another slide. But the key thing here is that that expectation does not change. It just kind of like makes the the variance in the estimator more narrow, which is a good thing because you're going to be drawing these estimates on real data, doing updates, and you don't want your gradient to bounce you around a bunch more than you need to.

60:35 Okay. Yeah. So, here's just a visualization of this. if we have a baseline of zero, here's our gradient estimates. They have this expected value, which is8. And as we drag the baseline, we can see that we're more we get these estimates for the gradient which converge on this on this true value for the expectation of the gradient. And in all cases we are going to be doing an update which makes the model more correct. We're going to be doing an update which shifts theta to be more positive which is going to increase the probability of a put more reward or put more probability on the thing with reward one. But we're just going to be doing that more quickly if we have lower variance here in this estimate.

61:30 And here's a graph that shows that. So the purple line here is going to be doing draws of the estimate for the gradient without using the baseline. So no B term just advantages sorry just scales of one or zero. And then the green line is going to be estimates that subtract a baseline just a very simple baseline of 0.5 which was you know which looked like this on this plot. So it's moved them a little bit closer to this true expectation.

62:03 So we'll just like draw some samples and this simulates like doing steps of policy gradient on the exact same example over and over again. So here are sampling in this case we sampled guess A. So we got positive gradient and we got a much larger gradient if we weren't using a baseline. So we'll be taking a very large step and we sample A again. And this line here is going to be our running average over all of the samples of the gradient so far. So when we sample B, now we have a sample of the gradient that is below the true value. So we're going to be closer, our running average is going to be closer to that true value. But you can see here we're we have less variance when we have the baseline. So our running average is already closer to that true value.

63:00 And this continues like this. We happen to sample a lot of guess B's in a row. We sample guess A, it goes back up again. And it does converge eventually, but we're generally closer to it when we've reduced the variance with this baseline. Any questions about this? Okay. Hopefully this wasn't too in the weeds, but like baselines have always been a little bit mysterious for me and I wanted to make these visualizations also to convince myself how they work.

63:29 hopefully they're helpful for you too. The intuitions for this very simple two armed bandit setting do transfer generally to these more complex you know sequential decision-m settings too. Yeah. So baselines are helpful because they can mark a trajectory as better or worse than expected. And if you've chosen a good baseline, a lot of reinforcement learning is about choosing good ways of representing this advantage function. It can help to reduce gradient variance and and stabilize your learning. but you have some downsides too and a lot of choices about what you use for the baseline.

64:10 So with and I'll skip over this other than just to say if you are working in a setting that gives you rewards at intermediate time steps rather than just at the end of the trajectory you'll you'll replace your reward scaling with the reward to go like the sum over rewards over future time steps and this will also help to reduce the variance. we'll have this slide here and I we can also talk about this offline if it's something that's relevant to you. but it's how you would handle not having all of your rewards come at the end of the task.

64:52 Okay. So, with the last bit of time that we have, we'll cover a simplified version of GRPO which is group relative policy optimization I believe. and it's a way of e of getting this advantage function of estimating a baseline that uses rollouts from the model itself. and this was this was actually introduced pretty recently. So it was in one of the papers from deepseek that was training reasoning models with a big focus on efficiency. So a lot of past work on applying RL to language models required a separate model to predict this baseline. Predicting baselines to you know reduce the variance is very important but was challenging and people were using separate models for it. But DeepSeek was very interested in efficiency and training models without using much compute and having a separate model takes a lot of compute. So they proposed this approach which gets rid of the separate model and just uses multiple samples from the model itself to estimate the expected value the expected reward under the policy which is really kind of elegant. so the idea is that we can sample multiple times from the model for a given instance and then get the reward on each of those. and then we can center that reward by taking the mean of the reward across all of these instances in the group. We get a mean. We basically use that mean as our baseline. And this makes sense, right? Because you know the baseline is your expected a good baseline is your expected reward under the policy. And a really good estimate of that is just running the policy multiple times and getting the reward taking the mean of those. So we get our centered reward here. And in the original GRPO paper, they're also standardizing this. so they're not just shifting it, they're rescaling it. they compute the standard deviation over the rewards in this group of trajectories and then they divide by that standard deviation to get the advantage functions. Now there's like actually it seems like dividing by the reward is not necessary and a lot of implementations of that that call themselves GRPO actually get rid of this division by the standard deviation.

67:31 So it more closely matches the advantage term. there's a paper we'll talk about in a few slides called DRGP Dr. gpo or done right gpo which talks more about the motivations for that but actually in a lot of implementations of gpo they won't divide by the standard deviation they will just subtract the mean but yeah the gpo update is very simple looks a lot like policy gradient that we defined before but so you have like a task instance which is you know like one RL problem to solve you will generate multiple rollouts from it. So you will solve the same problem multiple times, resetting the environment, getting a full trajectory by sampling from your agent and this will typically Aorva in his lecture will give some guidance on like choosing values for this group size I believe, but this is typically something like eight. and you'll get the mean of rewards over those rollouts and you'll get your advantage estimates from that. we'll show two different ways of getting those advantage estimates in another slide, but that gives you a group of rollouts for one trajectory along with advantages for them. If you have room in your GPU to fit in more examples, you can also get groups of the same size for additional instances. And you'll be computing group means only within each example.

69:05 So this is important right because the value function should just be an estimate of how it should be an estimate of the reward on this particular example. so our advantages should only look at the reward within this example. But you will get an advantage for each trajectory by scaling the reward for that trajectory using the the formula that shifts by the mean and maybe divides by the standard deviation of the rewards within that instance.

69:39 So for policy gradient loss we were just taking a single sample for each for each instance and then rescaling. But in gpo we are going to be having multiple instances in our batch and then multiple trajectories within the group for each instance and we'll get an advantage for each of those. We'll use that to rescale our log probabilities before differentiating. Any questions about this? Yeah.

70:12 >> is this unbiased? Yes, because you your advantage function here depends on nothing except the prompt and the model. So your baseline function here which you're subtracting from the reward is just rollouts from the language model. So this is unbiased. Okay, cool.

70:45 Yeah. So, Dr. GRPO makes two changes to to GRPO. And I should note too, GRPO also, so what what I wrote here is not the loss for GRPO as defined in the paper. the loss in the paper has some additional terms which I think Graham will talk more about in his lecture which also like clip the gradient so that you don't take steps too far in one direction or another and also regularize the policy so that you don't deviate too much from the policy that you started with. Both of those are additional things that help stabilize training. We're not going to talk about them today. We're just talking about the advantage functions within GRPO. but yeah, this will help you get started on assignment three when that comes out.

71:34 Okay, great. So, we're good on time. so, DRGRPO makes two changes to GRPO. and you will also do these in your assignment. and they seem generally helpful to do because they help avoid some biases within GRPO. And the two changes are just the first one we've talked about a little bit already. It doesn't divide by a standard deviation. It just subtracts the mean of rewards within the group. And yeah, the motivation for that is that if you're dividing by the standard deviation, your advantages are no longer aware of the difficulty of the problem. because you're sort of like rescaling all ranges of difficulties to be standardized. so that can actually be bad. GRPO also divided by the number of tokens generated which can bias the model towards towards long wrong answers and so DRGRPO just divides by a constant which is the max number of tokens that could be generated.

72:46 Okay. So, same figure that we saw before. How does this actually look when you're implementing it? You have your generated tokens in blue. you compute an advantage function which is going to take your reward for this trajectory and just subtract off the mean reward from all of your other trajectories. And then you get a per token reward which just normalizes that by a constant C, which is typically the number of tokens that could have been generated.

73:16 And you do this for all of the rewards in your all the trajectories in your group which could have varied rewards positive or negative. And the loss looks like this. Any questions about this? Yeah. And yeah, so just to note again the original paper also has like a a ratio which helps regularize the policy and clipping, but we won't talk about those for now.

73:46 Okay, so we had a question before about like what if all of your rewards are zero? that's bad because if all of your trajectories fail, that means it's just probably too hard for the current policy and you're not going to get any signal from that even if you have a baseline. you also have a problem if all of your trajectories succeed because they're too easy for the current policy. So like you know if this is binary rewards and one is the best that we could possibly do it's bad and generally it's bad if we don't have any spread within our examples because then we won't have any signal from the advantage function as estimated by gRPO.

74:33 So what could we do if every you know sample trajectory fails? We could choose different task instances. We could inject some some knowledge into the prompt that encourages the model to produce more guided solutions that might be correct. We could add demonstrations. there's ways of doing off policy RL that put like correct demonstrations into the group. or more simply, you could just fine-tune the model first on those correct demonstrations before running RL. Or you could have a more informative reward which encourages the model to make correct intermediate decisions before producing an entire trajectory. That's called reward shaping. So the like giving a model reward for finding the correct file could be an example of that.

75:23 we'll skip this because we'll potentially talk more about PO later on. But PO basically just instead of using the trajectories to compute the advantage function has a separate model which usually shares a backbone with the policy to predict the the baseline. Okay. So very quickly just comparing the methods that we've seen so far. every method has the agent take actions in the environment and then we'll just compute a loss using those actions produced by the agent. All of these look like selftraining in some way but they just differ in the way that they rewe those generated actions. So restm we assume that we have binary rewards of zero or one and we just filter to the trajectories which are correct and then do sftt on those. reinforce uses the reward of the trajectory to rewe the loss which corresponds to reweing the gradient which we spent a lot of time showing is correct. And then GRPO uses these advantage functions where the reward is is shifted and potentially also rescaled by the reward of other trajectories sampled by the model for the same example.

76:50 And all of these methods help to address the three problems that we had at the very beginning of the lecture. So task mismatch should now be solved because we're now maximizing reward directly. expected reward directly data mismatch. We can now learn from trajectories that have reward zero as long as we're doing baselines in the appropriate way. And also the on policy nature of these methods helps with exposure bias because the model's like essentially driving at training time. So it should be learning to recover from its own errors. it just comes at the expense of you know you have to actually just try stuff in the environment which can be costly can introduce some noise into your gradients and we've covered some ways of trying to make that process more robust but in the next two lectures starting next week we'll we'll talk about some others too okay that's all we