Imagination-Augmented Agents for Deep Reinforcement Learning
deep learning reinforcement learning deep mind academic paper research
Commentary of
https://arxiv.org/abs/1707.06203

Abstract
We introduce Imagination-Augmented Agents (I2As), a novel architecture for deep reinforcement learning combining model-free and model-based aspects. In contrast to most existing model-based reinforcement learning and planning methods, which prescribe how a model should be used to arrive at a policy, I2As learn to interpret predictions from a learned environment model to construct implicit plans in arbitrary ways, by using the predictions as additional context in deep policy networks. I2As show improved data efficiency, performance, and robustness to model misspecification compared to several baselines.

Authors
Théophane Weber, Sébastien Racanière, David P. Reichert, Lars Buesing, Arthur Guez, Danilo Jimenez Rezende, Adria Puigdomènech Badia, Oriol Vinyals, Nicolas Heess, Yujia Li, Razvan Pascanu, Peter Battaglia, David Silver, Daan Wierstra
 Hi, today we're taking a look at Imagination Augmented Agents for deep reinforcement learning. This is a paper by DeepMind and has been in the news a bit recently, so we're going to have a look at what it's all about. Basically they claim that agents who have a model of the world perform better usually than agents who don't. But of course usually we don't have a model of the world, so they make the agent learn a model of the world which you can then use to plan. Now this learning of the model can of course be imperfect because it's learned and so they provide a way to work with imperfect environment models and combine them with a model-free approach. So what do we mean by models and model-free? Basically what you can say is if you have a model of the world, you have kind of a machine, say a box, and in this box you have a state S and you feed the state to the machine and you feed an action and the model of the world will tell you what did S' the new state is going to be. So this is in the case where you exactly know how your environment works. Now in a model-free approach what you would do is you would plan basically you would have a state and you would put that through some kind of a layered neural network and out would come what action should I take right now. So in the model-based approach you're trying to try out all these actions and tell you look which one gives me kind of a desired final state. And in the model-free approach you simply use the rewards to go directly and say here's my state, what should my action be? So this paper is a combination of both. The basic architecture is here, so let's start from the very right. We have two paths divided along this line. The final policy, so which actions you're going to take and what kind of values you can expect is going to be a result of two different models that are combined. There's a model-free path which means this is what we talked about. Simply here is the state and you simply feed it through this neural network thing, blah, blah, blah, blah, blah, blah, out comes a policy or an action you should take. But then there's also this other path and this is the imagination path. Basically consists a bunch of these rollout encoders and these rollout encoders is just the agent imagining the future. So the agent doing some actions and looking at how they will perform. So as this is done, there's this imagination core thingy. What this consists of is a policy network and an environment model. This environment model is really the core of the entire thing. So this environment model you basically learn from what you've seen so far. So far you've taken certain actions here in certain states. You use this to learn the environment model that gives you from one state the next state and the next reward. So that's what you learn. Of course also using neural networks and whatnot. You use that environment model to imagine the future. So here in this imagination core, basically you put in your state, you get out some new state and some reward. You feed the new state and you imagine another action. Of course the actions aren't random. The actions you also take via this thing. And this is where it loops all back. This is now a model free policy network that works with the environment model. So basically in your imagination you only use, if you look at the very right here, you only use this right path. Because your imagination doesn't need to be super exact or super well planned, you can use the model free approach that we kind of know kind of works for some problems. You use this to generate your actions that you imagine. And you use an environment model in order to look how these actions will play out. And that's how you imagine one step of the future. And you simply repeat this a couple of steps. And then you have an entire what's called a rollout, which consists of these pairs of states and rewards. And what you do then is you encode this rollout via this encoder, which is in this case an LSTM or something like this I think. You encode all these states into one vector, into one embedding basically for this rollout. And this embedding describes kind of this future imagined path. Of course, what you're going to hope is that somehow this encoding captures how you will do in the future and how good this will be. So these states and rewards. Once you have a couple of these rollouts, so once you've imagined a couple of different futures, you then aggregate them in this aggregator. I think in their case, they just concatenate these rollout encodings. And then you feed this too to the big aggregator on top. So the big aggregator on top can now combine the model free path and the imagined futures. So if the big aggregator thinks that the imagination isn't correct, it can resort to the model free path, but it can also think that maybe it's correct, or it can be kind of if it's sure it's correct, it can fully trust these rollouts and perform actions according to that. All of this is of course trained end to end. There's a tiny piece we haven't looked at yet, namely how this here, this policy network on the left is learned. And this is simply learned by, and I have to pay attention that I'm doing the right thing here. So you take this big thing here, your final policy network, and you perform, you kind of learn to copy its actions simply from the input. So from this model free input over here, you take this input and you take, excuse me, and you take the output of your big policy network and you try to simply make a neural network that copies the outputs given these inputs. And that's kind of your small policy network in here that's simply model free. So the loop closes in a way that you use your learned model to then again imagine the future. But of course for imagining the future, within imagining the future, you can't have another instance of this network because it would be infinite recursion. So you can only have a model free network. All right. That's it for the model. Of course, yeah, there's a couple of tricks and how to encode these things. Basically they perform experiments and this is maybe what you've seen in the media so far of this game. And this game is a game where you have to push around the brown boxes onto the red squares using the green avatar that you have. So this game is difficult because first of all, the levels are generated randomly. So there's no way you can like hard code anything. And second of all, if you push a box, say this box here, if you were to push it to the right into the corner, you would have no way of getting it out again. That's why I have to plan ahead and avoid such mistakes because they're not fixable. So once you make the mistakes, you can't go back and that's where planning comes in so handy. If you imagine this future and if your model is correct or approximately correct, then you can avoid such mistakes. Of course, that's the difficulty in this game and that's where the planning helps. Note that they don't code in how the game works. So all these models get is pixel input of the game and they have to kind of imagine the pixel output they're going to get. So that's increased difficulty. So technically the method is model free in the sense that there's really no coded model of the world, just the pixels. So they have performance comparisons where if you and I find this on the right here interesting, you can see according to the unrolled depth, so how much steps into the future you imagine. You can see it kind of flattens out after only about five steps. Whereas the game usually lasts for about 50 steps, they say. So only imagining five steps is already really helpful. What I don't like here is that they compare to what they say this copy model because this here is a standard model free comparison. So it's just a model free agent and of course, or not of course, but it performs worse right here. Because it has no imagination, but it also has less parameters. So they're trying to compare it to something with the same amount of parameters and say, oh, we have this copy model agent here. And what the copy model agent is doing is simply, for the environment model, it's the same architecture, but for the environment model, it simply predicts the output as the input. So it simply says, oh, you do this action, the environment is going to be exactly the same as it is now. And I don't like it because basically this entire branch here becomes rather useless. And so even though you have parameters in here, they're not useful. So to say that this is a comparison with the model of the same amount of parameters, I don't know, technically true. Another thing that they do is they pre-train the environment model with a model free agent. So first they code a model free agent, then they pre-train the environment model to then use with this agent. So it's not fully learned and I can imagine they tried and it didn't work. And this is how you get it to work. So they also experiment with imperfect models. So they train the environment model only imperfectly. And as you can see here, this is kind of the output you can get. Say you have duplicates, you have kind of errors, you have twice your character here, you have like boxes within the wall or all kinds of things. And they basically show that if you try to classically plan using these models, these bad models, you get nowhere. Basically this is a Monte Carlo sampler planner using a poor model and its performance degrades significantly from when you use the good model, which is right here. And the imagination agent is not affected by kind of the bad model, except that it takes kind of longer to reach its high inaccuracy. All right, so there's a couple of other experiments and a couple of Pac-Man experiments where they show you can learn one model to transfer kind of to play different games in this Pac-Man world. And that just works the more if you have very sparse rewards, which you can imagine, yes, if you need to plan then that's what you get. You get the ability to earn more sparse rewards because you can kind of look ahead. All right, so I think I'll conclude here with the discussion of this paper. I quite liked it and it's a cool method, combines many things and I'll see you next time.
 Hi, today we're taking a look at Imagination Augmented Agents for deep reinforcement learning.
 This is a paper by DeepMind and has been in the news a bit recently, so we're going to
 have a look at what it's all about.
 Basically they claim that agents who have a model of the world perform better usually
 than agents who don't.
 But of course usually we don't have a model of the world, so they make the agent learn
 a model of the world which you can then use to plan.
 Now this learning of the model can of course be imperfect because it's learned and so they
 provide a way to work with imperfect environment models and combine them with a model-free
 approach.
 So what do we mean by models and model-free?
 Basically what you can say is if you have a model of the world, you have kind of a machine,
 say a box, and in this box you have a state S and you feed the state to the machine and
 you feed an action and the model of the world will tell you what did S' the new state is
 going to be.
 So this is in the case where you exactly know how your environment works.
 Now in a model-free approach what you would do is you would plan basically you would have
 a state and you would put that through some kind of a layered neural network and out would
 come what action should I take right now.
 So in the model-based approach you're trying to try out all these actions and tell you
 look which one gives me kind of a desired final state.
 And in the model-free approach you simply use the rewards to go directly and say here's
 my state, what should my action be?
 So this paper is a combination of both.
 The basic architecture is here, so let's start from the very right.
 We have two paths divided along this line.
 The final policy, so which actions you're going to take and what kind of values you
 can expect is going to be a result of two different models that are combined.
 There's a model-free path which means this is what we talked about.
 Simply here is the state and you simply feed it through this neural network thing, blah,
 blah, blah, blah, blah, blah, out comes a policy or an action you should take.
 But then there's also this other path and this is the imagination path.
 Basically consists a bunch of these rollout encoders and these rollout encoders is just
 the agent imagining the future.
 So the agent doing some actions and looking at how they will perform.
 So as this is done, there's this imagination core thingy.
 What this consists of is a policy network and an environment model.
 This environment model is really the core of the entire thing.
 So this environment model you basically learn from what you've seen so far.
 So far you've taken certain actions here in certain states.
 You use this to learn the environment model that gives you from one state the next state
 and the next reward.
 So that's what you learn.
 Of course also using neural networks and whatnot.
 You use that environment model to imagine the future.
 So here in this imagination core, basically you put in your state, you get out some new
 state and some reward.
 You feed the new state and you imagine another action.
 Of course the actions aren't random.
 The actions you also take via this thing.
 And this is where it loops all back.
 This is now a model free policy network that works with the environment model.
 So basically in your imagination you only use, if you look at the very right here, you
 only use this right path.
 Because your imagination doesn't need to be super exact or super well planned, you can
 use the model free approach that we kind of know kind of works for some problems.
 You use this to generate your actions that you imagine.
 And you use an environment model in order to look how these actions will play out.
 And that's how you imagine one step of the future.
 And you simply repeat this a couple of steps.
 And then you have an entire what's called a rollout, which consists of these pairs of
 states and rewards.
 And what you do then is you encode this rollout via this encoder, which is in this case an
 LSTM or something like this I think.
 You encode all these states into one vector, into one embedding basically for this rollout.
 And this embedding describes kind of this future imagined path.
 Of course, what you're going to hope is that somehow this encoding captures how you will
 do in the future and how good this will be.
 So these states and rewards.
 Once you have a couple of these rollouts, so once you've imagined a couple of different
 futures, you then aggregate them in this aggregator.
 I think in their case, they just concatenate these rollout encodings.
 And then you feed this too to the big aggregator on top.
 So the big aggregator on top can now combine the model free path and the imagined futures.
 So if the big aggregator thinks that the imagination isn't correct, it can resort to the model
 free path, but it can also think that maybe it's correct, or it can be kind of if it's
 sure it's correct, it can fully trust these rollouts and perform actions according to
 that.
 All of this is of course trained end to end.
 There's a tiny piece we haven't looked at yet, namely how this here, this policy network
 on the left is learned.
 And this is simply learned by, and I have to pay attention that I'm doing the right
 thing here.
 So you take this big thing here, your final policy network, and you perform, you kind
 of learn to copy its actions simply from the input.
 So from this model free input over here, you take this input and you take, excuse me, and
 you take the output of your big policy network and you try to simply make a neural network
 that copies the outputs given these inputs.
 And that's kind of your small policy network in here that's simply model free.
 So the loop closes in a way that you use your learned model to then again imagine the future.
 But of course for imagining the future, within imagining the future, you can't have another
 instance of this network because it would be infinite recursion.
 So you can only have a model free network.
 All right.
 That's it for the model.
 Of course, yeah, there's a couple of tricks and how to encode these things.
 Basically they perform experiments and this is maybe what you've seen in the media so
 far of this game.
 And this game is a game where you have to push around the brown boxes onto the red squares
 using the green avatar that you have.
 So this game is difficult because first of all, the levels are generated randomly.
 So there's no way you can like hard code anything.
 And second of all, if you push a box, say this box here, if you were to push it to the
 right into the corner, you would have no way of getting it out again.
 That's why I have to plan ahead and avoid such mistakes because they're not fixable.
 So once you make the mistakes, you can't go back and that's where planning comes in so
 handy.
 If you imagine this future and if your model is correct or approximately correct, then
 you can avoid such mistakes.
 Of course, that's the difficulty in this game and that's where the planning helps.
 Note that they don't code in how the game works.
 So all these models get is pixel input of the game and they have to kind of imagine
 the pixel output they're going to get.
 So that's increased difficulty.
 So technically the method is model free in the sense that there's really no coded model
 of the world, just the pixels.
 So they have performance comparisons where if you and I find this on the right here interesting,
 you can see according to the unrolled depth, so how much steps into the future you imagine.
 You can see it kind of flattens out after only about five steps.
 Whereas the game usually lasts for about 50 steps, they say.
 So only imagining five steps is already really helpful.
 What I don't like here is that they compare to what they say this copy model because this
 here is a standard model free comparison.
 So it's just a model free agent and of course, or not of course, but it performs worse right
 here.
 Because it has no imagination, but it also has less parameters.
 So they're trying to compare it to something with the same amount of parameters and say,
 oh, we have this copy model agent here.
 And what the copy model agent is doing is simply, for the environment model, it's the
 same architecture, but for the environment model, it simply predicts the output as the
 input.
 So it simply says, oh, you do this action, the environment is going to be exactly the
 same as it is now.
 And I don't like it because basically this entire branch here becomes rather useless.
 And so even though you have parameters in here, they're not useful.
 So to say that this is a comparison with the model of the same amount of parameters, I
 don't know, technically true.
 Another thing that they do is they pre-train the environment model with a model free agent.
 So first they code a model free agent, then they pre-train the environment model to then
 use with this agent.
 So it's not fully learned and I can imagine they tried and it didn't work.
 And this is how you get it to work.
 So they also experiment with imperfect models.
 So they train the environment model only imperfectly.
 And as you can see here, this is kind of the output you can get.
 Say you have duplicates, you have kind of errors, you have twice your character here,
 you have like boxes within the wall or all kinds of things.
 And they basically show that if you try to classically plan using these models, these
 bad models, you get nowhere.
 Basically this is a Monte Carlo sampler planner using a poor model and its performance degrades
 significantly from when you use the good model, which is right here.
 And the imagination agent is not affected by kind of the bad model, except that it takes
 kind of longer to reach its high inaccuracy.
 All right, so there's a couple of other experiments and a couple of Pac-Man experiments where
 they show you can learn one model to transfer kind of to play different games in this Pac-Man
 world.
 And that just works the more if you have very sparse rewards, which you can imagine, yes,
 if you need to plan then that's what you get.
 You get the ability to earn more sparse rewards because you can kind of look ahead.
 All right, so I think I'll conclude here with the discussion of this paper.
 I quite liked it and it's a cool method, combines many things and I'll see you next time.