Dreamer v2: Mastering Atari with Discrete World Models (Machine Learning Research Paper Explained)
deep learning machine learning arxiv explained neural networks ai artificial intelligence paper reinforcement learning deep reinforcement learning dreamer dreamer v2 dreamer rl dreamer reinforcement learning google reinforcement learning deepmind reinforcement learning google ai world model world model reinforcement learning google deepmind world model google deepmind reinforcement learning atari reinforcement learning atari world model rainbow muzero
#dreamer #deeprl #reinforcementlearning

Model-Based Reinforcement Learning has been lagging behind Model-Free RL on Atari, especially among single-GPU algorithms. This collaboration between Google AI, DeepMind, and the University of Toronto (UofT) pushes world models to the next level. The main contribution is a learned latent state consisting of one discrete part and one stochastic part, whereby the stochastic part is a set of 32 categorical variables, each with 32 possible values. The world model can freely decide how it wants to use these variables to represent the input, but is tasked with the prediction of future observations and rewards. This procedure gives rise to an informative latent representation and in a second step, reinforcement learning (A2C Actor-Critic) can be done purely - and very efficiently - on the basis of the world-model's latent states. No observations needed! This paper combines this with straight-through estimators, KL balancing, and many other tricks to achieve state-of-the-art single-GPU performance in Atari.

OUTLINE:
0:00 - Intro & Overview
4:50 - Short Recap of Reinforcement Learning
6:05 - Problems with Model-Free Reinforcement Learning
10:40 - How World Models Help
12:05 - World Model Learner Architecture
16:50 - Deterministic & Stochastic Hidden States
18:50 - Latent Categorical Variables
22:00 - Categorical Variables and Multi-Modality
23:20 - Sampling & Stochastic State Prediction
30:55 - Actor-Critic Learning in Dream Space
32:05 - The Incompleteness of Learned World Models
34:15 - How General is this Algorithm?
37:25 - World Model Loss Function
39:20 - KL Balancing
40:35 - Actor-Critic Loss Function
41:45 - Straight-Through Estimators for Sampling Backpropagation
46:25 - Experimental Results
52:00 - Where Does It Fail?
54:25 - Conclusion

Paper: https://arxiv.org/abs/2010.02193
Code: https://github.com/danijar/dreamerv2
Author Blog: https://danijar.com/project/dreamerv2/
Google AI Blog: https://ai.googleblog.com/2021/02/mastering-atari-with-discrete-world.html

ERRATA (from the authors): 
- KL balancing (prior vs posterior within the KL) is different from beta VAEs (reconstruction vs KL)
- The vectors of categoricals can in theory represent 32^32 different images so their capacity is quite large

Abstract:
Intelligent agents need to generalize from past experience to achieve goals in complex environments. World models facilitate such generalization and allow learning behaviors from imagined outcomes to increase sample-efficiency. While learning world models from image inputs has recently become feasible for some tasks, modeling Atari games accurately enough to derive successful behaviors has remained an open challenge for many years. We introduce DreamerV2, a reinforcement learning agent that learns behaviors purely from predictions in the compact latent space of a powerful world model. The world model uses discrete representations and is trained separately from the policy. DreamerV2 constitutes the first agent that achieves human-level performance on the Atari benchmark of 55 tasks by learning behaviors inside a separately trained world model. With the same computational budget and wall-clock time, DreamerV2 reaches 200M frames and exceeds the final performance of the top single-GPU agents IQN and Rainbow.

Authors: Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, Jimmy Ba

Links:
TabNine Code Completion (Referral): http://bit.ly/tabnine-yannick
YouTube: https://www.youtube.com/c/yannickilcher
Twitter: https://twitter.com/ykilcher
Discord: https://discord.gg/4H8xxDF
BitChute: https://www.bitchute.com/channel/yannic-kilcher
Minds: https://www.minds.com/ykilcher
Parler: https://parler.com/profile/YannicKilcher
LinkedIn: https://www.linkedin.com/in/yannic-kilcher-488534136/
BiliBili: https://space.bilibili.com/1824646584

If you want to support me, the best thing to do is to share out the content :)

