World Models
deep learning reinforcement learning deep reinforcement learning deep rl schmidhuber environment model imagination vae rnn lstm
Authors: David Ha, Jürgen Schmidhuber

Abstract:
We explore building generative neural network models of popular reinforcement learning environments. Our world model can be trained quickly in an unsupervised manner to learn a compressed spatial and temporal representation of the environment. By using features extracted from the world model as inputs to an agent, we can train a very compact and simple policy that can solve the required task. We can even train our agent entirely inside of its own hallucinated dream generated by its world model, and transfer this policy back into the actual environment.

https://arxiv.org/abs/1803.10122
 Hi, today we're looking at World Models by David Ha and Jürgen Schmidhuber. This is a paper that's concerned with reinforcement learning and especially with the problem of, say, you have an environment that you interact with and you kind of need to learn to act in it, but it could be, for example, very expensive to always query the environment. So let's say you have a robot and it needs to do something in the world, and you kind of, to have a robot execute something and then observe it, is quite expensive, costs electricity and so on. So you would like to sort of minimize how many times this happens. So here, searching for a good picture, they're concerned with problems, for example, like this. This is a race car simulator. There's an OpenAI gym environment for that. The other one that they use is a so-called like a doom experiment where, as you look at this, there's a couple of monsters and they're shooting fireballs at you and the task is just to kind of avoid the fireballs. So the entire point of the paper is that I don't actually need to interact with the environment in order to learn it. I can simply kind of learn a model of the environment and then learn using that model. So basically, I can learn how the environment works and then simply use my imagination of the environment, my model, in order to learn from that so I don't have to interact with the real environment anymore. So how do they do this? They do it in multiple stages. Here, first thing they do is they collect a bunch of samples from the environment. So they go to the environment, they simply do a random policy and then they collect a bunch of samples. I think the process is outlined down here somewhere. We saw it before. Here, collect 10,000 rollouts from a random policy. Next, they train a VAE here to kind of learn the environment. So that's where that comes in. This is all done in stages, not end-to-end. The VAE is simply a model that takes, in this case, a video frame here. It sends it through an encoder neural network to obtain what's called a latent representation, which is a much smaller dimensional representation. So if the image is 64 by 64 pixels, then the latent code could be as little as 100 or even 10 dimensional. So you see that there's quite a bit of compression going on. This is a variational autoencoder. It's not really important here that it's variational since the difference is the variational autoencoder is kind of a stochastic process, whereas the regular autoencoder isn't. But they introduce stochasticity later again. So it's not particularly important. So it's a variational autoencoder, which means they obtain a latent representation that defines distribution over outputs. So they send this sample from this latent distribution that they obtain, and then they feed this to the decoder. And the decoder kind of gives back what it thinks the encoder encoded. So the decoder tries to reconstruct as close as possible this original frame that was given to the encoder. But of course it can't because we've compressed it so much to this lower dimensional representation here. So it kind of does its best effort. So what you hope to achieve with this is that kind of the decoder learns, for example, there's always here. This is the ceiling right here. It's always gray. So basically, you shouldn't actually need to encode this in your Z. If it's always gray, the decoder should learn this by itself. So your hope is that the Z, the latent representation, will simply end up containing just the information that's kind of different or between the individual frames, which here I guess would be kind of the fireballs coming and your position relative to them. That's what's changing if you think about this environment. So your hope is that the latent representation captures only that, whereas all the static parts that are irrelevant or never change are kind of captured by the encoder and the decoder architecture by itself. So yeah, it's important to note the encoder and decoder are obviously always the same for all the frames, whereas the Z representation, of course, is there is one per frame, so each frame will give you a different Z. And that's so you can imagine how that works or how that's going to be useful. So they train this on like a randomly collected sample of the environment until they're confident they now have a good model of the environment. And then what they do next is they use this in order to train an RNN. So again, they kind of have their compression model of the environment. What they do now is they use these Z states you see here, here, here, here that they get from that. And they train how these latent representations evolve over time. So with an RNN here goes over time. So the RNN will always kind of predict what's the next state of the environment going to be. But importantly, maybe compared to environment models that we've discussed before in the, for example, imagination augmented agent paper, there we always try to directly predict the future pixels, so to say, of the future frame. Here, the environment model is over the latent representation. Of course, this means that the this is a much smaller space. So if your compression model is good, then this should be much easier to learn than, say, like a full end to end environment model. So this model learns how your latent states evolve over time, given your actions. So you can imagine the Z being an abstract representation of your state and then your action. And then this goes into the RNN and the RNN will predict what's the next latent representation. And there is what's called a temperature parameter to control the stochasticity. I've already told you this, there is a stochasticity built into this. So the RNN will simply output like some vector, what it thinks is the next thing going to be. And they don't use this directly as the next step, but they parameterize a kind of a mixture of Gaussian distributions coupled with a decoder here in order to give a random distribution over the next state. And they control the amount of randomness with the temperature parameter. They argue that this comes in handy later. So all right, so what do we have? We have a system that can compress the environment into what we would call an essential part. Every frame we extract what's important in that frame. Then next we have a model that can predict, given a state and an action, what's the next state going to be, the next latent state. So technically we now have an environment model, right, given a state. We can simply, given a state and a policy, we can simply use this model to roll forward. So the last component is the actual policy. And the actual policy here, as you can see, is in their case simply a linear model. The linear model will take the z, which is the latent representation of the current state, and the h, which is the current state of the RNN that models the environment over time. And it simply is a linear function of the two, gives you the action probabilities, or I guess the log-its of the actions. So it's a really, really simple controller over these things. And they do this in order to show that the main part of the work is being done by this environment model. And given the environment model, you only need very few parameters basically to then learn a policy. Here is what I said in a diagram. So the observation goes into the compression of the VAE, the latent representation of that goes into the RNN together with the hidden state from the last step. And this will output a new hidden state, which goes here into the controller, and we also directly take this z into the controller. And then from these two, we perform an action, which now we have a choice. It could go to the environment, right, give you the next observation, but also, or at the same time, since you kind of need to update your RNN, it can go here and update your RNN because it will need to predict the next hidden state. The thing is, we can also now leave away this path, which means we can simply take our RNN and our kind of imagine the next latent representation, put it through the decoder part of the VAE and use that as an observation. I hope this makes sense. It's rather intuitive, right? You have a model of the environment. You can simply use this instead of the real environment. So, there's a bit of pseudo code here, and they do a bunch of experiments, right? So, we're primarily interested, so they say, they see here, okay, our compression works, and this is the real frame, and this is the reconstructed frame, kind of looks, you know, captures the essence of what's going on. And I actually want to go down here, the Visdome experiment. So, what they do here in the car racing experiment is they kind of learn this entire thing, right? And then they learn a policy in the real world, in the environment, using this model up here, this procedure where they always go to the environment, and here is the exact experiment set up. So, first they collect, again, rollouts for a random policy, they train the VAE, they train the RNN, and then they learn the controller using the entire model, but in kind of the real world. So, they always interact with the environment, but because they also have their kind of latent representation of the observation, and not directly the observation, they get a higher score. And also, the policy that they use in the real environment transfers to the environment model. So, the policy they learn in the true environment, it transfers to the imagined, so if they use the imagined model as an environment, it also performs well. In the next experiment, they're going to try to do this the other way around. They're going to try to learn only using their model of the environment, and then see whether or not the policy transfers to the true environment. So, that's what they do here. They collect, again, a sample from the environment, they train the VAE, they train the RNN, and then they simply use this virtual environment, what they call it, in order to learn a policy, and at the end, they try to transfer, use the learn policy on the actual environment. And given the results, you see here, there we go. So, you see the kind of best it does, I would say, is about here, where the actual score is, you can see in this, and also in this setting, is higher than the kind of previous best algorithm in the OpenAI GIMP, when you go from virtual to actual. So, what this means is kind of, yeah, you can train using this imagined model, and then it will actually transfer, but there's a crucial thing, and that is this kind of temperature thing here. You can see a lot of times they actually don't manage to reach a good score, if this parameter is wrong. What does this parameter do? This parameter controls, as we discussed, the stochasticity of the model. So, basically, the environment model doesn't directly imagine a future state, but it imagines a distribution over future states. And the higher this parameter, the more stochastic this distribution is, basically the more uniform, I guess, the more entropy you have in these future states. We've seen this temperature parameter here. Which is important, because they go into length explaining why in this entire page here that we skipped. Here you see just text, there. Cheating the world model, which basically they say, okay, if you have a wrong model, if you have a model that's wrong of the environment, and you train a policy on it, necessarily, it's going to probably find a policy that exploits the wrongness of this model. So you might be able to walk through walls or fly or ignore the fireballs. Or basically, find that if you stand next to a wall, in your imagination, you'll never get hit. Something like this, which isn't true in the real world. So the policy will exploit that. And to counter this, they simply basically turn up this temperature parameter, giving them a more stochastic procedure. Meaning they imagine a lot of kind of different futures, and they train their policy on all of them, or in expectation over a sample of them. Which means that if the environment model is wrong, this kind of... I want to say if it's wrong, this corrects for it. It doesn't. But if it's wrong, you still sample different futures. So if it has one wrong future, you still have the other ones to kind of punish the policy, if it tries to exploit this one mistake. At least that's the reasoning behind it. So that's how they do this. You can interact with their trained environment models online somehow. They also give a kind of a look at what they would like to have. Instead of collecting the environment model from random rollout, they would try to train it, then to use it again to collect more data, to train more environment model, then use the environment, better environment model to train more the policy, and so on in a stepwise fashion. But they don't actually do it, they simply describe it. And the rest of the paper is a bit of related work and discussion. It's very prosaically written, kind of different from what you're used to if you read a lot of these papers. But yeah, I hope you can now you know what's going on and see you next time.
 Hi, today we're looking at World Models by David Ha and Jürgen Schmidhuber.
 This is a paper that's concerned with reinforcement learning and especially with the problem of,
 say, you have an environment that you interact with and you kind of need to learn to act in it,
 but it could be, for example, very expensive to always query the environment.
 So let's say you have a robot and it needs to do something in the world,
 and you kind of, to have a robot execute something and then observe it, is quite expensive, costs electricity and so on.
 So you would like to sort of minimize how many times this happens.
 So here, searching for a good picture, they're concerned with problems, for example, like this.
 This is a race car simulator. There's an OpenAI gym environment for that.
 The other one that they use is a so-called like a doom experiment where, as you look at this,
 there's a couple of monsters and they're shooting fireballs at you and the task is just to kind of avoid the fireballs.
 So the entire point of the paper is that I don't actually need to interact with the environment in order to learn it.
 I can simply kind of learn a model of the environment and then learn using that model.
 So basically, I can learn how the environment works and then simply use my imagination of the environment,
 my model, in order to learn from that so I don't have to interact with the real environment anymore.
 So how do they do this? They do it in multiple stages.
 Here, first thing they do is they collect a bunch of samples from the environment.
 So they go to the environment, they simply do a random policy and then they collect a bunch of samples.
 I think the process is outlined down here somewhere. We saw it before.
 Here, collect 10,000 rollouts from a random policy. Next, they train a VAE here to kind of learn the environment.
 So that's where that comes in. This is all done in stages, not end-to-end.
 The VAE is simply a model that takes, in this case, a video frame here.
 It sends it through an encoder neural network to obtain what's called a latent representation,
 which is a much smaller dimensional representation.
 So if the image is 64 by 64 pixels, then the latent code could be as little as 100 or even 10 dimensional.
 So you see that there's quite a bit of compression going on.
 This is a variational autoencoder. It's not really important here that it's variational since the difference is
 the variational autoencoder is kind of a stochastic process, whereas the regular autoencoder isn't.
 But they introduce stochasticity later again. So it's not particularly important.
 So it's a variational autoencoder, which means they obtain a latent representation that defines distribution over outputs.
 So they send this sample from this latent distribution that they obtain, and then they feed this to the decoder.
 And the decoder kind of gives back what it thinks the encoder encoded.
 So the decoder tries to reconstruct as close as possible this original frame that was given to the encoder.
 But of course it can't because we've compressed it so much to this lower dimensional representation here.
 So it kind of does its best effort.
 So what you hope to achieve with this is that kind of the decoder learns, for example, there's always here.
 This is the ceiling right here. It's always gray.
 So basically, you shouldn't actually need to encode this in your Z.
 If it's always gray, the decoder should learn this by itself.
 So your hope is that the Z, the latent representation, will simply end up containing just the information that's kind of different or between the individual frames,
 which here I guess would be kind of the fireballs coming and your position relative to them.
 That's what's changing if you think about this environment.
 So your hope is that the latent representation captures only that,
 whereas all the static parts that are irrelevant or never change are kind of captured by the encoder and the decoder architecture by itself.
 So yeah, it's important to note the encoder and decoder are obviously always the same for all the frames,
 whereas the Z representation, of course, is there is one per frame, so each frame will give you a different Z.
 And that's so you can imagine how that works or how that's going to be useful.
 So they train this on like a randomly collected sample of the environment until they're confident they now have a good model of the environment.
 And then what they do next is they use this in order to train an RNN.
 So again, they kind of have their compression model of the environment.
 What they do now is they use these Z states you see here, here, here, here that they get from that.
 And they train how these latent representations evolve over time.
 So with an RNN here goes over time.
 So the RNN will always kind of predict what's the next state of the environment going to be.
 But importantly, maybe compared to environment models that we've discussed before in the, for example,
 imagination augmented agent paper, there we always try to directly predict the future pixels, so to say, of the future frame.
 Here, the environment model is over the latent representation.
 Of course, this means that the this is a much smaller space.
 So if your compression model is good, then this should be much easier to learn than, say, like a full end to end environment model.
 So this model learns how your latent states evolve over time, given your actions.
 So you can imagine the Z being an abstract representation of your state and then your action.
 And then this goes into the RNN and the RNN will predict what's the next latent representation.
 And there is what's called a temperature parameter to control the stochasticity.
 I've already told you this, there is a stochasticity built into this.
 So the RNN will simply output like some vector, what it thinks is the next thing going to be.
 And they don't use this directly as the next step, but they parameterize a kind of a mixture of Gaussian distributions
 coupled with a decoder here in order to give a random distribution over the next state.
 And they control the amount of randomness with the temperature parameter.
 They argue that this comes in handy later.
 So all right, so what do we have?
 We have a system that can compress the environment into what we would call an essential part.
 Every frame we extract what's important in that frame.
 Then next we have a model that can predict, given a state and an action, what's the next state going to be, the next latent state.
 So technically we now have an environment model, right, given a state.
 We can simply, given a state and a policy, we can simply use this model to roll forward.
 So the last component is the actual policy.
 And the actual policy here, as you can see, is in their case simply a linear model.
 The linear model will take the z, which is the latent representation of the current state,
 and the h, which is the current state of the RNN that models the environment over time.
 And it simply is a linear function of the two, gives you the action probabilities, or I guess the log-its of the actions.
 So it's a really, really simple controller over these things.
 And they do this in order to show that the main part of the work is being done by this environment model.
 And given the environment model, you only need very few parameters basically to then learn a policy.
 Here is what I said in a diagram.
 So the observation goes into the compression of the VAE,
 the latent representation of that goes into the RNN together with the hidden state from the last step.
 And this will output a new hidden state, which goes here into the controller,
 and we also directly take this z into the controller.
 And then from these two, we perform an action, which now we have a choice.
 It could go to the environment, right, give you the next observation, but also,
 or at the same time, since you kind of need to update your RNN, it can go here
 and update your RNN because it will need to predict the next hidden state.
 The thing is, we can also now leave away this path,
 which means we can simply take our RNN and our kind of imagine the next latent representation,
 put it through the decoder part of the VAE and use that as an observation.
 I hope this makes sense. It's rather intuitive, right? You have a model of the environment.
 You can simply use this instead of the real environment.
 So, there's a bit of pseudo code here, and they do a bunch of experiments, right?
 So, we're primarily interested, so they say, they see here, okay, our compression works,
 and this is the real frame, and this is the reconstructed frame, kind of looks, you know,
 captures the essence of what's going on.
 And I actually want to go down here, the Visdome experiment.
 So, what they do here in the car racing experiment is they kind of learn this entire thing, right?
 And then they learn a policy in the real world, in the environment, using this model up here,
 this procedure where they always go to the environment, and here is the exact experiment set up.
 So, first they collect, again, rollouts for a random policy, they train the VAE, they train the RNN,
 and then they learn the controller using the entire model, but in kind of the real world.
 So, they always interact with the environment, but because they also have their kind of latent representation
 of the observation, and not directly the observation, they get a higher score.
 And also, the policy that they use in the real environment transfers to the environment model.
 So, the policy they learn in the true environment, it transfers to the imagined,
 so if they use the imagined model as an environment, it also performs well.
 In the next experiment, they're going to try to do this the other way around.
 They're going to try to learn only using their model of the environment,
 and then see whether or not the policy transfers to the true environment.
 So, that's what they do here. They collect, again, a sample from the environment,
 they train the VAE, they train the RNN, and then they simply use this virtual environment,
 what they call it, in order to learn a policy, and at the end, they try to transfer,
 use the learn policy on the actual environment.
 And given the results, you see here, there we go.
 So, you see the kind of best it does, I would say, is about here,
 where the actual score is, you can see in this, and also in this setting,
 is higher than the kind of previous best algorithm in the OpenAI GIMP,
 when you go from virtual to actual.
 So, what this means is kind of, yeah, you can train using this imagined model,
 and then it will actually transfer, but there's a crucial thing,
 and that is this kind of temperature thing here.
 You can see a lot of times they actually don't manage to reach a good score,
 if this parameter is wrong. What does this parameter do?
 This parameter controls, as we discussed, the stochasticity of the model.
 So, basically, the environment model doesn't directly imagine a future state,
 but it imagines a distribution over future states.
 And the higher this parameter, the more stochastic this distribution is,
 basically the more uniform, I guess, the more entropy you have in these future states.
 We've seen this temperature parameter here.
 Which is important, because they go into length explaining why in this entire page here that we skipped.
 Here you see just text, there.
 Cheating the world model, which basically they say, okay, if you have a wrong model,
 if you have a model that's wrong of the environment, and you train a policy on it, necessarily,
 it's going to probably find a policy that exploits the wrongness of this model.
 So you might be able to walk through walls or fly or ignore the fireballs.
 Or basically, find that if you stand next to a wall, in your imagination, you'll never get hit.
 Something like this, which isn't true in the real world.
 So the policy will exploit that.
 And to counter this, they simply basically turn up this temperature parameter,
 giving them a more stochastic procedure.
 Meaning they imagine a lot of kind of different futures,
 and they train their policy on all of them, or in expectation over a sample of them.
 Which means that if the environment model is wrong, this kind of...
 I want to say if it's wrong, this corrects for it. It doesn't.
 But if it's wrong, you still sample different futures.
 So if it has one wrong future, you still have the other ones to kind of punish the policy,
 if it tries to exploit this one mistake. At least that's the reasoning behind it.
 So that's how they do this.
 You can interact with their trained environment models online somehow.
 They also give a kind of a look at what they would like to have.
 Instead of collecting the environment model from random rollout,
 they would try to train it, then to use it again to collect more data,
 to train more environment model, then use the environment,
 better environment model to train more the policy, and so on in a stepwise fashion.
 But they don't actually do it, they simply describe it.
 And the rest of the paper is a bit of related work and discussion.
 It's very prosaically written, kind of different from what you're used to if you read a lot of these papers.
 But yeah, I hope you can now you know what's going on and see you next time.