Deep RL Bootcamp Lecture 4B Policy Gradients Revisited

Andrej Karpathy walks through policy gradients using Pong as a testbed, showing how to train a 1.3M-parameter neural network on raw pixels. He covers the core trick: weighting action log-probabilities by advantage signals, then implements it in 130 lines of NumPy, training the agent in two days on a MacBook.
Chapters
Hello, I'm very glad to be here. My name is Andrej Karpathy. Until recently, I was a Stanford PhD student, then I went to OpenAI. I was there for about a year and a half, and I worked on deep reinforcement learning and unsupervised learning. And as of 2 months ago, I joined Tesla where I work on Autopilot. So kind of went much more from the direction of very academic to very applied. Unfortunately, at Tesla, I don't get to apply some of this technology that you're hearing about too much because even though we have a huge fleet and we could do lots of kind of parallel collection of episodes and rollouts, unfortunately you can't explore it all, right? Because I wouldn't want to run anything epsilon-greedy on the car. So unfortunately I didn't get to apply this technology, as much fun as it is. So in this lecture I'd like to talk a bit about policy gradients again, but I'd like to especially focus on the intuition behind policy gradients and why it works and how it works, and especially how it's different from supervised learning, or rather I should say how similar it is to supervised learning. And also I'll give you a bit of a kind of concrete idea of how you would implement this in NumPy, kind of in raw Python. So we're going to be talking specifically about Pong. So hopefully I don't have to explain the rules of Pong. We're going to be playing against some kind of an AI opponent, and we're going to be in control of one of the paddles, and then there's a ball and we're trying to get it past the opponent. We get -1 reward if the ball comes past us, or we get a +1 reward if we get the ball past the opponent. Now for our purposes, of course, this is just a Markov decision process. So that every single, every single configuration of these two cattle and the ball is a state in a Markov decision process, and then we have actions that we can take, and in this case we're only going to consider a binary action. Am I going either up or am I going down? So there will be two arrows coming out of every single state, and for some of the transitions we experience a reward. So the reward in this case is every time 0, except for in the final terminal state. So the way we're going to approach this is we're going to define a policy network, and So the policy network is going to take the raw pixels of our Pong, and we're going to have a multi-layer perceptron here, in this case just with a single hidden layer, and we're going to have a hidden layer of some neurons that are fully connected to the raw pixels of this input, and then they're going to be connected to a single output node that gives us the probability of going up. So again, as kind of Peter brought up and also Rocky before him, we always work with stochastic policies because stochastic policies allow you to do exploration. and also they make the problem, the optimization problem, smooth. So we're only predicting the probability of moving up in this case. So just to make this very concrete, suppose we're working, for example, with 80x80 arrays which contain the raw pixels, and this is grayscale instead of RGB. So we have an 80x80 grid that is binary and tells us the position of the paddles and the ball, and we're going to be plugging that into the neural network. The neural network might look something like— as follows. So this is just a simple multilayer perceptron. We take the pixels of this grid, we stretch them out into a single giant column. So this is a column of 6,400 numbers, 80 times 80. And then we do a matrix multiply with some weight matrix W1 to get a hidden state. We threshold every value below 0 to 0, that's the ReLU nonlinearity. And then we have another matrix multiply here of W2, which is again parameters of the network, to get our probability, or I should say logit, corresponding to going up, and then we apply the sigmoid function here to squash this to a range of 0 to 1 so that we interpret that as a probability of moving up. So very simple kind of vanilla multilayer perceptron here. Now just to make this kind of very concrete and get an idea of what's involved in terms of the optimization, effectively suppose we have on the order of 200 nodes in the hidden layer, then we're looking at 200 nodes each looking at 80 by 80 grid of pixels, and these are the biases, and then in the second layer we just have 200 connections to these neurons and a bias. So we're looking at about 1.3 million parameters with this kind of a simple network, which is actually, like, quite substantial, I suppose. 1.3 million is quite a bit. And of course our task is to find the setting of these parameters so that that network plays Pong very well. And I should mention one more thing is that you can't actually play Pong just from a static frame, so— because you kind of need to have some motion information. So what we can do in that case is we can either concatenate two of these frames or we can, for example, take a difference of the new frame from the previous frame or something like that. Okay, so one thing to kind of point out here is that this is actually a really difficult optimization problem, right? So the neural network does not see this. It is not used to objects like paddles and balls that move around. It sees 6,400 numbers that change over time, and these are just like bits flipping in random locations because we don't actually have even a convolutional network, just this straight-up multilayer perceptron, so we've destroyed all kind of spatial information. And then we get a reward of +1 or -1 kind of some of the time as this thing kind of like jitters around and bits are flipping at random. You're taking actions at random in the beginning and sometimes you get a reward, and you need to find a setting of 1.3 million parameters. It's kind of a daunting task and almost a little bit surprising that we can get it to work. So one thing I'd like to also point out is that in this setting where we're working with games and we have a simulator and we have an evaluation criterion, In fact, the solution to these problems is in fact kind of trivial if you're willing to pay with compute. So in fact, we do have a very general learning algorithm that can learn any single game, and the learning algorithm is you repeat lots of times, you sample a random policy and you run it for a while, and then you see if it works, and you keep track of the best one you've seen so far. This is like the best learning algorithm. As long as the solution is in the space of neural networks, you're going to find something that works. And then you return the best policy. So the only thing I'd like to point out is really we're in a business of trying to make this more efficient because ultimately we do have the solution, which is random search. You don't have to do random search completely at random. You could also do stochastic local search and things like that. So the problem with this is we have a solution, but it's of course extremely inefficient. So here we would have to pay with a lot of computation, and then if you want to do a lot of compute, you're using a lot of energy, and then you're basically burning forests. So, This forest— you can have a really good Pong policy, but unfortunately you have to kind of sacrifice a forest or some of its inhabitants, and of course we don't like that, so we're going to talk about something more efficient, something more environmentally friendly, which in this case will be policy gradients, and we can both optimize our Pong agents but also preserve the species of this earth. Okay, so let's talk now about policy gradients and how we can make this more efficient. So before I dive into policy gradients, I'd really like to tie this to supervised learning and kind of step back and ask the question of, well, the core problem really is that we don't know how to act. Like, we have some policy that we sample at random, but really at every single position we don't actually know what to do, up or down, in the beginning. So suppose I actually did tell you, like, at this state the correct thing to do is to go up, at this state the correct thing is to go down. So we actually have a labeled dataset, someone gave it to us, And in that case, we want an agent that can perform well in this task. And if we knew the correct answer for each state, then what we do is we just train a neural network, just like you would train a neural network on ImageNet. We use supervised learning. We plug in the images here. We get the probability of going up, and then we know the correct answer in every single case. So we're just going to maximize the log probability of the correct action to take, here y's, so these could be up or down, given our image. So if we were just told what to do at every single state, this is almost like a trivial problem. We maximize the log probability using supervised learning objective. Now, the problem for us, of course, is that we do not know if we should go up or down, and this feels like we're kind of stuck. But we're actually going to do a very simple trick to get reinforcement learning to work, and that is the insight, really, of reinforcement learning is that what we're going to do is we don't know what to do, so we're just going to try a bunch of stuff, and we're going to see what happens, and then we're going to do more of the stuff that worked. Okay? So we're going to start with a random thing, we're going to do some stuff, and then whatever stuff happens to work, because we have the rewards, we're just going to make that stuff more likely in the future. We're going to iterate this until we succeed. So let's try to make this more concrete, otherwise this is quite simple and straightforward. So what we're going to do is we're going to take our policy, we initialize it randomly, and now what we're going to do is collect rollouts, right? So we're going to run a single episode of the game and we're going to collect a rollout. So take this policy, initialize random, and then we start in a starting state. We get a probability of going up. We flip a coin, a biased coin I suppose. You either go down or up. So in this single rollout, we happen to go down, up, up, down, up, up. We have a sequence of states and then we ended up with a win in this case. This would be really unlikely if you were to start with a random policy, but let's just suppose that in this particular case we've succeeded. So we're not only going to do a single rollout, we're in fact going to collect an entire batch of rollouts. So we're going to run the policy, collect an entire batch of rollouts, and some of them we're going to win and some of them we're going to lose. So we basically have a gradient. It's— we have some information telling us where something happened that was good and where something happened that was bad. So in particular, we're not quite sure what happened here. We took a bunch of actions, but in these two rollouts we won. So something that we did here was good, and conversely, something that we did here was not very good. We're just not quite sure what that was. So here's what we're going to do, basically, is we're going to basically do supervised learning. So this set of rollouts here is our fake dataset. We've collected a fake dataset by rolling out the network, and now we have a dataset of the pairs of what was the state and what did I do, and now we're just not quite sure what was— and these are kind of like fake labels now. So the actions we happen to take are now fake labels for our dataset that we collected by doing rollouts. And now what we're going to do is intuitively for the episodes where we won, we want to take all those actions that we've taken and we're going to increase their log probability, just like we were doing supervised learning. So we're going to increase log probability of the actions where things worked, and then here in the cases where we did not— things did not work out, we're going to instead of increasing the log probability of those actions, we're going to decrease them. So we're just gonna maximize now the negative of the log probability of of those states. And so really what's going on in— and this is policy gradients. I mean, that's pretty much it. So just to summarize, in supervised learning, we are given the correct labels, so we can just maximize the log probability of the correct action given the images. In reinforcement learning, we don't have the labels, so we instead roll out our policy, our current policy, and the actions we happen to take during the rollouts, those become our fake labels. For this little small dataset that we've collected for ourselves. And then what we're going to do once we collect that batch is we're going to maximize this expression. So it's the log probability of the action we've taken exactly as here, but now we have this additional scalar multiplication here, and we want this AI intuitively to be high if we want to encourage that action in the future, and we want that AI to be small if we want to discourage that action in the future. So this number here, I'll just call it advantage. It has a bit of a more kind of mathematical precise sense, but in this case advantage is just a scalar, and depending on how a rollout turns out, that scalar can be high or low, and then that's going to modulate whether or not we want to increase or decrease any particular action that we've performed in the past. So positive advantage will make that action more likely in the future. Negative advantage will make it less likely in the future. So I'm going to go into some of the bells and whistles of what's involved and some additional pieces on top of this, but at this point I'd just like to maybe see if there are any questions on this point, because we're just doing supervised learning. The labels, or fake labels I suppose, are the actions we happened to do during the rollout, and we're scaling that objective based on the advantage of how things turned out. Any questions? [Speaker:AUDIENCE] But things were going well for a while, and then you made some mistakes. [Speaker:ANDREJ KARPATHY] Yeah. [Speaker:AUDIENCE] You can do a capture of that, right? [Speaker:ANDREJ KARPATHY] Yeah. [Speaker:AUDIENCE] You get a bad— [Speaker:ANDREJ KARPATHY] Correct. So right now the advantage— correct. So what you're pointing out is, We may have done some things correct and some things incorrect and some things correct, and how is this working? So right now the advantage that we're using, this AI that I put here, we're using that to be +1 or -1 for the entire rollout depending on the final outcome in this simplest example. When I go to the bells and whistles, we're going to see how we can do, for example, discounted rewards and things like that. That's going to help us try to disambiguate between those. But for now, we just have the advantage will be distributed equally to all actions. You're equally blaming all the actions in the entire rollout. And the reason this will work to some extent is because suppose you have millions and billions of these, then if there are some actions that usually are predictive of you winning at the end, eventually they'll kind of win out through brute force, and they will be encouraged more than discouraged, and so that's how this will actually all average up in the network training, and it will kind of work out. Go ahead. [Speaker:ANDREJ KARPATHY] Yep. Yes, so some of these states could be the same states. Yeah. Yeah, go ahead. [INAUDIBLE] You get a bunch of loses here in the batch? Oh yeah, let's say like you do it like 100 times and they're all lose because of the game order. That's right. Would that be still— Yeah, so in this case, if you— so we're going to— in the bells and whistles, I'm going to describe how we modulate that advantage. In this case, if we were to adopt the simple setting of just +1 if you won and -1 if you lost, then Yeah, what's going to happen is you're going to discourage everything you've done, and so I don't know what that does to the neural network. I mean, everything you've done throughout will just be discouraged into future, so you're just increasing the probability of the other thing at every single state, which is maybe the correct thing to do, right, because everything you've done, nothing worked, so you're kind of— your network will search through different parts of the state space because you've decreased those probabilities now. Yeah, go ahead. The advantage is kind of like anything. It's some kind of a scalar that tells you how much you want to encourage or discourage the action that you happen to take. Yes, so there are some principled ways of choosing the advantage, and I'm going to show you one example, and Peter has shown some examples of the advantage. Like, for example, you can take the the empirical return and subtract a baseline. Any of those— there are multiple ways of using advantage. Yeah, go ahead. [INAUDIBLE] Yep. [INAUDIBLE] Yep. [INAUDIBLE] Yep. [INAUDIBLE] So the question is, how do you deal with continuous-valued action spaces? So we're just increasing and decreasing log probabilities. So as long as you can assign— as long as you can have a distribution over the continuous space, suppose you use a Gaussian and you have a Gaussian policy, so you predict the mean and the standard deviation, and now you sample from it. And now what we're doing is we're increasing or decreasing the log probability of that sample depending on how things turned out. So we've sampled— we have a Gaussian, and now we sample something off the mean, of course. And now if this thing worked well, then the log probability— what would happen is you would nudge the mean towards the value that worked. And if instead this was a bad sample, then by decreasing its log probability, you would nudge the mean in the opposite direction, and you're less likely to sample that thing again. So basically there's no difference. As long as you have a probability distribution over the action space, this formula just applies in a straightforward manner. Yeah, go ahead. [Speaker:AUDIENCE] Would this have really particularities being part of the state space that for older drones Uh, with discontinuities or, um, discontinuities? [INAUDIBLE] Discontinuities in, uh— [INAUDIBLE] So in this case, we have a fairly smooth optimization problem because The neural net has these parameters. If we nudge them a tiny bit, it's going to change the probability distribution, but hopefully kind of your expected behavior is kind of similar. And in fact, TRPO, one of the things that we're going to cover in a bit, it explicitly tries to make sure that when you make a small parameter change in your network, you're not actually changing your distribution of your actions for each state too much, because we want that smoothness. [Speaker:AUDIENCE] How does that attribute feel though? How do you decide which applications to respond? [Speaker:ANDREJ KARPATHY] So yeah, so as I mentioned, in many of these cases we really have a huge amount of data. So we start off random, and we stumble into solutions that happen to work. And then if they worked, then we're going to increase them. And next time around, you're slightly more likely to sample them. And so we continuously move this distribution of our probability of actions, and we just continuously do something that is more rewarding. [Speaker:AUDIENCE] So you mentioned the neural net 9 times, but in the second one you do it 7 times? [Speaker:ANDREJ KARPATHY] Yeah. Yeah, so in this example, that would be the case. And you can get more complex, and I'm going to get into some of the bells and whistles. But at least for the purposes of this slide, unless there are more questions, the rough idea is we're doing supervised learning, except we don't have a data set. So the fake labels come from the actions we happen to sample. And then the only difference is in addition that's the advantage, which tells you if we should actually be increasing the probability, which in supervised learning case is always the case. But in this case, sometimes things don't work. So actually we want to decrease their probability. When am I supposed to end, by the way? At 3 o'clock? [Speaker:ANDREJ KARPATHY] We're running 15 minutes behind. [Speaker:ANDREJ KARPATHY] Oh. [Speaker:ANDREJ KARPATHY] We're at 3:15. [Speaker:ANDREJ KARPATHY] OK, so maybe I should— is it a really good question? Okay, it doesn't look like— doesn't look like it's a really good question. Okay, great. It's a really good question. Okay, go ahead. Yeah. Yep, so we take 6,400 numbers at every single dot here. We put them through the network and we get a probability of 0.7 of going up, and then we sample from that distribution. So we are likely to sample a 1, and in this case 1 would be up, for example. Yeah, and that happens at every single one of these states. We saw a particular image, we take a particular action, and now we're gonna wait out and see if it worked or not, and depending on that, we're gonna increase or decrease it in the future. [Speaker:AUDIENCE] [INAUDIBLE] [Speaker:ANDREJ KARPATHY] So in this case, because of backpropagation, if we don't have a recurrent neural network, we're not actually backpropagating through the rollouts. Suppose we would just be truncating every one of these independently. These are all kind of like a single— every arrow here is a single example. But if you had a recurrent policy, then this would be— you would have to backprop through that. OK, cool. So let me show you some other stuff. So one additional trick that kind of was brought up is it's kind of odd to assign equal blame to every single— whoops. Okay. It's kind of not very philosophically appealing to assign equal blame to every single action in the rollout because some of them may have been good and just some of them happened to be really bad, and then we're equally blaming everything, so that doesn't seem very optimal. So what we do often in practice is we do discounting. And it's just a modulation of blame. It's a heuristic for modulating this blame in a time-sensitive manner. In particular, we make the assumption that we blame the action— we blame each action kind of in an exponentially decreasing manner. So here, for example, we took the action down and suddenly we have a reward of -1. Then the discounted reward at this point would be -1, the raw reward, but then this action, the down before that, its discounted reward would be only -0.9. So we're assuming that basically we're using a gamma of 0.9 in this case. So we're discounting backwards, and this action is less blamed because it's further away in time. And so we're just making this assumption that actions have immediate effects most of the time, and then exponentially so they have less effect into the future. And really it's just kind of a heuristic. So in the case of Pong, let's see what happens with discounted reward. So suppose we're playing Pong. This is over time. We start off at the episode, and then suppose here at some point the ball gets past us, so we're going to get a negative 1 reward, but there will be several frames where the ball is kind of going past us, and there's nothing we can do anymore, but there will probably be 10 more frames during which this ball is, you know, on its way, and we're almost certainly getting negative 1 reward. There's nothing we can do anymore. So we get a -1 reward here, and what you'll see with the discounted reward now, and this would kind of factor into the advantage now, is we blame this action a lot, we blame all these actions, but really these were all kind of useless actions, and the real important actions happened here because we missed the ball. So in this case, this is a bad heuristic almost. So things would— so this heuristic kind of breaks down, but on average usually this this is a good idea in practice, because your effects usually do have an immediate impact. And if you were to use a value function, then what would happen in this case is if you have a value function estimator, then it could in this case learn that when the ball is past you, then the value function would immediately tell you that, okay, you're about to lose. Your value suddenly becomes -0.9 and converges all the way to -1. and basically the value function would tell— would learn eventually that everything is hopeless, everything you've tried, the expected return is always just -1, and then, and then what would happen is the algorithm would kind of learn to anticipate that, and you can never change that. So now you have the empirical return is always equal to the value function, so when you subtract that baseline, then you're just getting 0. and then eventually you would not be introducing this noise of I tried lots of stuff here and nothing works, and it's not going to contribute any gradient in that case, because in that case our advantage is the empirical return minus the baseline. So if that is 0 because your baseline is now correctly predicting that you are about to lose, then you'll get no— you'll get 0 in there, and then that 0 is going to multiply your lock probability, so eventually you'll stop learning and you'll stop introducing random noise into your gradient estimator. Okay, so at this point I'd like to briefly kind of go through a piece of code just to make this even more concrete. So there's a pointer— this is from a while ago now, but I at one point wrote a 130-line gist on— the only dependency is NumPy, and I'd like to briefly take you through it. It basically runs policy gradients on Pong, and it's a complete implementation. So let's see what this looks like. I'm going to put it here on the left, and we're going to briefly step through some of the— what's involved with this concrete example. So in the beginning we do some initializations here. This is not really important. Nothing crazy is going on here. We initialize a gym environment using OpenAI gym framework, and we initialize some stuff, and then we start a training loop. So while True, we're going to iterate now the learning algorithm. So in the beginning we're going to— we're going to take a frame from gym. So this is an observation that we're going to preprocess into our vector. And so the preprocessing function, as you can see, it just kind of like crops out unimportant stuff. I downsample, so I only take every odd— every even pixel, and I make it binary, and then I return as float. And then in the preprocessing, right after the preprocessing, the actual observation that we're going to put the network on is going to be the current preprocess version minus the previous preprocess version. So we're trying to incorporate some of the temporal information in there, so we're going to really see a difference image instead of the raw image so that we have some motion information in the network. So now what we're going to do is we're going to— we've preprocessed the image, now we need to put it into neural network. So the next line of code is going to call our policy forward on this preprocessed vector. The policy forward is the multilayer perceptron that I've described and that I've put up before, and we're going to return the probability of moving up, and we're going to return the hidden state. The reason I'm returning this hidden state is because this is no— we're not using any neural network library, and so I need the hidden state because later we need to cache all of this because later we're going to use backpropagation. So I need to know what the hidden state is and I need to kind of cache it so that I can do backpropagation properly. And then once we get our probability, we're going to flip a coin, and with that probability we're either going to take action 2 or action 3. This is because in the gym simulator action 2 is up and action 3 is down. And, yeah, and here's the sigmoid implementation. So basically we forward the neural network. and then we cache a lot of stuff. Now here's kind of like where this gist, I think, becomes most complicated, and this is, I think, where— is that we've taken that action, and I'm actually like doing a tiny piece of backpropagation in the second— in the next line. So this is a tiny piece of backprop and bookkeeping for backpropagation. Now if you are familiar with neural network frameworks and you understand that PyTorch is the superior way of building neural networks, then— and if you were using PyTorch for this, for this example, then you would not have to do any of this bookkeeping. There would be no lines of code here because I'm just doing this so that I can backpropagate later. And in particular, what I'm doing here is I'm trying to compute the derivative of the log probability of the taken action, the one that we took, with respect to the output of the network before the sigmoid. So really, I'm trying to evaluate the log probability of the action that we took given this image. I'm trying to just evaluate this tiny piece here without the advantage scaling. So just this piece, the log probability of the taken action, and this is, this is us forwarding the S, is kind of the output of the network. Then we put that to compute the sigmo— after sigmoid, that's a probability, and then we sample the label. So we sample from that probability, and basically if you do some of the math, and we really have a small logistic regression classifier here, it's a small binary classifier, The loss function in this case is— I'm not sure how much time I have actually to go into this. So if you just think through this just very briefly, if we suppose that p comes out to 0.5, so we had a half chance of going up or down based on the network. And now suppose that we actually sampled y being 1 so that we're actually going to go up in this particular case. Then what we're trying to evaluate here is the log probability of the action that we actually took and the gradient of that with respect to S so that we can later backpropagate down through the network. So basically we're trying to dlogp of the taken action with respect to S, and the probability of up in this case, so if y is 1 and the probability is then really the probability of up is just log p. And so if you just do the math of log p, dl by ds, because p is a function of s, if you backpropagate, you'll just find that it's 1 minus p. And if we happen to take action 0, then it's going to turn out this way. There's a more compact kind of formulation of this, and this is all kind of mechanical calculus to derive some of these. But really, it's kind of intuitive. So suppose we took up and p was 0.5. Then the log probability here is log of 0.5, and then if we were to increase the probability of the taken action, then with respect to the output of this network, then 1 minus 0.5 will be 0.5, or a positive number. So what's this telling us is that, okay, the log probability was 0.5, and now the action I took was 1, or going up. So if I wanted to increase the log probability of the correct action, then I should increase S. If S went up, positive number of 0.5, then the log probability of the action I took would also go up. So that's the gradient information there. And I would need significantly more hand waving to potentially explain this, to potentially explain this better. Okay, so what we've done so far is we've forwarded the network, we've preprocessed the image, we forwarded the network, we did a tiny piece of backprop here, and now we have the action that we would like to take, so we pass it into the environment, and we get the new information in terms of the observation, the reward, and whether or not the episode is over. So we just cache some of it, and now if the episode is over and the simulator tells us whether or not it's over, then I have some code here that basically does maintenance. And again, if you are using PyTorch or something like that that builds up this graph, you don't have to do this. I have to collect together the entire batch of all the images that I've seen, all the actions I've taken, and all the rewards I've received, and I have to kind of put it together into NumPy arrays and make a batch out of it. Because this is kind of now our training data that we've collected. And now in the next line, I'm going to do the discounting. So what we're going to do here is we're going to discount EPR, which is the episode rewards just over time that we've observed. The discounting does this thing, so we exponentially decay any reward over time. So we take the reward, we discount it backwards, and then here I'm trying to compute a kind of a low-variance advantage. And the way I'm going to do this here in this particular example is we have a discounted reward, and then I'm going to— I want that to be the advantage, but before that becomes the advantage, I'm going to subtract the mean off this so that it's zero-centered. So the advantage, half of it is kind of positive and half of it is negative. And then in addition, I'm going to divide by the standard deviation to put everything to kind of a single scale. So we have the discounted rewards, and then if you imagine just standardizing this so that it's unit Gaussian distributed, really, so we shift it so that in encouraging or discouraging half the actions, and then we shrink it so that we always have the same kind of magnitude of advantage. So if you could just imagine, let me go back here. I have a discounted reward, and then I'm going to subtract the mean, and I'm going to make sure that all of these are unit, zero mean and unit variance. And I'm not doing this across a single rollout, I'm doing this across an entire batch of rollouts. So that's an attempt to decrease the variance, to decrease the kind of— to make everything learn a bit faster. And then, okay, so we have this discounted reward, we compute the advantage, and then we actually modulate our gradient by the advantage, and then we call the backward pass of the policy. And so here we get the gradients of the weight matrix W1 and W2. If we do a small update on w_1 and w_2, then we're going to do the correct thing. The network will be more likely to take actions that increase reward and less likely to increase— and less likely to take the actions that did not work well according to the advantage. And I'm using the RMSProp update algorithm to do the update, and I think that's it. So that's just print statements here. So that's the entire script. So really to summarize, It's 130 lines of some preprocessing code, forward the neural network. I'm doing a bit of backprop inline here. Call the environment, call the environment, keep track of everything, and compute the discounted reward, and then compute the advantage. And once we have the advantage, we can actually do the backpropagation. So we have the loss, and we backprop, and then we get a parameter update. We actually perform the update, and now we iterate the whole process again. So we expect to get better about policy over time. So in summary, we initialize policy network, and then we collect a bunch of rollouts. We increase the probability of actions that work well, and if we ever break out of the loop, then we get to eventually possibly profit. And I'm just going to show you some example here. This is running this code, this gist, on Pong, just to show you that we're successful after the 130 lines. So it's very uncertain. I think in retrospect I should have had 3 actions. I should have had a go down, go up, and stay still, because this thing must go up or must go down. So I almost feel bad for it, and that's why it looks so neurotic. It's like it can't decide, but it's really just trying to go back and forth to stay at the same position, but it learns it fine. [Speaker:AUDIENCE] How long did that take to train? [Speaker:ANDREJ KARPATHY] So this took on my MacBook, I think, 2 days. It's like quite some time. Like, it's not instantaneous at all. And the reason for that is there's 1.3 million parameters. Like, it really is quite a large optimization problem, and we initialize them at random, and it does need to run for a while to actually converge well. A lot of these algorithms in general have to run for a while, which is why I'm so concerned about the environment and all those lines, because, yeah, it does take a while. Sorry. Oh, wait, sorry. I should just point out— oh, crap. Yeah, so that is the end of the talk. So yeah, thank you. And I'll take questions. OK, go ahead. [SPEAKING CHINESE] Sorry? [SPEAKING CHINESE] Oh, yeah. Yes, so the discounted reward, so So you want to get an advantage, right? And so either you are equally blaming every single action, or we're doing this discounting where if you get a reward, then we're going to smooth out that reward across time. And so now I would like to use the discounted reward directly as the advantage. Like, you could do that, but the problem is that— well, there's no real problem. Like, that would actually kind of work, and that was my first way of doing it. And then what I found was what works a bit better is since the discounted rewards, they kind of— Like, the AI will modulate the log probability, right? It will multiply the log probability. So I'm just basically scaling it into a nice range. Like, the discounted rewards might be off mean, might be not zero mean, and they might be kind of all over the place depending on the scale of the rewards. And so I'm just trying to make it zero mean. Yeah, I'm just making it zero mean so that there's equal amount of push and pull, and I'm kind of putting it into a nice range. Like, that's pretty much it. [Speaker:AUDIENCE] Same baseline thing. [Speaker:ANDREJ KARPATHY] Yeah, but otherwise the advantage would be just the discounted reward. You can just think of it that way.
More from Andrej

Skill Issue: Andrej Karpathy on Code Agents, AutoResearch, and the Loopy Era of AI

Andrej Karpathy: Tesla AI, Self-Driving, Optimus, Aliens, and AGI | Lex Fridman Podcast #333