If you want to support me financially (completely optional and voluntary, but a lot of people have asked for this):
SubscribeStar: https://www.subscribestar.com/yannickilcher
Patreon: https://www.patreon.com/yannickilcher
Bitcoin (BTC): bc1q49lsw3q325tr58ygf8sudx2dqfguclvngvy2cq
Ethereum (ETH): 0x7ad3513E3B8f66799f507Aa7874b1B0eBC7F85e2
Litecoin (LTC): LQW2TRyKYetVC8WjFkhpPhtpbDM4Vw7r9m
Monero (XMR): 4ACL8AGrEo5hAir8A9CeVrW8pEauWvnp1WnSDZxW7tziCDLhZAGsgzhRQABDnFy8yuM9fWJDviJPHKRjV4FWt19CJZN9D4n
 Hi there. What you're seeing here are predictions by a world model learned for Atari reinforcement learning. On the top you see what really happened during an episode of play. And on the bottom you see the predictions of this world model. The world model just gets five frames at the beginning, which you don't even see here as a conditioning. And then it predicts 45 frames of gameplay. It's astounding how accurate it is, not only in terms of how the game evolves, but also in terms of what the agent will actually do. So the world model, the specific world model you see here is part of the Dreamer V2 algorithm from the paper Mastering Atari with Discrete World Models by Danijar Hafner, Timothy Lilikrub, Mohamed Nourouzi and Jimmy Ba of Google Brain, DeepMind and the University of Toronto. So these kind of world models, they enable you to do very quick reinforcement learning once you have the model, you can use it to imagine yourself playing the game instead of actually playing the game. And therefore you can do much more efficient reinforcement learning. And this paper details how to get an accurate world model for Atari, which was sort of out of reach until now, especially considering that they only do single GPU reinforcement learning. So the result, as you can see here, is going to be an algorithm that is the top single GPU agent right now, competing, outperforming other, so here is Dreamer V2 outperforming other algorithms such as Rainbow, IQN, DQN. And the special thing here is that Dreamer V2 is a model based algorithm, whereas the current or the previous best ones, especially single GPU best ones, were model free algorithms. And you can see the next best model based algorithms were, are not really competitive in Atari, right? This is specifically Atari. So Dreamer V2 is an evolution of Dreamer V1, which worked well for things like continuous control, but Atari still seemed a bit out of reach. So the difference between model based reinforcement learning and model free reinforcement learning is that model based reinforcement learning first learns a model of the world, it learns how the world acts, and then it uses that model to learn what actions to perform, whereas model free algorithms, they simply act in the world and they learn to predict the best actions as they act in the world. So there's your difference. And how does Dreamer V2 do that? On the high level, it has two stages. Stage one is learn a world model from past experience. And then stage two is use that world model, as we said, for reinforcement learning. And the reinforcement learning here is going to be just actor critic learning. Very straightforward. There's a little modification with a pass through estimator. But the real difference is going to be in how the world model is learned. And the novel contribution or the main contribution here is this latent state, which consists of a this stochastic latent state, which other than other world models, which model the latent states as something like Gaussian random variables. This paper models the latent state as categorical random variables. And that turns out to work pretty well for Atari. So that's step one. Learn world model. Step two, do reinforcement learning in the model. So not using any data anymore. And you can repeat those two steps as many times as you want. So you start out with a set of data, then you learn an actor, and then you use that actor to collect more data and so on until you have a really good actor. And the world model is really accurate for that actor. So that's the overview. And it's going to turn out, as we already saw, to beat other, at least single GPU models by quite a bit. So we'll go through the paper through the individual steps and discuss what's new and how it all works. The code is also available. I'll link to it. And the blog post I've shown you here has some more explanatory graphics. If you like content like this, as always, don't hesitate to click like and share it with all your friends, especially the Atari gamers, because they are outperformed, as you can see here. All right. So world models. Pretty quickly in reinforcement learning, as you all hopefully or maybe know, you have an agent that is interacting with an environment. And the agent can... So the environment always provides the agent with an observation, which would be an image in an Atari game. And the agent decides to do one of many available actions in response to receiving the observation. The environment then responds with a reward for that action. So either you die, which is like negative reward, or you collect a coin, which is positive reward, or you win the game, which is like a thousand reward. And it also gives the agent a new observation, the next observation. And the agent, again, responds by performing another action and so on. So you have this cycle. And the goal of reinforcement learning agent is usually to maximize all the rewards that it collects during playing with the environment. And you want to repeat that many times for many episodes to have the agent learn to get as to do the actions that are as good as possible in terms of reward. All right. Now, in classic, let's say classic, in model-free reinforcement learning, one way to do this is to take this right here as you play the game. As you play the game, you collect data, right? So let's assume we collect data as we act in the world. And from this data, we can learn something. So model-free learns from the raw experience. So an episode will always be a series of images, right? And actions you have performed. So here is an image and I have performed action one and then came a next image and I've performed action two. So what classic reinforcement learning would do is it would say, okay, from this transition doing this action, I have gotten five reward. And from this transition in this action, I've gotten negative three reward. So I'm going to have to do this action one more often because it gave me a lot of reward after I observe this thing here, right? The combination of this thing, I need to do action one more. And when I'm in this situation, I need to do action two less and so on. Okay, so you're simply trying to put this image that you get into a neural network that tries to predict action one as often as possible. And you want the same network when you input this next image to not predict action two. So like anything but action two. So that's going to be that's kind of the logic between of the classic model-free reinforcement learning. Usually this is implemented in a sort of an LSTM fashion or it's one way of doing it. So you have an LSTM that tracks a hidden state. Why do you need a hidden state? Because you might not see everything in the image there is, right? This is not necessarily Markovian. So there might be information that you need to remember for a long time, like when an enemy leaves the screen and then comes back, you want to track it. Do you have an LSTM or some kind of RNN and then you want to feed the images into that one by one. And then you simply so with an encoder, which is usually kind of a convolutional neural network, I want to draw it like this. And then you try to predict the here the good actions and here you try to not predict the bad action and so on. So this is a simple classifier. Ultimately, it's an LSTM with a classifier on top. And the classifier simply tries to either predict a class of action one or not or predict anything else. So and you train it via back propagation through time. And that's it. Now, here is a little bit different. So why? Why is this maybe not a good idea? Well, all you have is the signal of the reward for given actions. And that means it is it is fairly hard to generalize in these kinds of things. So when you imagine you have your screen right here and there's an opponent kind of here, there's an opponent here and you are down here and the opponent shoots. Right. You have to move out of the way. You have to move over here. Now, RL is completely capable of learning that. However, take the next situation over here. Now, the opponent is here, shoots and you are down here. You have to, again, learn to move out of the way for a classic RL algorithm. These two things are identity are completely different states. Like this is there's nothing equal about the two. Like this is a completely different thing. And it has to sort of learn by force. Look, in this situation, there, you know, you need to move. And in this situation, you also need to move. Now, given that that is a convolutional neural network, it might after a while learn the fact that it, you know, these two situations have something in common. But in essence, these are two different things. And you have to learn purely from the reward, purely from the fact that you're going to die if you don't move to get out of the way in two situations. And of course, this situation can be replicated all over. However, if you have a world model, right, imagine now we have a world model over here and the world model accurately learns to predict the future. Now we know that, you know, we are here. This is here. Now we can imagine ourselves forward and we're going to see we're going to get hit. And that means we need to go out of the way. So doing this explicitly would be called planning. We are not going to do planning in this paper. OK, we are still going to do the classic RL. But you can see what advantages a world model could do. Now, the advantage of the world model we have in this paper is that it is going to enable this left hand process much faster because we don't even we don't need to interact with the world anymore to learn all of this stuff. We can simply do this in imagination while dreaming, so to say. That's why it's called dreamer and learn the stuff on the left. So it's not that the world model is used for explicit planning for explicit thinking ahead, it's just going to rapidly speed up this process on the left. It's technically model free reinforcement learning in a learned model, which is, I guess why it's called model based. OK, so how do we learn the world model? This is quite a complex thing. So the backbone, as you can see, is this H chain right here. So the H chain, that is your classic keep where the model keeps track of a latent state. So you everything that's kind of going on in the game right now, you want to save into the latent state. So the model is going to learn a latent state transition. And this specifically is using a GRU recurrent neural network with a gated recurrent unit. So it's not an LSTM, but it's kind of the little brother of the LSTM that is sometimes a bit easier to train. Sorry, Jurgen. But this is the backbone. So from step to step, we somehow we get an observation and we somehow want to incorporate that information and keep track of it. Now, how how we do it? So you basically, this is it, right? Usually you just feed this into an encoder, which in this case is going to be a convolutional neural network. And then you combine that, you put that as an input into your recurrent cell. Let's disregard everything else for a moment. How do you actually train the thing? So in model three reinforcement learning, you would simply predict the reward or the action that maximizes the reward like you would predict the best action to do in actor critic. Or you can actually predict the Q value in Q learning, not in model based. We're trying to learn a model. So what we're going to do is we're going to try to predict here. We're going to try to predict the image. Now, this can be, in fact, the next image or it can be the same image. And I don't even remember which one it is. OK. It predicts. I don't know. So it can I'm going to guess it. I'm going to guess it reconstructs the same image. OK. So here you can see the image predictor. Oh, yeah. So XT is predicted from H T and ZT. So we want to reconstruct the same image first and foremost. So we input an image and we want to get out the same image. This is like an like an auto encoder. So the representation we're going to get in the middle here somehow needs to be able to represent the image very well. And we also want to predict the reward. Here, we're also going to get an action. It's you can see it here more. So we're going to get an action. Remember, we are learning from experience. We have done this here a bunch of times and we have a data set of experience. So we know what actions we took. We're going to learn a model that tells us given we're in this state and perform a certain action, what's going to happen. So we're going to learn the reward and the image. And it might not make too much sense with the same frame. But if you look at the next frame, it makes a bit more sense. So given image X1, we want to encode it somehow. Right. And then through the GRU over here, we are informed. Well, while after X1 happened, we did in this episode, we did a one. And then we got reward R2. And the resulting image was X2. Okay, so we're trying to predict given an observation and a latent state, this H1, we're trying to end an action. We're trying to predict what reward we got and what the game looked like after we performed the action. This is trained in back propagation through time. So not only do we predict one future image, but we actually predict a sequence of rewards and images. Okay, so that's how we're going to learn a world model. Input observations and actions and output rewards and observations. Okay. And that's exactly what you saw at the beginning in these videos. So the model was simply input a bunch of frames here and then rolled out for a number of steps. And we looked at the output of this. This is, by the way, this is a D convolutional neural network, a D convolutional, you know, like in a DC GAN type of type of network. Okay. Now, what are these special parts right here? These special parts are what makes this model work. So the hidden state, as you can see, the thing I circled in red in the middle is not just the recurrent neural network hidden state. It is actually a combination of two things. They call this a combination of a fixed state of a deterministic state and a stochastic state. So what you're going to have is you're going to have the state, which is a vector. This is the H. Let's call that H zero. Okay. Of the of the LSTM. Now you're going to get an action into this, as we saw before, the action is combined with this. And you ask yourself, given that action and the hidden state. And now we don't just want to know what's the next hidden state, like in a normal RNN. What we're going to predict is actually this Z variable right here. And this Z variable is a description of the current state, a stochastic description of the current state in a very specific form. So the H is simply a vector, right? You can store in it whatever you want. But the Z, which is going to be concatenated to the H, it's going to be both is going to be predicted from the H. And it is also going to be concatenated to the H for further processing. So you're going to predict this thing together with the image X down here. You're going to predict that Z thing. And you're also going to concatenate it to H for further processing. So the red circle is going to be the concatenation and not even that. OK, maybe I should explain what it is. So it is going to be of this form. It is going to be a collection of categorical variables, each having, you know, 32. So it's 32 categorical variables, each having 32 possible classes. And the model can decide absolutely by itself what the categorical variables are for and what each of the classes mean. So, for example, in the Space Invaders game, right, one categorical could be the location of the agent. Location, right. And the 32 different values it could take are maybe going to be, you know, if this value is if it's this value, then it means the agent is somewhere down here in this quadrant or in this tile. If it's this value right here, the agent is going to be in here and so on. So these are categorical values and they can, you know, take one of these 32 different values. They can only take one. So that's the difference between these and like a Gaussian latent variable, because these stochastic states used to be modeled in like, say, you know, we have 32 Gaussians, like in a VAE. We have 32 of these latent variables. Now we make them categorical. And that turns out to be pretty good for this Atari games. So the other could be the enemy. Does the enemy shoot? Is, you know, has the enemy fired a shot? Now, maybe we don't need 32 variables right here. Like this could simply mean this could simply mean yes, and this could simply mean no. But also, you know, we can make use. We can encode actually 16 different enemies. So we can encode has this enemy shot that we see here or has an enemy that is potentially here fired a shot or has an enemy that is potentially here fired a shot. Right. We can we can encode this in that. Now I can see that you can see the problem, right. Two enemies can shoot at the same time. And in a categorical variable, you can only have one value. However, it might still be enough to just encode, you know, whichever enemy has shot most recently or least recently into this variable. And you can still play the game with that information. Okay. So you can see here that so it's 32 variables. So 32, we can have 32 here and each can have 32 different values. And, you know, the state is going to be described by by having each of these 32 variables be, you know, in one position or another, as you can see right here. Hey, it's Janek from the future. I forgot the whole video to show you this. So I'm doing it now. They have a pretty good explanation of why categorical variables might be important for a thing like Atari. And that is because sometimes you have pretty big junctures in the world state. So maybe, you know, you do very similar actions or maybe slightly different actions from the same states. But, you know, the slightly different action results in different changes in the world. And that means your prediction sort of has to capture all of that. So when your predictions is just a Gaussian, a Gaussian can only sort of have a mean and a variance. It cannot predict multimodal distributions. However, a categorical distribution can like it can be spiky. It can be very concentrated on one particular thing, or it can actually be a superposition of many different states. And when you sample from that, you actually have your multimodality. So it's again something that is kind of very suited to certain environments, but not others. And, you know, when it fits, then it seems to work pretty well. But this is in the blog post. If you want to look at this graphic yourself. All right. Back to past Janek. Bye bye. You can see that the entire observation sequence, the observations, they never get into the system except through these z variables. So this is an extreme compression. Every observation that you get in is going to be described by this extremely compressed format. And they hypothesize that, you know, because it's so compressed, because it's so sparse, it might actually force the model to learn pretty good latent variables. And that's also why it's so fast, because you never touch the observations again. You only work in this latent space. So what actually happens is the CNN is going to predict a distribution. So for each of the 32 variables is going to predict a distribution of the 32 values that variable could take. And one here and one and so on. It's going to predict 32 distributions of that. And then there is a sampling step. So this is now sampled from this. This is the sign for sampling from. And that gives you not 32 distributions, but it actually gives you 32 just straight. OK, here, here, here. So this is why it's called the stochastic part. So and that I'll actually make that blue. So you realize that is going to be fed here. So this deterministic state H is going to be used to predict this distribution. The distribution is going to be sampled from. And then this sample is going to be concatenated together with H. And that will finally make our actual latent state. So the latent state here is this concatenation out of the deterministic and out of a sample of the stochastic. And that ensures that you sort of keep your your options because it's sampled about the world model. You always draw from this distribution, which you can entropy regularize. Right. But you also have the deterministic information that you pull through. OK, so that's how the hidden state comes to be. And there is one node we haven't left out right yet. OK, during learning, during actual reinforcement learning, what you want to do is the following. You simply want to start off with a single observation or actually a hidden state that you've seen during training of the world model. And from that point on, you don't want to have anything to do with observation. So you see right here, since we we learned a reward predictor, right, we can simply use that reward predictor instead of the real environment. So and we don't want observations anymore. So what you want to do is you simply want to use this backbone here to predict the these latent states. So you simply want to unroll these latent states. Now, usually in order to do that, you need the observation. You can see here clearly the next latent state is a result of the previous one and the action and the observation. Now, if you don't want to do this, it means you have to predict the observation, but you can't predict the observation because that will be slow. And we already know that doesn't really work. So you want to predict this Z variable. We've said that observation, the next observation is going to be fed into the algorithm through this by means of constructing such a Z variable. So if you could predict that variable without seeing the observation, you could you don't need the observation anymore. And that's exactly the last output right here. You can see each H state is not only used to construct that Z variable together with the observation. We also predict the same Z variable, but without looking at the observation. OK, of course, that's going to be not as good. Like the latent representation is going to be much better when you actually see what happens in the game. However, in order to do dream reinforcement learning, we need to be able to completely detach from the observations. And that's why we also predict at the same time. So we predict the same variable, but without seeing the observation. And then we're going to introduce a loss function that makes it such that these two are going to be very close together. So the agent now has to do a trade off. And the trade off is, do I want to get the best information out of my observation? Do I want to represent it as accurately as possible in order to reconstruct it really well? And in order to predict the reward really well? Or do I want to be able to predict this thing without seeing the observation, which means that, you know, I have to I have to not rely as much on the image. I have to rely more on learning the actual dynamics of the world and what happens when I perform actions in them. That's what exactly what this KL divergence here is going to do. So the model has to find a trade off between the two. And if you engineer that trade off correctly, you are able to use the just the predicted Z variables instead of the true ones, at least for a certain number of steps. I think they do 15 steps into the future during learning. And of course, the errors accumulate because you're never able to predict that Z exactly. However, it's enough to do good reinforcement learning. And this sparsity here, it helps very much. OK, I know this is a lot, but, you know, to shortly recap, learning world model means that you input observations and you learn to predict the future. So you learn to predict the future observations. You learn to predict the future rewards, given actions, given actions that you perform. You start off with a random agent or any agent you want. You simply want to learn what happens when I do something. Now, the way you predict that is going to be through a recurrent neural network, the latent state of which is going to be a combination of a classic latent state of an RNN and concatenated with a sample from a stochastic, very, very compressed state that you obtain from a CNN encoder combined with the last hidden state. So the combination of a sample from this and the deterministic state is going to be your compact world model state from which you predict the future. And in addition to that, you also try to predict this stochastic state just from the deterministic hidden state and the action without knowing what the actual next observation is or the current observation, I guess. And that means you can then use those prediction values at reinforcement learning time in order to be completely decoupled from the observations. And now, yeah, we we we sort of have it. So what if you learn a world model like this, what you can do now is you don't need the observations anymore. You maybe need one start observation and you simply unroll into the future and you do reinforcement learning in this completely imaginary like this is a dream. Now, this is a dream. This is just dream, a dream. Now, it's it's also completely not cheated. Yeah. So the reinforcement learning they do right here is going to be something like, you know, a to see or a three, see, it's going to be an actor critic method and advantage actor critic method. That's a pretty basic but very strong reinforcement learning algorithm where you learn sort of two models. You learn the critic that accumulates that tries to predict the future reward. So they try to predict these values right here. And you learn an actor that is trying to make the critic really, really happy. Now, you swap this once you have a good agent, you go back, you collect more data because your world model is never going to be accurate. It's never going to replace actually playing the environment. Your world model only has data from where the agent goes. Right. That's where it learns from. So it's crucial that once you have a better agent, you update your world model because now the agent does different things and it goes places that the world model has never seen. Right. If you know, if you have this, if you have like a maze game. Okay. And the mazes. I don't know. I'm not good at mazes, but you know, you're here. And once you crash into a wall, you're done. So the agent, it will just be random at the beginning. So like crash a lot into these walls and so on. You just do random actions. So the world model, if it just learns from that experience, it is going to learn maybe that there's a wall right here. But this thing we don't know. Right. Now, if you get a little bit of reward, maybe there's a coin right here. Okay. And every now and then this stupid random agent actually finds the coin. Right. It walks over here and finds the coin and gets a reward. The reinforcement learning means that it's going to do that more often. So now the agent is going to walk over here more and more often. But you only do that in the world model. The world model only knows up until here because that's where the agent has gone the farthest. Now that the agent goes further, right, you actually need to go back to the environment and let the agent run in the true environment. Because now that agent's going here, you know, it's going to explore a bit more. Because, you know, it learned it learned only seeing this. And now it learns a bit more. You record, you build out your world model. It's like, ah, there's the wall goes until here, but then there's a free space and then maybe something comes here and so on. So working with world model is not is not super easy. And it only is going to this is very specific. And this is going to be my my criticism right here in that all of this seems quite specific to Atari. Reinforcement learning is such a big field and such a general algorithm that you're going to build in some kind of prior knowledge about the world. But it seems like the some reinforcement learning papers, I never know how much is this all applicable to other oral environments. It seems like this is specifically for Atari. And learning these world models in this fashion is only going to work if, you know, every now and then you find a reward, you still have the explore exploit dilemma. If your world model isn't accurate, then, you know, you're not going to do accurate RL and so on. And maybe the density of rewards isn't going to be enough for you to actively push yourself up in these cycles. And, you know, there's another problem with these latent variables, they're categorical, which I think, you know, is super cool because it gives you a sparse representation. But you only learn it from the images. In fact, they say they can even leave away the reward predictor for the world model. So you learn to reconstruct the images. However, if two images are very close to each other, but they mean different things in the game. So, you know, two images can be super duper close, like an enemy can be here or slightly off, right? But if it's slightly off, it doesn't hit you. And therefore, you know, you're all good. Now, these two states are still pretty close because if you move a bit, you're likely to get hit. But sometimes a little bit of a change in image can mean actually a big change in game state and vice versa, which is actually even worse. A big change in image can mean like it doesn't matter. Like if everything in the image rotates around, but your agent still has nothing and is at the same place, it means nothing to you as a human. Yet an algorithm like this that whose goal it is to predict the future as accurately as possible, it will devote a lot of attention to accurately predict the future or predict variances in the future. Even though they might not be relevant. So in this in this task of or in this bottleneck of encoding everything into a very compact state, you might actually lose important information. And that means all of all of the like two states that are very, very far like need to be differentiated are going to be just the same in this representation. And that means your agent will never really learn because one is bad and one is good. So the mean reward is zero. And it says, well, when I get to that state, my mean reward is kind of zero and it's just kind of a big variance. And then the world model will never learn the difference because it has bigger things to worry about. So this is it's all very specific. And you'll see this in the in the loss term right here. So this is the loss function for learning the world model. And you can see they have an image reconstruction loss right here. This is a this is a cross entropy loss. So it's this is your approximation distribution. This is what really happened. Yeah, it's a it's kind of a probabilistic way of writing things. So these are cross entropy losses when you see log P of the expectation of under Q. They have a loss predicting the reward. They have a loss predicting the discount, which is mainly made for predicting when an episode ends in the in the imagined trajectory. And then they have this transition loss coupled with the entropy regularizer. So the transition loss is going to be for predicting these Z states and the entropy regularizer is for keeping the distribution in the Z states not peaked. So you want to kind of retain that stochasticity and this together you might recognize as the KL divergence between the P and Q. And that's this connection right here. So I'm going to minimize the KL, which is the same as saying I want this thing to be as accurate. I want I want I want these things to be as close as possible to each other, but the entropy should should still be given. And yeah, as you can see here, you can you can you can decompose that. So this is going to be this is going to be the KL divergence between the two distributions. I don't have a better way of explaining that without writing it down. You can already see they have a massive amount of hyperparameters, right? Like here's one, here's one, here's one, here's one, here's one. OK, so even within the KL divergence, they have actually two one hyperparameter for the KL divergence and one to trade off the entropy with the actual cross with the transition log loss with the cross entropy there. And they do ablations and see that that is really important that you have that trade off that you're able to make that trade off. And it's the same as the beta variational autoencoder, by the way. It's an entire paper about why you need an additional hyperparameter here. Like that's the entire paper of beta VAs, which I found funny. But, you know, it seems to be important. So you can see right here, this is KL balancing. So you have one, you have one term for making the prior close to the posterior, the prior being the one where you just see H and the posterior being the one where you see H and X. And you have another term for making the posterior close to the prior and you trade them off with these variables right here. Then the reinforcement learning itself, again, has a bunch of hyperparameters. So it is doing TD lambda learning. And you can look that up. TD lambda learning basically means you are here in your state and you're going to predict the value, sorry, the reward. Going to the next state and you're going to predict the value at that state. And then you're also going to predict from the same state the reward two steps forward and the value at that state. And you're also going to predict the reward three steps forward and the value at that state. And at the end, you're going to sum all of that up into one number that is kind of an aggregate of all of this. And that's going to be your prediction. That's what you regress on in your value predictor. And the actor tries to maximize that. So there's another parameter lambda that tells you how you aggregate these things. Right. And also H for how many steps you do that. There's also going to be in the actor loss function. They decided not only do they want the classic reinforce loss as you have, you actually want the straight through estimator of the distribution. And so a straight through estimator is when you want to backprop through sampled things. Normally, the reinforced gradients, what they do is if your actor outputs a distribution, let's say over three actions. Right. You don't all you can say is that I did action to here and it gave me seven reward. Right. So you want to make that more likely because seven is pretty good. Actually, you subtract the baseline. But, you know, let's say after the baseline, it's seven. So you simply act like you have a target distribution of this and scale it by seven. That's reinforced gradients. What you could also do is you could actually regress on directly through the softmax operation right here. Because this here is a sampling step. You cannot backprop through sampling steps. The way you can do it is that you you take the signal, the loss signal here, but you act as if this was your output and not this. OK, so you act as if you had made actions in proportion to their distribution and not actually sampled one particular action. This is going to give you a biased signal, but it has much lower variance. Whereas if you sample and then scale, it's going to be unbiased, but much higher variance. So they do these straight through estimators not only here, but actually also in this step up here. And you can see how that works in modern deep learning frameworks. So you have your distribution in terms of your logits. So what you can do is you sample from them and forward propagate should be the sample. Right. So the trick is to do plus and minus the same thing. So the forward propagation signal is simply your sample, as you can see right here. Now, the sample, this operation, it has no gradient. Oh, you can't see that it has no gradient. So the deep learning framework will simply not backprop through it. So if you were to just use the sample in your graph, you won't get a gradient. But what you can do is you can actually calculate the probabilities here, like the thing you want to back propagate, and then do plus that and minus stop gradient of that. You can see right here, this has no gradient. This has no gradient. So the gradient is going to be as if you had forward propagated this probes variable. But on the forward pass, the probes variable exactly cancels out with itself. And the sample is forward propagated. This is called a straight through estimator. It gives you a biased gradient, but much less variance than if you had to, you know, if you scale the sample like the reinforced gradients. So they use this in the world model. And they use this actually in the actor loss right here. And you can see there is another hyperparameter. Here is another hyperparameter. And then they have an entropy regularizer to facilitate exploration, which is normal, but gives you another regularizer. And not only do they have, sorry, hyperparameter, not only do they have these three additional hyperparameters, they scale two of them during training. So they now have a schedule to scale them. So this straight through estimator, they actually scale it to zero over the course of training. But yet two more hyperparameters, namely how fast you want to decay those things. So this whole thing is a giant bucket of hyperparameters. And so they say, while the unbiased reinforced gradients can help a better final solution. However, we find that using only reinforced gradients for optimizing the policy also works well. It might just not work as fast or as well, but it also works well. You know that in general, this is reinforcement learning, but this is a bit, you know, the amount of hyperparameters here is quite staggering. And I'm going to guess that this took a lot of work to even get off the ground. Right. So here you can see how this compares to other algorithms. Specifically blue here is Dreamer V2. And they do suggest a bunch of different things. So they have task median gamer normalized. So gamer is a professional human level gamer. And gamer normalized means you simply divide by what that professional gamer can do. So you can see that it can even exceed, you know, this gamer. So here is over 1.5 times over 55 different Atari games. Very good. However, these Atari games, some of them are actually unbounded. And in some of them, a machine can just be so much better than a human that usually these scores are dominated by very, very few games where the machine just excels, you know, hugely. And other games are like zero and both the median score and the mean score. They are not really meaningful. At least that's what this paper here argues. So they propose two modifications. So the first modification, actually, this is from a different paper as well, says you shouldn't normalize by, you know, kind of a professional gamer. You should actually normalize by the human world record. So this is record normalized. You can see it gives a cleaner score. And then they say, well, given that a few games still the the machine can just outperform humans so much. What you should do is actually you should never allow. So you just you should just clip the machine score at where the human world record is. So the reasoning behind this, I can imagine, is something like what's the difference between the human world record and the professional gamer world record? Well, the human world record, the professional gamer is already pretty good at gaming in general, let's say. But the human world record holder has probably figured out every single detail of that particular game and is pushing it with like exploits and whatnot. I don't know if you've seen legend like Ocarina of Time speed runs lately, but they're crazy. So that is going to be human world record. And it's probably going to be better to normalize by this because, you know, the machine will necessarily find these kind of exploits. They will it will probably find them as well. However, there are some things that where the machine you have to be where you have to be like pixel and microsecond accurate where the machine can do it and the human can't. So clipping it might make sense. I'm not really sure about this, like there's arguments to be made that you maybe shouldn't normalize by the human world record because, you know, you don't want to give credence to like exploits. But the gamer kind of represents more how the game is intended to be played. I don't know. They just suggest this new score just so happens to be that in this new score, they are, you know, other than here, they are just dominating at all time points. Yeah, let's let's leave them that they do a quite a number of ablations, especially they find out that, for example, if they do latent variables as categorical that outperforms Gaussian latent variables by a lot. So and that's, you know, that's kind of a reasoning why they use the categorical variables. The KL balancing simply means that additional parameter in the KL term, if they enable it, you can see it helps a lot. Image gradients. So when they they wonder, can we learn the world models from predicting images or from predicting rewards or from both? So they do both as a default. But if they leave away the image gradients, it doesn't work anymore. However, if they leave away the reward gradients, you can see it still works pretty well. Again, this is all quite Atari specific. And it also means that you can see right here, right? The Atari game lends itself to this kind of to exactly this kind of model. So how much this is a success for general reinforcement learning is questionable. However, what you can say is that if an environment lends itself to be world model learned by this kind of latent categorical variables, like so if the image state is going to be if changes in the image are going to be a good indicator of actual changes in relevant world variables, then you know, you might you might be very suited with a model like this. And so they compare this to other algorithms, for example, to use zero, which doesn't run on a single GPU. I think it is better, but it doesn't run on a single GPU. And it uses kind of a lot more Atari frames than the the dreamer algorithm. So you see again that you just need to find the correct category and you can be state of the art. So if this is like single GPU, Atari, no, I don't want to I don't want to dunk on this. This is pretty cool work. And if you look at the code, it took a lot of effort. Like you can see that from the code. OK, the last thing I want to look at is where does it succeed and where does it fail? So you can see a comparison, for example, dreamer V2 versus IQN or dreamer V2 versus Rainbow. And you can see and particularly interesting is where does it fail? And it fails in video pinball. And actually, I don't have it pulled up right here. But if you look it up, so if you look it up, you can probably see why. Because this video pinball thing. Thanks. Thanks, YouTube. This video pinball thing, it has a lot of changes in image without really doing much changes in the world state. So what actually matters is like this little tiny ball, this little tiny, you know, it's kind of a bunch of pixels. And the rest, you know, kind of moves around. And OK, maybe it doesn't move too much right here. But still, you know, there's this new cross that appears and so on. So a world model that learns to, you know, there's kind of flashes over the whole image, a world model that learns to accurately predict the world. Maybe is going to not focus so much on that little ball, but maybe is going to focus more on the rest of the image if that changes well. And also, you can see maybe the reward. Now, again, a flash, the reward doesn't change all too much. Yeah, it does, maybe. But, you know, any any time it bumps somewhere. So my hypothesis is going to be that in games where what actually matters consists of very few changes in the actual image. And there are lots of other big image changes that don't really matter so much for the immediate reward, maybe for the future, but not for the immediate. This algorithm is going to not be as good. And that is one example is this video pinball. And I might be wrong on this, but it's kind of a hypothesis. So the code for this is going to is available right here. Check it out as well as you should check out the blog post. They have a lot of ablations right here, as you can see, and graphs for the individual games turning off and on different variables. And you might as well give it a try if you have a reinforcement learning problem that has an environment similar to Atari. All right. That was everything I had to say for this pretty cool paper. Check it out. Bye bye.
 Hi there. What you're seeing here are predictions by a world model learned for Atari reinforcement learning.
 On the top you see what really happened during an episode of play.
 And on the bottom you see the predictions of this world model.
 The world model just gets five frames at the beginning, which you don't even see here as a conditioning.
 And then it predicts 45 frames of gameplay.
 It's astounding how accurate it is, not only in terms of how the game evolves,
 but also in terms of what the agent will actually do.
 So the world model, the specific world model you see here is part of the Dreamer V2 algorithm
 from the paper Mastering Atari with Discrete World Models by Danijar Hafner, Timothy Lilikrub,
 Mohamed Nourouzi and Jimmy Ba of Google Brain, DeepMind and the University of Toronto.
 So these kind of world models, they enable you to do very quick reinforcement learning
 once you have the model, you can use it to imagine yourself playing the game instead of actually playing the game.
 And therefore you can do much more efficient reinforcement learning.
 And this paper details how to get an accurate world model for Atari, which was sort of out of reach until now,
 especially considering that they only do single GPU reinforcement learning.
 So the result, as you can see here, is going to be an algorithm that is the top single GPU agent right now,
 competing, outperforming other, so here is Dreamer V2 outperforming other algorithms such as Rainbow, IQN, DQN.
 And the special thing here is that Dreamer V2 is a model based algorithm,
 whereas the current or the previous best ones, especially single GPU best ones, were model free algorithms.
 And you can see the next best model based algorithms were, are not really competitive in Atari, right?
 This is specifically Atari. So Dreamer V2 is an evolution of Dreamer V1,
 which worked well for things like continuous control, but Atari still seemed a bit out of reach.
 So the difference between model based reinforcement learning and model free reinforcement learning is that
 model based reinforcement learning first learns a model of the world, it learns how the world acts,
 and then it uses that model to learn what actions to perform,
 whereas model free algorithms, they simply act in the world and they learn to predict the best actions as they act in the world.
 So there's your difference. And how does Dreamer V2 do that? On the high level, it has two stages.
 Stage one is learn a world model from past experience.
 And then stage two is use that world model, as we said, for reinforcement learning.
 And the reinforcement learning here is going to be just actor critic learning. Very straightforward.
 There's a little modification with a pass through estimator.
 But the real difference is going to be in how the world model is learned.
 And the novel contribution or the main contribution here is this latent state,
 which consists of a this stochastic latent state, which other than other world models,
 which model the latent states as something like Gaussian random variables.
 This paper models the latent state as categorical random variables.
 And that turns out to work pretty well for Atari.
 So that's step one. Learn world model. Step two, do reinforcement learning in the model.
 So not using any data anymore. And you can repeat those two steps as many times as you want.
 So you start out with a set of data, then you learn an actor,
 and then you use that actor to collect more data and so on until you have a really good actor.
 And the world model is really accurate for that actor.
 So that's the overview. And it's going to turn out, as we already saw,
 to beat other, at least single GPU models by quite a bit.
 So we'll go through the paper through the individual steps and discuss what's new and how it all works.
 The code is also available. I'll link to it.
 And the blog post I've shown you here has some more explanatory graphics.
 If you like content like this, as always, don't hesitate to click like and share it with all your friends,
 especially the Atari gamers, because they are outperformed, as you can see here.
 All right. So world models.
 Pretty quickly in reinforcement learning, as you all hopefully or maybe know,
 you have an agent that is interacting with an environment.
 And the agent can... So the environment always provides the agent with an observation,
 which would be an image in an Atari game.
 And the agent decides to do one of many available actions in response to receiving the observation.
 The environment then responds with a reward for that action.
 So either you die, which is like negative reward, or you collect a coin, which is positive reward,
 or you win the game, which is like a thousand reward.
 And it also gives the agent a new observation, the next observation.
 And the agent, again, responds by performing another action and so on.
 So you have this cycle. And the goal of reinforcement learning agent is usually to maximize all the rewards
 that it collects during playing with the environment.
 And you want to repeat that many times for many episodes to have the agent learn to get as to do the actions
 that are as good as possible in terms of reward.
 All right. Now, in classic, let's say classic, in model-free reinforcement learning,
 one way to do this is to take this right here as you play the game.
 As you play the game, you collect data, right?
 So let's assume we collect data as we act in the world.
 And from this data, we can learn something.
 So model-free learns from the raw experience.
 So an episode will always be a series of images, right?
 And actions you have performed.
 So here is an image and I have performed action one and then came a next image and I've performed action two.
 So what classic reinforcement learning would do is it would say,
 okay, from this transition doing this action, I have gotten five reward.
 And from this transition in this action, I've gotten negative three reward.
 So I'm going to have to do this action one more often because it gave me a lot of reward after I observe this thing here, right?
 The combination of this thing, I need to do action one more.
 And when I'm in this situation, I need to do action two less and so on.
 Okay, so you're simply trying to put this image that you get into a neural network that tries to predict action one as often as possible.
 And you want the same network when you input this next image to not predict action two.
 So like anything but action two.
 So that's going to be that's kind of the logic between of the classic model-free reinforcement learning.
 Usually this is implemented in a sort of an LSTM fashion or it's one way of doing it.
 So you have an LSTM that tracks a hidden state.
 Why do you need a hidden state?
 Because you might not see everything in the image there is, right?
 This is not necessarily Markovian.
 So there might be information that you need to remember for a long time, like when an enemy leaves the screen and then comes back, you want to track it.
 Do you have an LSTM or some kind of RNN and then you want to feed the images into that one by one.
 And then you simply so with an encoder, which is usually kind of a convolutional neural network, I want to draw it like this.
 And then you try to predict the here the good actions and here you try to not predict the bad action and so on.
 So this is a simple classifier.
 Ultimately, it's an LSTM with a classifier on top.
 And the classifier simply tries to either predict a class of action one or not or predict anything else.
 So and you train it via back propagation through time.
 And that's it.
 Now, here is a little bit different.
 So why?
 Why is this maybe not a good idea?
 Well, all you have is the signal of the reward for given actions.
 And that means it is it is fairly hard to generalize in these kinds of things.
 So when you imagine you have your screen right here and there's an opponent kind of here, there's an opponent here and you are down here and the opponent shoots.
 Right. You have to move out of the way.
 You have to move over here.
 Now, RL is completely capable of learning that.
 However, take the next situation over here.
 Now, the opponent is here, shoots and you are down here.
 You have to, again, learn to move out of the way for a classic RL algorithm.
 These two things are identity are completely different states.
 Like this is there's nothing equal about the two.
 Like this is a completely different thing.
 And it has to sort of learn by force.
 Look, in this situation, there, you know, you need to move.
 And in this situation, you also need to move.
 Now, given that that is a convolutional neural network, it might after a while learn the fact that it, you know, these two situations have something in common.
 But in essence, these are two different things.
 And you have to learn purely from the reward, purely from the fact that you're going to die if you don't move to get out of the way in two situations.
 And of course, this situation can be replicated all over.
 However, if you have a world model, right, imagine now we have a world model over here and the world model accurately learns to predict the future.
 Now we know that, you know, we are here.
 This is here.
 Now we can imagine ourselves forward and we're going to see we're going to get hit.
 And that means we need to go out of the way.
 So doing this explicitly would be called planning.
 We are not going to do planning in this paper.
 OK, we are still going to do the classic RL.
 But you can see what advantages a world model could do.
 Now, the advantage of the world model we have in this paper is that it is going to enable this left hand process much faster because we don't even we don't need to interact with the world anymore to learn all of this stuff.
 We can simply do this in imagination while dreaming, so to say.
 That's why it's called dreamer and learn the stuff on the left.
 So it's not that the world model is used for explicit planning for explicit thinking ahead, it's just going to rapidly speed up this process on the left.
 It's technically model free reinforcement learning in a learned model, which is, I guess why it's called model based.
 OK, so how do we learn the world model?
 This is quite a complex thing.
 So the backbone, as you can see, is this H chain right here.
 So the H chain, that is your classic keep where the model keeps track of a latent state.
 So you everything that's kind of going on in the game right now, you want to save into the latent state.
 So the model is going to learn a latent state transition.
 And this specifically is using a GRU recurrent neural network with a gated recurrent unit.
 So it's not an LSTM, but it's kind of the little brother of the LSTM that is sometimes a bit easier to train.
 Sorry, Jurgen. But this is the backbone.
 So from step to step, we somehow we get an observation and we somehow want to incorporate that information and keep track of it.
 Now, how how we do it?
 So you basically, this is it, right? Usually you just feed this into an encoder, which in this case is going to be a convolutional neural network.
 And then you combine that, you put that as an input into your recurrent cell.
 Let's disregard everything else for a moment.
 How do you actually train the thing?
 So in model three reinforcement learning, you would simply predict the reward or the action that maximizes the reward like you would predict the best action to do in actor critic.
 Or you can actually predict the Q value in Q learning, not in model based.
 We're trying to learn a model.
 So what we're going to do is we're going to try to predict here.
 We're going to try to predict the image.
 Now, this can be, in fact, the next image or it can be the same image.
 And I don't even remember which one it is.
 OK.
 It predicts.
 I don't know. So it can I'm going to guess it.
 I'm going to guess it reconstructs the same image.
 OK. So here you can see the image predictor.
 Oh, yeah. So XT is predicted from H T and ZT.
 So we want to reconstruct the same image first and foremost.
 So we input an image and we want to get out the same image.
 This is like an like an auto encoder.
 So the representation we're going to get in the middle here somehow needs to be able to represent the image very well.
 And we also want to predict the reward.
 Here, we're also going to get an action.
 It's you can see it here more.
 So we're going to get an action.
 Remember, we are learning from experience.
 We have done this here a bunch of times and we have a data set of experience.
 So we know what actions we took.
 We're going to learn a model that tells us given we're in this state and perform a certain action, what's going to happen.
 So we're going to learn the reward and the image.
 And it might not make too much sense with the same frame.
 But if you look at the next frame, it makes a bit more sense.
 So given image X1, we want to encode it somehow.
 Right. And then through the GRU over here, we are informed.
 Well, while after X1 happened, we did in this episode, we did a one.
 And then we got reward R2.
 And the resulting image was X2.
 Okay, so we're trying to predict given an observation and a latent state, this H1, we're trying to end an action.
 We're trying to predict what reward we got and what the game looked like after we performed the action.
 This is trained in back propagation through time.
 So not only do we predict one future image, but we actually predict a sequence of rewards and images.
 Okay, so that's how we're going to learn a world model.
 Input observations and actions and output rewards and observations.
 Okay. And that's exactly what you saw at the beginning in these videos.
 So the model was simply input a bunch of frames here and then rolled out for a number of steps.
 And we looked at the output of this.
 This is, by the way, this is a D convolutional neural network, a D convolutional, you know, like in a DC GAN type of type of network.
 Okay. Now, what are these special parts right here?
 These special parts are what makes this model work.
 So the hidden state, as you can see, the thing I circled in red in the middle is not just the recurrent neural network hidden state.
 It is actually a combination of two things.
 They call this a combination of a fixed state of a deterministic state and a stochastic state.
 So what you're going to have is you're going to have the state, which is a vector.
 This is the H. Let's call that H zero.
 Okay. Of the of the LSTM.
 Now you're going to get an action into this, as we saw before, the action is combined with this.
 And you ask yourself, given that action and the hidden state.
 And now we don't just want to know what's the next hidden state, like in a normal RNN.
 What we're going to predict is actually this Z variable right here.
 And this Z variable is a description of the current state, a stochastic description of the current state in a very specific form.
 So the H is simply a vector, right? You can store in it whatever you want.
 But the Z, which is going to be concatenated to the H, it's going to be both is going to be predicted from the H.
 And it is also going to be concatenated to the H for further processing.
 So you're going to predict this thing together with the image X down here.
 You're going to predict that Z thing.
 And you're also going to concatenate it to H for further processing.
 So the red circle is going to be the concatenation and not even that.
 OK, maybe I should explain what it is. So it is going to be of this form.
 It is going to be a collection of categorical variables, each having, you know, 32.
 So it's 32 categorical variables, each having 32 possible classes.
 And the model can decide absolutely by itself what the categorical variables are for and what each of the classes mean.
 So, for example, in the Space Invaders game, right, one categorical could be the location of the agent.
 Location, right. And the 32 different values it could take are maybe going to be, you know, if this value is if it's this value,
 then it means the agent is somewhere down here in this quadrant or in this tile.
 If it's this value right here, the agent is going to be in here and so on.
 So these are categorical values and they can, you know, take one of these 32 different values.
 They can only take one. So that's the difference between these and like a Gaussian latent variable,
 because these stochastic states used to be modeled in like, say, you know, we have 32 Gaussians, like in a VAE.
 We have 32 of these latent variables. Now we make them categorical.
 And that turns out to be pretty good for this Atari games.
 So the other could be the enemy. Does the enemy shoot?
 Is, you know, has the enemy fired a shot? Now, maybe we don't need 32 variables right here.
 Like this could simply mean this could simply mean yes, and this could simply mean no.
 But also, you know, we can make use. We can encode actually 16 different enemies.
 So we can encode has this enemy shot that we see here or has an enemy that is potentially here fired a shot or has an enemy that is potentially here fired a shot.
 Right. We can we can encode this in that.
 Now I can see that you can see the problem, right.
 Two enemies can shoot at the same time. And in a categorical variable, you can only have one value.
 However, it might still be enough to just encode, you know, whichever enemy has shot most recently or least recently into this variable.
 And you can still play the game with that information.
 Okay. So you can see here that so it's 32 variables.
 So 32, we can have 32 here and each can have 32 different values.
 And, you know, the state is going to be described by by having each of these 32 variables be, you know, in one position or another, as you can see right here.
 Hey, it's Janek from the future.
 I forgot the whole video to show you this.
 So I'm doing it now. They have a pretty good explanation of why categorical variables might be important for a thing like Atari.
 And that is because sometimes you have pretty big junctures in the world state.
 So maybe, you know, you do very similar actions or maybe slightly different actions from the same states.
 But, you know, the slightly different action results in different changes in the world.
 And that means your prediction sort of has to capture all of that.
 So when your predictions is just a Gaussian, a Gaussian can only sort of have a mean and a variance.
 It cannot predict multimodal distributions.
 However, a categorical distribution can like it can be spiky.
 It can be very concentrated on one particular thing, or it can actually be a superposition of many different states.
 And when you sample from that, you actually have your multimodality.
 So it's again something that is kind of very suited to certain environments, but not others.
 And, you know, when it fits, then it seems to work pretty well.
 But this is in the blog post. If you want to look at this graphic yourself.
 All right. Back to past Janek. Bye bye.
 You can see that the entire observation sequence, the observations,
 they never get into the system except through these z variables.
 So this is an extreme compression.
 Every observation that you get in is going to be described by this extremely compressed format.
 And they hypothesize that, you know, because it's so compressed, because it's so sparse,
 it might actually force the model to learn pretty good latent variables.
 And that's also why it's so fast, because you never touch the observations again.
 You only work in this latent space.
 So what actually happens is the CNN is going to predict a distribution.
 So for each of the 32 variables is going to predict a distribution of the 32 values that variable could take.
 And one here and one and so on.
 It's going to predict 32 distributions of that.
 And then there is a sampling step.
 So this is now sampled from this.
 This is the sign for sampling from.
 And that gives you not 32 distributions, but it actually gives you 32 just straight.
 OK, here, here, here.
 So this is why it's called the stochastic part.
 So and that I'll actually make that blue.
 So you realize that is going to be fed here.
 So this deterministic state H is going to be used to predict this distribution.
 The distribution is going to be sampled from.
 And then this sample is going to be concatenated together with H.
 And that will finally make our actual latent state.
 So the latent state here is this concatenation out of the deterministic and out of a sample of the stochastic.
 And that ensures that you sort of keep your your options because it's sampled about the world model.
 You always draw from this distribution, which you can entropy regularize.
 Right. But you also have the deterministic information that you pull through.
 OK, so that's how the hidden state comes to be.
 And there is one node we haven't left out right yet.
 OK, during learning, during actual reinforcement learning, what you want to do is the following.
 You simply want to start off with a single observation or actually a hidden state that you've seen during training of the world model.
 And from that point on, you don't want to have anything to do with observation.
 So you see right here, since we we learned a reward predictor, right, we can simply use that reward predictor instead of the real environment.
 So and we don't want observations anymore.
 So what you want to do is you simply want to use this backbone here to predict the these latent states.
 So you simply want to unroll these latent states.
 Now, usually in order to do that, you need the observation.
 You can see here clearly the next latent state is a result of the previous one and the action and the observation.
 Now, if you don't want to do this, it means you have to predict the observation, but you can't predict the observation because that will be slow.
 And we already know that doesn't really work.
 So you want to predict this Z variable.
 We've said that observation, the next observation is going to be fed into the algorithm through this by means of constructing such a Z variable.
 So if you could predict that variable without seeing the observation, you could you don't need the observation anymore.
 And that's exactly the last output right here.
 You can see each H state is not only used to construct that Z variable together with the observation.
 We also predict the same Z variable, but without looking at the observation.
 OK, of course, that's going to be not as good.
 Like the latent representation is going to be much better when you actually see what happens in the game.
 However, in order to do dream reinforcement learning, we need to be able to completely detach from the observations.
 And that's why we also predict at the same time.
 So we predict the same variable, but without seeing the observation.
 And then we're going to introduce a loss function that makes it such that these two are going to be very close together.
 So the agent now has to do a trade off.
 And the trade off is, do I want to get the best information out of my observation?
 Do I want to represent it as accurately as possible in order to reconstruct it really well?
 And in order to predict the reward really well?
 Or do I want to be able to predict this thing without seeing the observation, which means that, you know, I have to I have to not rely as much on the image.
 I have to rely more on learning the actual dynamics of the world and what happens when I perform actions in them.
 That's what exactly what this KL divergence here is going to do.
 So the model has to find a trade off between the two.
 And if you engineer that trade off correctly, you are able to use the just the predicted Z variables instead of the true ones, at least for a certain number of steps.
 I think they do 15 steps into the future during learning.
 And of course, the errors accumulate because you're never able to predict that Z exactly.
 However, it's enough to do good reinforcement learning.
 And this sparsity here, it helps very much.
 OK, I know this is a lot, but, you know, to shortly recap, learning world model means that you input observations and you learn to predict the future.
 So you learn to predict the future observations.
 You learn to predict the future rewards, given actions, given actions that you perform.
 You start off with a random agent or any agent you want.
 You simply want to learn what happens when I do something.
 Now, the way you predict that is going to be through a recurrent neural network, the latent state of which is going to be a combination of a classic latent state of an RNN and concatenated with a sample from a stochastic,
 very, very compressed state that you obtain from a CNN encoder combined with the last hidden state.
 So the combination of a sample from this and the deterministic state is going to be your compact world model state from which you predict the future.
 And in addition to that, you also try to predict this stochastic state just from the deterministic hidden state and the action without knowing what the actual next observation is or the current observation, I guess.
 And that means you can then use those prediction values at reinforcement learning time in order to be completely decoupled from the observations.
 And now, yeah, we we we sort of have it.
 So what if you learn a world model like this, what you can do now is you don't need the observations anymore.
 You maybe need one start observation and you simply unroll into the future and you do reinforcement learning in this completely imaginary like this is a dream.
 Now, this is a dream.
 This is just dream, a dream. Now, it's it's also completely not cheated.
 Yeah. So the reinforcement learning they do right here is going to be something like, you know, a to see or a three, see, it's going to be an actor critic method and advantage actor critic method.
 That's a pretty basic but very strong reinforcement learning algorithm where you learn sort of two models.
 You learn the critic that accumulates that tries to predict the future reward.
 So they try to predict these values right here.
 And you learn an actor that is trying to make the critic really, really happy.
 Now, you swap this once you have a good agent, you go back, you collect more data because your world model is never going to be accurate.
 It's never going to replace actually playing the environment.
 Your world model only has data from where the agent goes.
 Right. That's where it learns from.
 So it's crucial that once you have a better agent, you update your world model because now the agent does different things and it goes places that the world model has never seen.
 Right. If you know, if you have this, if you have like a maze game.
 Okay. And the mazes. I don't know. I'm not good at mazes, but you know, you're here.
 And once you crash into a wall, you're done. So the agent, it will just be random at the beginning. So like crash a lot into these walls and so on.
 You just do random actions. So the world model, if it just learns from that experience, it is going to learn maybe that there's a wall right here.
 But this thing we don't know. Right.
 Now, if you get a little bit of reward, maybe there's a coin right here. Okay.
 And every now and then this stupid random agent actually finds the coin. Right. It walks over here and finds the coin and gets a reward.
 The reinforcement learning means that it's going to do that more often.
 So now the agent is going to walk over here more and more often.
 But you only do that in the world model. The world model only knows up until here because that's where the agent has gone the farthest.
 Now that the agent goes further, right, you actually need to go back to the environment and let the agent run in the true environment.
 Because now that agent's going here, you know, it's going to explore a bit more.
 Because, you know, it learned it learned only seeing this.
 And now it learns a bit more. You record, you build out your world model.
 It's like, ah, there's the wall goes until here, but then there's a free space and then maybe something comes here and so on.
 So working with world model is not is not super easy.
 And it only is going to this is very specific.
 And this is going to be my my criticism right here in that all of this seems quite specific to Atari.
 Reinforcement learning is such a big field and such a general algorithm that you're going to build in some kind of prior knowledge about the world.
 But it seems like the some reinforcement learning papers, I never know how much is this all applicable to other oral environments.
 It seems like this is specifically for Atari.
 And learning these world models in this fashion is only going to work if, you know, every now and then you find a reward, you still have the explore exploit dilemma.
 If your world model isn't accurate, then, you know, you're not going to do accurate RL and so on.
 And maybe the density of rewards isn't going to be enough for you to actively push yourself up in these cycles.
 And, you know, there's another problem with these latent variables, they're categorical, which I think, you know, is super cool because it gives you a sparse representation.
 But you only learn it from the images.
 In fact, they say they can even leave away the reward predictor for the world model.
 So you learn to reconstruct the images.
 However, if two images are very close to each other, but they mean different things in the game.
 So, you know, two images can be super duper close, like an enemy can be here or slightly off, right?
 But if it's slightly off, it doesn't hit you.
 And therefore, you know, you're all good.
 Now, these two states are still pretty close because if you move a bit, you're likely to get hit.
 But sometimes a little bit of a change in image can mean actually a big change in game state and vice versa, which is actually even worse.
 A big change in image can mean like it doesn't matter.
 Like if everything in the image rotates around, but your agent still has nothing and is at the same place, it means nothing to you as a human.
 Yet an algorithm like this that whose goal it is to predict the future as accurately as possible, it will devote a lot of attention to accurately predict the future or predict variances in the future.
 Even though they might not be relevant.
 So in this in this task of or in this bottleneck of encoding everything into a very compact state, you might actually lose important information.
 And that means all of all of the like two states that are very, very far like need to be differentiated are going to be just the same in this representation.
 And that means your agent will never really learn because one is bad and one is good.
 So the mean reward is zero.
 And it says, well, when I get to that state, my mean reward is kind of zero and it's just kind of a big variance.
 And then the world model will never learn the difference because it has bigger things to worry about.
 So this is it's all very specific.
 And you'll see this in the in the loss term right here.
 So this is the loss function for learning the world model. And you can see they have an image reconstruction loss right here.
 This is a this is a cross entropy loss.
 So it's this is your approximation distribution.
 This is what really happened.
 Yeah, it's a it's kind of a probabilistic way of writing things.
 So these are cross entropy losses when you see log P of the expectation of under Q.
 They have a loss predicting the reward.
 They have a loss predicting the discount, which is mainly made for predicting when an episode ends in the in the imagined trajectory.
 And then they have this transition loss coupled with the entropy regularizer.
 So the transition loss is going to be for predicting these Z states and the entropy regularizer is for keeping the distribution in the Z states not peaked.
 So you want to kind of retain that stochasticity and this together you might recognize as the KL divergence between the P and Q.
 And that's this connection right here.
 So I'm going to minimize the KL, which is the same as saying I want this thing to be as accurate.
 I want I want I want these things to be as close as possible to each other, but the entropy should should still be given.
 And yeah, as you can see here, you can you can you can decompose that.
 So this is going to be this is going to be the KL divergence between the two distributions.
 I don't have a better way of explaining that without writing it down.
 You can already see they have a massive amount of hyperparameters, right?
 Like here's one, here's one, here's one, here's one, here's one.
 OK, so even within the KL divergence, they have actually two one hyperparameter for the KL divergence and one to trade off the entropy with the actual cross with the transition log loss with the cross entropy there.
 And they do ablations and see that that is really important that you have that trade off that you're able to make that trade off.
 And it's the same as the beta variational autoencoder, by the way.
 It's an entire paper about why you need an additional hyperparameter here.
 Like that's the entire paper of beta VAs, which I found funny.
 But, you know, it seems to be important.
 So you can see right here, this is KL balancing.
 So you have one, you have one term for making the prior close to the posterior, the prior being the one where you just see H and the posterior being the one where you see H and X.
 And you have another term for making the posterior close to the prior and you trade them off with these variables right here.
 Then the reinforcement learning itself, again, has a bunch of hyperparameters.
 So it is doing TD lambda learning.
 And you can look that up.
 TD lambda learning basically means you are here in your state and you're going to predict the value, sorry, the reward.
 Going to the next state and you're going to predict the value at that state.
 And then you're also going to predict from the same state the reward two steps forward and the value at that state.
 And you're also going to predict the reward three steps forward and the value at that state.
 And at the end, you're going to sum all of that up into one number that is kind of an aggregate of all of this.
 And that's going to be your prediction.
 That's what you regress on in your value predictor.
 And the actor tries to maximize that.
 So there's another parameter lambda that tells you how you aggregate these things.
 Right. And also H for how many steps you do that.
 There's also going to be in the actor loss function.
 They decided not only do they want the classic reinforce loss as you have,
 you actually want the straight through estimator of the distribution.
 And so a straight through estimator is when you want to backprop through sampled things.
 Normally, the reinforced gradients, what they do is if your actor outputs a distribution, let's say over three actions.
 Right. You don't all you can say is that I did action to here and it gave me seven reward.
 Right. So you want to make that more likely because seven is pretty good.
 Actually, you subtract the baseline.
 But, you know, let's say after the baseline, it's seven.
 So you simply act like you have a target distribution of this and scale it by seven.
 That's reinforced gradients. What you could also do is you could actually regress on directly through the softmax operation right here.
 Because this here is a sampling step. You cannot backprop through sampling steps.
 The way you can do it is that you you take the signal, the loss signal here,
 but you act as if this was your output and not this.
 OK, so you act as if you had made actions in proportion to their distribution and not actually sampled one particular action.
 This is going to give you a biased signal, but it has much lower variance.
 Whereas if you sample and then scale, it's going to be unbiased, but much higher variance.
 So they do these straight through estimators not only here, but actually also in this step up here.
 And you can see how that works in modern deep learning frameworks.
 So you have your distribution in terms of your logits.
 So what you can do is you sample from them and forward propagate should be the sample.
 Right. So the trick is to do plus and minus the same thing.
 So the forward propagation signal is simply your sample, as you can see right here.
 Now, the sample, this operation, it has no gradient.
 Oh, you can't see that it has no gradient.
 So the deep learning framework will simply not backprop through it.
 So if you were to just use the sample in your graph, you won't get a gradient.
 But what you can do is you can actually calculate the probabilities here,
 like the thing you want to back propagate, and then do plus that and minus stop gradient of that.
 You can see right here, this has no gradient.
 This has no gradient.
 So the gradient is going to be as if you had forward propagated this probes variable.
 But on the forward pass, the probes variable exactly cancels out with itself.
 And the sample is forward propagated.
 This is called a straight through estimator.
 It gives you a biased gradient, but much less variance than if you had to, you know,
 if you scale the sample like the reinforced gradients.
 So they use this in the world model.
 And they use this actually in the actor loss right here.
 And you can see there is another hyperparameter.
 Here is another hyperparameter.
 And then they have an entropy regularizer to facilitate exploration,
 which is normal, but gives you another regularizer.
 And not only do they have, sorry, hyperparameter,
 not only do they have these three additional hyperparameters,
 they scale two of them during training.
 So they now have a schedule to scale them.
 So this straight through estimator, they actually scale it to zero over the course of training.
 But yet two more hyperparameters, namely how fast you want to decay those things.
 So this whole thing is a giant bucket of hyperparameters.
 And so they say, while the unbiased reinforced gradients can help a better final solution.
 However, we find that using only reinforced gradients for optimizing the policy also works well.
 It might just not work as fast or as well, but it also works well.
 You know that in general, this is reinforcement learning, but this is a bit,
 you know, the amount of hyperparameters here is quite staggering.
 And I'm going to guess that this took a lot of work to even get off the ground.
 Right. So here you can see how this compares to other algorithms.
 Specifically blue here is Dreamer V2.
 And they do suggest a bunch of different things.
 So they have task median gamer normalized.
 So gamer is a professional human level gamer.
 And gamer normalized means you simply divide by what that professional gamer can do.
 So you can see that it can even exceed, you know, this gamer.
 So here is over 1.5 times over 55 different Atari games.
 Very good. However, these Atari games, some of them are actually unbounded.
 And in some of them, a machine can just be so much better than a human that usually these scores are dominated by very,
 very few games where the machine just excels, you know, hugely.
 And other games are like zero and both the median score and the mean score.
 They are not really meaningful.
 At least that's what this paper here argues.
 So they propose two modifications.
 So the first modification, actually, this is from a different paper as well, says you shouldn't normalize by, you know, kind of a professional gamer.
 You should actually normalize by the human world record.
 So this is record normalized. You can see it gives a cleaner score.
 And then they say, well, given that a few games still the the machine can just outperform humans so much.
 What you should do is actually you should never allow.
 So you just you should just clip the machine score at where the human world record is.
 So the reasoning behind this, I can imagine, is something like what's the difference between the human world record and the professional gamer world record?
 Well, the human world record, the professional gamer is already pretty good at gaming in general, let's say.
 But the human world record holder has probably figured out every single detail of that particular game and is pushing it with like exploits and whatnot.
 I don't know if you've seen legend like Ocarina of Time speed runs lately, but they're crazy.
 So that is going to be human world record.
 And it's probably going to be better to normalize by this because, you know, the machine will necessarily find these kind of exploits.
 They will it will probably find them as well.
 However, there are some things that where the machine you have to be where you have to be like pixel and microsecond accurate where the machine can do it and the human can't.
 So clipping it might make sense.
 I'm not really sure about this, like there's arguments to be made that you maybe shouldn't normalize by the human world record because, you know, you don't want to give credence to like exploits.
 But the gamer kind of represents more how the game is intended to be played.
 I don't know. They just suggest this new score just so happens to be that in this new score, they are, you know, other than here, they are just dominating at all time points.
 Yeah, let's let's leave them that they do a quite a number of ablations, especially they find out that, for example, if they do latent variables as categorical that outperforms Gaussian latent variables by a lot.
 So and that's, you know, that's kind of a reasoning why they use the categorical variables.
 The KL balancing simply means that additional parameter in the KL term, if they enable it, you can see it helps a lot.
 Image gradients. So when they they wonder, can we learn the world models from predicting images or from predicting rewards or from both?
 So they do both as a default.
 But if they leave away the image gradients, it doesn't work anymore.
 However, if they leave away the reward gradients, you can see it still works pretty well.
 Again, this is all quite Atari specific.
 And it also means that you can see right here, right?
 The Atari game lends itself to this kind of to exactly this kind of model.
 So how much this is a success for general reinforcement learning is questionable.
 However, what you can say is that if an environment lends itself to be world model learned by this kind of latent categorical variables,
 like so if the image state is going to be if changes in the image are going to be a good indicator of actual changes in relevant world variables,
 then you know, you might you might be very suited with a model like this.
 And so they compare this to other algorithms, for example, to use zero, which doesn't run on a single GPU.
 I think it is better, but it doesn't run on a single GPU.
 And it uses kind of a lot more Atari frames than the the dreamer algorithm.
 So you see again that you just need to find the correct category and you can be state of the art.
 So if this is like single GPU, Atari, no, I don't want to I don't want to dunk on this.
 This is pretty cool work.
 And if you look at the code, it took a lot of effort.
 Like you can see that from the code.
 OK, the last thing I want to look at is where does it succeed and where does it fail?
 So you can see a comparison, for example, dreamer V2 versus IQN or dreamer V2 versus Rainbow.
 And you can see and particularly interesting is where does it fail?
 And it fails in video pinball.
 And actually, I don't have it pulled up right here.
 But if you look it up, so if you look it up, you can probably see why.
 Because this video pinball thing.
 Thanks. Thanks, YouTube.
 This video pinball thing, it has a lot of changes in image without really doing much changes in the world state.
 So what actually matters is like this little tiny ball, this little tiny, you know, it's kind of a bunch of pixels.
 And the rest, you know, kind of moves around.
 And OK, maybe it doesn't move too much right here.
 But still, you know, there's this new cross that appears and so on.
 So a world model that learns to, you know, there's kind of flashes over the whole image, a world model that learns to accurately predict the world.
 Maybe is going to not focus so much on that little ball, but maybe is going to focus more on the rest of the image if that changes well.
 And also, you can see maybe the reward.
 Now, again, a flash, the reward doesn't change all too much.
 Yeah, it does, maybe.
 But, you know, any any time it bumps somewhere.
 So my hypothesis is going to be that in games where what actually matters consists of very few changes in the actual image.
 And there are lots of other big image changes that don't really matter so much for the immediate reward, maybe for the future, but not for the immediate.
 This algorithm is going to not be as good.
 And that is one example is this video pinball.
 And I might be wrong on this, but it's kind of a hypothesis.
 So the code for this is going to is available right here.
 Check it out as well as you should check out the blog post.
 They have a lot of ablations right here, as you can see, and graphs for the individual games turning off and on different variables.
 And you might as well give it a try if you have a reinforcement learning problem that has an environment similar to Atari.
 All right. That was everything I had to say for this pretty cool paper.
 Check it out. Bye bye.