Can Wikipedia Help Offline Reinforcement Learning? (Paper Explained)
deep learning machine learning arxiv explained neural networks ai artificial intelligence paper
#wikipedia #reinforcementlearning #languagemodels

Transformers have come to overtake many domain-targeted custom models in a wide variety of fields, such as Natural Language Processing, Computer Vision, Generative Modelling, and recently also Reinforcement Learning. This paper looks at the Decision Transformer and shows that, surprisingly, pre-training the model on a language-modelling task significantly boosts its performance on Offline Reinforcement Learning. The resulting model achieves higher scores, can get away with less parameters, and exhibits superior scaling properties. This raises many questions about the fundamental connection between the domains of language and RL.

OUTLINE:
0:00 - Intro
1:35 - Paper Overview
7:35 - Offline Reinforcement Learning as Sequence Modelling
12:00 - Input Embedding Alignment & other additions
16:50 - Main experimental results
20:45 - Analysis of the attention patterns across models
32:25 - More experimental results (scaling properties, ablations, etc.)
37:30 - Final thoughts

Paper: https://arxiv.org/abs/2201.12122
Code: https://github.com/machelreid/can-wikipedia-help-offline-rl
My Video on Decision Transformer: https://youtu.be/-buULmf7dec

Abstract:
Fine-tuning reinforcement learning (RL) models has been challenging because of a lack of large scale off-the-shelf datasets as well as high variance in transferability among different environments. Recent work has looked at tackling offline RL from the perspective of sequence modeling with improved results as result of the introduction of the Transformer architecture. However, when the model is trained from scratch, it suffers from slow convergence speeds. In this paper, we look to take advantage of this formulation of reinforcement learning as sequence modeling and investigate the transferability of pre-trained sequence models on other domains (vision, language) when finetuned on offline RL tasks (control, games). To this end, we also propose techniques to improve transfer between these domains. Results show consistent performance gains in terms of both convergence speed and reward on a variety of environments, accelerating training by 3-6x and achieving state-of-the-art performance in a variety of tasks using Wikipedia-pretrained and GPT2 language models. We hope that this work not only brings light to the potentials of leveraging generic sequence modeling techniques and pre-trained models for RL, but also inspires future work on sharing knowledge between generative modeling tasks of completely different domains.

Authors: Machel Reid, Yutaro Yamada, Shixiang Shane Gu

Links:
Merch: http://store.ykilcher.com
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
LinkedIn: https://www.linkedin.com/in/ykilcher
BiliBili: https://space.bilibili.com/2017636191

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
 Can Wikipedia help offline reinforcement learning? This is the title of the paper that we're going to look at today. This paper is borderline preposterous in the results that it presents. Language model pre-training helps reinforcement learning, which is crazy. The two domains have almost nothing in common with each other, and yet there seems to be some transfer from language to reinforcement learning. This is not just about pre-training on any old task. The authors here have tried various things, and there seems to be something special about language. So here is how the video looks. This video right here is a paper review. It presents me going through the paper together with you, explaining the paper, explaining what I think about the paper, what kind of questions I have, and so on. After this video, you'll have a good understanding of what the paper contains, what its main claims are, maybe also what I think its weaknesses are. In the next video, which will be released tomorrow, I will interview the authors of this paper, which is very cool. The authors will have seen my review and are directly able to respond to criticisms, to any questions that are raised there, and this is so valuable. We're able to directly dive in and get you the best possible insight into the behind-the-scenes stuff and into the research process about this paper. I invite you to watch both videos, although feel free to choose whichever one you like most. As always, let me know what you think in the comments, leave a like if you do, and I'll see you around. Bye. Hello there. Today, we're going to look at Can Wikipedia Help Offline Reinforcement Learning by Michelle Reed, Yutaro Yamada, and Shixiang Shenggu. This paper is a special paper because it very counter-intuitively trains a language model. So it pre-trains a transformer to do language modeling, for example, Wikipedia text modeling. As you can see right here, language goes in, it does next word prediction, like you're used to from a language model like GPT-2, GPT-3, and so on. And then it takes that transformer and fine-tunes it to trajectory modeling. This is a special subfield of offline reinforcement learning where decision transformers have recently been introduced. So in offline reinforcement learning, you have some data set of trajectories, and then you try to do reinforcement learning just given on that data set. It turns out that if you pre-train something on language and then fine-tune it on these trajectories, that will turn out to be a much better model, like a much more performant model for getting you good reward at the end than if you just train this trajectory model here from scratch, which is very counter-intuitive because it means that somehow the language modeling task, like the language model pre-training, has a beneficial effect on the reinforcement learning tasks that comes later. To note that the reinforcement learning task has nothing to do with language. And even more special, they also try a bunch of other things. Most notably, they try to pre-train the image GPT model, and that does not result in good performance. So it's not just the fact that you have pre-trained on something, and it is really a very special result. So we're going to dive into the paper right here. The setup is fairly simple, and then there is a series of experiments that try to investigate this phenomenon. So they say that the offline reinforcement learning, as I said, has been seen as a sequence-to-sequence model. And I've already pre-annotated some stuff right here. Let me know how you like that. I thought I'd do it in this way. So I have the green, that is the current one, and the yellow is from my previous escapades on this paper. So they go into offline reinforcement learning, and that is being framed as simply supervised learning to fit return-augmented trajectories in an offline data set. What do they mean? They mean the setup of the decision transformer. I've made a video on the decision transformer. If you want to look at that, you can go after you watch this video. So the decision transformer says, well, see, you are an agent somehow. There is an environment. There is some interaction between the agent and the environment. And in offline reinforcement learning, we usually have a data set of this. So someone else has performed this, and they've distilled all the episodes into this data set. And their goal is to learn just from the data set. We can't actually interact with the environment. So in the data set, there are a number of trajectories, trajectories of the agent interacting with the environment. There's always some sort of a state coming back from the environment or an observation, if you will. The agent always gives some sort of an action back, and then there is a reward and the next state coming from the environment and so on. So that is naturally a sequence. And the sequence is there is a state, then there is an action, then there is a reward and a new state, then there is an action again, and then there is a reward and a new state. So this is a sequence. And since I have a data set of these sequences, I might as well throw that into a big transformer to do sequence modeling. Now, this has its own problems, which I've all discussed in the decision transformer video. For example, if the transformer has a context length of four, it cannot conceivably look back further than that, which is a classic problem in reinforcement learning, how to look back and forward infinite times. The decision transformer has the limited context window. It has sort of the caveats of language modeling. However, we understand language modeling very well, and therefore, we are quite able to do that. There is one modification that they do. What they do is they transform the rewards right here. They don't let the model model the rewards, they let it model the rewards to go. We're going to see that in just a bit. This here is interesting. What they say is that we look at whether transformer based pre-trained language models are able to be adapted to standard offline reinforcement learning tasks that have no relations to language. I've already told you that this is going to work out fairly well. That's the special message of this paper. They show consistent performance gains and significantly faster convergence. By faster convergence, they mean that a convergence point, like a non-improving the loss anymore, is reached after much many fewer steps than if you were to train from scratch, which makes sense for pre-training if it's in the same domain. But given that the pre-training is a completely different domain than the fine tuning, that is still a just a special thing. So here is how we're going to frame the problem. And if you've watched the decision transformer video, this should be familiar to you. We model a episode as a sequence in the following manner. This is almost as we've seen it, except the rewards right here. They are not individual rewards, but they are this thing right here, the sum of all the rewards at this end, the next steps, which they call the returns to go. So this, for example, says from here until the end of the episode, I'm going to gather 50 reward. Now, maybe you're in this state and you made an action that gave you a reward of one. So then this here would be 49. So you'd say, well, from here on out, I'm going to make 49 reward and so on. So the benefit of this is that at inference time, you can just put like a really high reward right here. So at inference time, you would always you would model these things you would get from the environment. So you'd start out with like just a big reward right here, just whatever the maximum you've observed plus 10% or something to just encourage your model to go very high. And you plug the state in here that the environment has given you and you let the model produce this one. So it's important that at training time, we do sequence modeling, really model the sequence of returns and state and action as a GPT, like next token prediction. However, at inference time, we obviously only predict the action and the environment is going to give us these two things, or the environment is going to give us the reward. And then we simply subtract the reward from the previous returns to go. And we plug that in here. And then we plug in the state we got from the environment. We let the model predict the next action right here and so on. So this is very cool because much like something like upside down reinforcement learning, this is conditioned on it like a desired reward. This also has advantages and disadvantages, but the advantage is we can control the reward we want at inference time. So we don't always have to go for a high, super high reward, but we can. Yeah, so this is the setup. You don't actually need to understand much more. But what we're going to do is we're going to model this as a sequence in our data set, and then at inference time, we just put some high returns to go. And that's it. We're going to use a transformer for that, for the sequence model. And they're going to use a bunch of different models right here. For example, GPT-2 small, which is a pre-trained model. They also pre-train their own that they call chibi-t, which is the same size. So that is the same parameter count as the original decision transformer to make it comparable to them. So the decision transformer is the one that introduced this transformer as sequence model for reinforcement learning. And they are going to see this chibi-t model has the exact same amount of parameters as the decision transformer, so they can directly compare what the language pre-training is going to gain them in the same model. They also use CLIP. However, they only, as far as I am aware, they only use the text encoder part of CLIP because that's an autoregressive model, which can do the sequence modeling. And they use image GPT, which is an autoregressive model that goes via image tokens. So an image GPT, it would split up the image into, no, not pixels, but chunks, I believe, either chunks or pixels. I don't even remember. And it would do the sequence model, essentially go through the image like this, and then like this, and then like this. So it framed the image as a sequence of either patches or pixels and go through it as a sequence model. So that's a sequence model too. We can pre-train it, and then we can apply it to this space. They do various things right here, other than just language modeling, sorry, other than just language or sequence prediction. Let's call that sequence prediction right here. Other than just sequence prediction for the reinforcement learning data, they do two more things. First of all, they want to align the input representations. So they have a set of language embeddings, which comes from the pre-training data set. Now, obviously, the pre-training data set has a tokenizer. That tokenizer generates tokens from the text, and every one of these tokens will have one of these embeddings associated with it. So V is the vocabulary size. However, obviously, in the reinforcement learning settings there, we don't have the same tokens. We don't have the same input modality even. And therefore, we don't need a tokenizer because it's already tokenized, right? Each of these things right here is a token. However, what we do need is now a new vocabulary, not a new vocabulary, but a new embedding matrix, so to say. So we have a different amount of tokens, so from one to the 3N tokens. And what we're going to want to do is, what they say at least, we want to have a set of linear projections that will map the return embeddings, the action embeddings, and the state embeddings to be very close in their cosine similarity to some embedding vector in the original setting. So that means they want to force, not force, they want to encourage the model to sort of reuse the embeddings that it used during the language model training. So for each of the input embeddings, they're going to find the maximum, the closest, nearest neighbor in cosine space of the embeddings of the original vocabulary. And then they're going to encourage the input embedding, the new input embedding, to be closer to that. So that is just a loss that they add during training. So you can see right here, this is the loss for the language or the sequence modeling decision transformer objective. This is the loss that encourages the embeddings to be close to the original language embeddings or to one of the original language embeddings. And this loss right here is the continuation of language modeling. So during training of the sequence prediction for reinforcement learning, they additionally also do, that's what they call language model co-training, continuing to train jointly on language modeling and trajectory modeling. This allows us to encourage, this allows us to encouraging, it probably should be encourage, the model's transformer backbone to be able to handle both language and trajectory simultaneously. OK, maybe it helps. This seems either like an idea that had been had at some point or something they had to put in after the fact just to make it even a bit better, or because maybe it didn't work, though they ablated it at some point. And it also works without. So that's almost it. Yeah, they describe a little bit their baselines and their setup. I was a bit confused here. It says it's a batch size of 65,000 tokens, which I don't, like, I don't, is that, I don't, batch size is usually not in tokens, like the sequence length would be in tokens. But in any case, they say for our additional objectives, we decay lambda 1 and lambda 2 to reach 0 after 5,000 steps. We tuned the initial values of lambda 1 and lambda 2. And, you know, these seem, they seem reasonable. But the fact that you have to, like, decay the additional losses after x many steps and so on, it points to a little bit of brittleness in them. And I'm not sure always how brittle these things are, because reinforcement learning is traditionally kind of a very brittle field. So the main, the main results we have right here, the top one is four games in Atari. The bottom one is, I believe, three environments in the, in the OpenAI gym that are, oh, sorry, the, this is a data set, the D4RL data set. All of this is offline reinforcement learning. On top, you also have the 1% DQN replay Atari data set. So as you can see, in many cases, the, both the Chibi-T and the GPT-2, by the way, GPT-2 is a lot larger than, so this is a lot larger in parameters than the Chibi-T model, and therefore also than the decision transformer model. So just, just saying that. So here, the pre-trained models outperform the other ones in quite a few tasks. However, there is also Qbert, where they still do outperform the decision transformer, as you can see. But the, they're, one of the baselines is just a lot stronger. The other baselines are just useless. That's kind of what I mean when I complain about, when I complain about reinforcement learning is that it is just weird. Like a bit of a different environment can make a large difference. But as you can see, the pre-language pre-trained models consistently outperform the decision transformer models. Also something to note right here, this is mean and variance across three seeds. So this is variance, I'm going to guess they mean standard deviation. And that is like a large number. So if that's the standard deviation, then the differences to the decision transformer, they are well, well within that. And that means, I mean, it is visible that across experiments, we see the same trend, right? That gives it credence. But also, this just seems extremely noisy. And yeah, I'm not going to say I'm going to sound like reviewer 2 when I say, well, you should make more experiments to estimate or to get smaller error bars. But it just seems like, I don't know, it seems like results that you can't really put a lot of weight on because they're very noisy. However, a bit, like a little bit less noisy are the experiments here on the bottom. You can see that the standard deviations here are quite a bit smaller than on top. That's also three seeds. I like how they wrote the number three here and the word three right here. That is just something that you never see until someone points it out. You can also see right here that the decision transformer, for example, is rather consistently outperformed. What's also interesting is that image GPT just sucks. Like you can see right here, like it just, it doesn't get anywhere on any of these tasks. Also clip very often underperforms. You can see, for example, here clip underperforms. And they do have some hypotheses on that. That being said, there are still a lot of times where the baselines here are quite a bit better or just better than all of these transformer-based models. So just pointing that out. Yeah. They do also analyze, and this I find really interesting, the attention pattern between the GPT-2 pre-trained model, the image GPT pre-trained model, and what I understand is a randomly initialized model that has just been fine-tuned. Yeah, randomly initialized model that has just been fine-tuned. So there's no pre-training. So all of these models are fine-tuned, but the random one hasn't been pre-trained. Interestingly, if you look at GPT-2, you can see these bands right here. And the bands are always in the distance of 3. So there's always 3 distance. Now, 3 should be an interesting number if you remember how the sequence is made right here. So there is always going to be 1, 2, 3. These tokens come in packets of 3, right? Their next return would be here. The next state would be here. The next action would be here. So every token in this attention pattern is most focused on multiples of 3 behind it in order to predict the next token. So there's always a lag of attention to multiples of 3, which means that essentially, if I want to predict the next return, probably the last return ends are the most important. If I want to predict the next action, maybe the last actions are important. This might also be a property of the environment. This is on Hopper. So on these continuous control tasks, I guess it's very often the case that I'm just going to repeat an action for a while if I want to achieve some goal. I don't know the frame rate exactly of these things. However, that seems to be something that is rather maybe viable to do. And therefore, looking at the last action can give me a lot of clues about the next action. Looking at the last state can give me a lot of clues about the next state. I would wonder how this changes if it's something like, well, I don't even know, anywhere where I don't naturally repeat my last action often. You can see this is the early layer. Then in the middle layer, the GPT-2, it seems to sort of focus on particular states that seem to be important, as you can see right here. So this is where the attention comes from. This is where it goes to. And you can see that it kind of decides that particular states are important, and it kind of remains at that. So it selects a few states that, or a few tokens that it chooses to attend particularly to. In contrast to that, our image GPT seems to have a large recency bias. So if you see this right here, there's really this band right here, which essentially means that every token attends to kind of the few tokens behind it in order to predict it. Then, well, the question is, is it even worth looking at stuff further down? Because this model clearly doesn't learn at all. So I would consider this and this just to be kind of random noise. The early layers might be interesting, though, because there is kind of a pattern. And maybe that is influenced by the pre-training. So in image GPT, since you have your image, and maybe it's in chunks, maybe it's in pixels, but I can imagine that if I want to predict a particular chunk, that maybe the last few that I've predicted, unless I cross a boundary right here and go one line down, the last few that I predicted are or might be particularly worth looking at. And rather distant chunks might be not worth looking at very much, other than in language modeling, where I often have to go a little bit more across the distance and the exact neighboring words might not be as important. So that might explain why image GPT has this particular recency bias pattern in its attention. What's also interesting is that the randomly initialized model, look at that. This is another interesting pattern. And you can see that it's very much the same as in the GPT example happens, except much more extreme. So you have these rows. For example, this row right here, you can see there is a hard attention for three back. It's really hard attention. Then there are rows where you can see right here, there is always these two, and then these two, and then these two, with particular attention on the first one and then also slight attention on the second one. And it's a special pattern. So no, I'm one off, sorry, in the one above. So this is the hard three. Then the one below is the, I'm going to call it the soft three. So there is one strong one and one weak one. And then the one even below that, there is one semi-strong, one weak, and one really weak. So what's happening? I'm not exactly, so what I don't know here is which of these tokens is returns, which ones is state, and which one is action. But I'm going to just guess, and I might be totally wrong right here, that the very strong bias here, that is going to be the returns to go, which would only focus on the last returns to go. And then after that would be the state tokens. So what the state tokens would do is, and you can see this, I'm just going to, so let's say this is the returns to go, the bright ones. And you can see that in the state tokens, there is, actually there is one missing here on the diagonal. So this diagonal one here is just completely blank, which means that it just kind of ignores the token behind it, which is the reward, right? So what it cares about is the last state, and it also cares about the last action maybe. I don't know how to interpret that very much otherwise. So if I want to predict the next state, I'm going to care about the last state, and the action after that, maybe that makes sense. If I want to predict the next action, then I might be able to care about all of the stuff beforehand a little bit. Again, I don't know if I'm interpreting this correctly. However, what I am able to say is that there is a very, very structured attention right here. There is this pattern of three is very prevalent, and it is in general very, very structured. So this seems to be actually the best kind of attention, right? It is very structured in the way it looks at the information. It learns exactly, aha, there is a structure to it. I'm going to attend to the different parts in this different structure. However, my hypothesis is, and that is not super duper discussed in the paper. I mean, it is discussed, but my hypothesis is that this bias here, it might be almost too strong. It might learn the exact structure of this stuff, but it might be too strong, and it might miss information. Because it, for example, says, well, I don't need to know anything in between here, because the most relevant thing for predicting the return is the last return, and therefore, I'm not even going to look at other stuff. Whereas the language model pre-training just kind of acts as a regularizer that says, well, you should maybe look at all of the stuff, even though you don't find it super useful in this particular data. Now, one thing that I didn't point out in the video that I wanted to point out right now is that if you look at GPT-2 at the very left column, what it does is it focuses particularly on the returns to go steps. It doesn't matter which step it is at. It always kind of looks back at the very first token, which is the returns to go of the whole episode, and among other things, also at the second and the third returns to go token. And this is important, because the returns to go is kind of an indicator of how the episode's going to go along. If the returns to go are low, it means that entirely different episode paths should be chosen in order to achieve that reward. Whereas if the returns to go is high, then I would have to do different actions to get that returns to go. So it makes a lot of sense to look at the returns to go tokens. And rather than, whereas you can see in the right hand column, the randomly initialized thing, it only really focuses on the returns to go in these middle layers whenever it needs to predict the next return. And so it's much more diffuse, and it doesn't condition all of what it does a lot on these returns, where it makes total sense to do that. Because in one instance, the language modeling is just sampling any sort of high likelihood trajectory. However, additionally in the GPT-2 case, it is almost like conditioning that sampling on the most relevant information that distinguishes between the different futures. I hope that makes sense. Why a model that would learn to focus in particular on this information would be better at sampling appropriate trajectories for the current episode. All right, back to my comments in the past. We know that language models retain large parts of their pre-training even during fine tuning. So the language modeling thing might just be like a very good prior. And I wonder if we could build these types of priors into the decision transformers if we didn't do language model pre-training, but just as sort of like a bias or a regularizer or something like this. Yeah, you can see that through the random attention at the end, you do not get this focus as you get with the language model thing that it focuses on particularly interesting last states, but you'd rather you do get like an attention matrix in the last layer that is kind of diffuse and sort of similar to the image GPT that just doesn't work at all. So yeah, that would be my maybe postulation that maybe it is possible to achieve the same effect by introducing the correct regularizers. However, I don't know. So they look at a few other things which I just quickly wanna go through. Because they have pre-trained, they can demonstrate that their model converges much more quickly. So instead of like three hours, their models of the same size needs 43 minutes and their model that is a lot larger, I believe GPT-2 is 144 times larger. It only uses an hour and 27 minutes. So still half of the time than this decision transformer. Now, I also wonder whether they have based their code base on the decision transformer or whether some of this difference is also due to just kind of like a better implementation. So yeah, that is that. They have some analysis right here. For example, they say they hypothesize that a generative training objective is useful. That's how they explain why CLIP might not be as effective because CLIP is ultimately a discriminative objective or a contrastive objective. They also say that there are underlying similarities between language modeling and trajectory modeling where there is a large difference between image modeling and trajectory modeling, which is it's a hypothesis. They say, yeah, there is the language modeling has a natural sequential nature. The versus image modeling is kind of a forced autoregressive task. I agree with that, but I'm not sure if there's really due to like language being particularly similar or whether, as I said it might just be a good prior. This would be an interesting question to investigate. And it might ultimately turn out to be the same thing. So, you know, interestingly the context size doesn't really matter. You can see right here, if they increase the context size they do get worse actually. So yeah, that's worse. It's just more noisy, which is special which actually means that these models aren't appropriate yet or we haven't really figured out how to appropriately use them yet, right? More information shouldn't necessarily give you less of a reward unless I guess maybe you have a fixed size data set and therefore you have less training data points. So maybe that's an effect of that. Interestingly, the pre-trained models, they do scale better which I guess you might have expected if you've been in deep learning the last few years but if you just take a decision transformer it will overfit after a while if you scale it up. So these are millions of parameters. You scale it up, it actually gets worse. Actually not sure if that's overfitting or just, you know it gets too big and then the average reward decreases. However, if you pre-train first, then it can handle and it will actually increase with more data. Interesting would be to see if that at some point actually declines again or if that sort of holds up if the language model pre-training for which there is like infinite data, right? In language model pre-training, you can get infinite data and therefore it could be that this just kind of gets you diminishing returns but not ever come down again. Yeah. They also experiment with freezing parameters and they say that this drastically reduces performance. So if they only train, if they only train, what do you say? Only action state and return projections being trained. So only this alignment of this projection of the, the projection of the token embeddings are being trained. That doesn't work much, which is also surprising because there is a lot of work that kind of shows that you don't have to train many parameters of these transformer models to effectively transform or transfer them from one task to the other. They say that this might be, this might be the case that this might be due to the task of generative modeling being harder as opposed to discriminative classification where this was previously applied. They have a lot of, yeah, they pose a lot of hypotheses here of why things might be and I feel each one of them could be its own research paper. Yeah, I'm gonna leave it at that for the paper explanation. I hope you got a little bit an intuition. I still find it very, very special and very cool that this even works and I think it's an, it's an like a sign of the times of our models just becoming the same models for all modalities. This would not even have been possible a few years ago where every modality would use very different models like CNN for images and RNNs for language and so on. Although RNNs were used for RL already, but given that our models converge and we're getting, we're learning so much more, this type of research is really cool. Yeah, let me know what you think is, have we overlooked something right here, like something that could easily explain why this works and gives good results that just no one kinda sees or are there more applications for this? Let us know what you think and bye bye.
 Can Wikipedia help offline reinforcement learning?
 This is the title of the paper that we're going to look at today.
 This paper is borderline preposterous in the results that it presents.
 Language model pre-training helps reinforcement learning, which is crazy.
 The two domains have almost nothing in common with each other,
 and yet there seems to be some transfer from language to reinforcement learning.
 This is not just about pre-training on any old task.
 The authors here have tried various things,
 and there seems to be something special about language.
 So here is how the video looks.
 This video right here is a paper review.
 It presents me going through the paper together with you, explaining the paper,
 explaining what I think about the paper, what kind of questions I have, and so on.
 After this video, you'll have a good understanding of what the paper contains,
 what its main claims are, maybe also what I think its weaknesses are.
 In the next video, which will be released tomorrow,
 I will interview the authors of this paper, which is very cool.
 The authors will have seen my review and are directly able to respond to criticisms,
 to any questions that are raised there, and this is so valuable.
 We're able to directly dive in and get you the best possible insight
 into the behind-the-scenes stuff and into the research process about this paper.
 I invite you to watch both videos,
 although feel free to choose whichever one you like most.
 As always, let me know what you think in the comments,
 leave a like if you do, and I'll see you around. Bye.
 Hello there.
 Today, we're going to look at Can Wikipedia Help Offline Reinforcement Learning
 by Michelle Reed, Yutaro Yamada, and Shixiang Shenggu.
 This paper is a special paper because it very counter-intuitively trains a language model.
 So it pre-trains a transformer to do language modeling, for example, Wikipedia text modeling.
 As you can see right here, language goes in, it does next word prediction,
 like you're used to from a language model like GPT-2, GPT-3, and so on.
 And then it takes that transformer and fine-tunes it to trajectory modeling.
 This is a special subfield of offline reinforcement learning
 where decision transformers have recently been introduced.
 So in offline reinforcement learning, you have some data set of trajectories,
 and then you try to do reinforcement learning just given on that data set.
 It turns out that if you pre-train something on language
 and then fine-tune it on these trajectories, that will turn out to be a much better model,
 like a much more performant model for getting you good reward at the end
 than if you just train this trajectory model here from scratch,
 which is very counter-intuitive because it means that somehow the language modeling task,
 like the language model pre-training, has a beneficial effect
 on the reinforcement learning tasks that comes later.
 To note that the reinforcement learning task has nothing to do with language.
 And even more special, they also try a bunch of other things.
 Most notably, they try to pre-train the image GPT model,
 and that does not result in good performance.
 So it's not just the fact that you have pre-trained on something,
 and it is really a very special result.
 So we're going to dive into the paper right here.
 The setup is fairly simple,
 and then there is a series of experiments that try to investigate this phenomenon.
 So they say that the offline reinforcement learning, as I said,
 has been seen as a sequence-to-sequence model.
 And I've already pre-annotated some stuff right here.
 Let me know how you like that.
 I thought I'd do it in this way.
 So I have the green, that is the current one,
 and the yellow is from my previous escapades on this paper.
 So they go into offline reinforcement learning,
 and that is being framed as simply supervised learning
 to fit return-augmented trajectories in an offline data set.
 What do they mean?
 They mean the setup of the decision transformer.
 I've made a video on the decision transformer.
 If you want to look at that, you can go after you watch this video.
 So the decision transformer says,
 well, see, you are an agent somehow.
 There is an environment.
 There is some interaction between the agent and the environment.
 And in offline reinforcement learning, we usually have a data set of this.
 So someone else has performed this,
 and they've distilled all the episodes into this data set.
 And their goal is to learn just from the data set.
 We can't actually interact with the environment.
 So in the data set, there are a number of trajectories,
 trajectories of the agent interacting with the environment.
 There's always some sort of a state coming back from the environment
 or an observation, if you will.
 The agent always gives some sort of an action back,
 and then there is a reward and the next state coming from the environment and so on.
 So that is naturally a sequence.
 And the sequence is there is a state, then there is an action,
 then there is a reward and a new state, then there is an action again,
 and then there is a reward and a new state.
 So this is a sequence.
 And since I have a data set of these sequences,
 I might as well throw that into a big transformer to do sequence modeling.
 Now, this has its own problems, which I've all discussed in the decision transformer video.
 For example, if the transformer has a context length of four,
 it cannot conceivably look back further than that,
 which is a classic problem in reinforcement learning,
 how to look back and forward infinite times.
 The decision transformer has the limited context window.
 It has sort of the caveats of language modeling.
 However, we understand language modeling very well,
 and therefore, we are quite able to do that.
 There is one modification that they do.
 What they do is they transform the rewards right here.
 They don't let the model model the rewards, they let it model the rewards to go.
 We're going to see that in just a bit.
 This here is interesting.
 What they say is that we look at whether transformer based pre-trained
 language models are able to be adapted to standard offline reinforcement
 learning tasks that have no relations to language.
 I've already told you that this is going to work out fairly well.
 That's the special message of this paper.
 They show consistent performance gains and significantly faster convergence.
 By faster convergence, they mean that a convergence point,
 like a non-improving the loss anymore,
 is reached after much many fewer steps than if you were to train from scratch,
 which makes sense for pre-training if it's in the same domain.
 But given that the pre-training is a completely different domain than the fine
 tuning, that is still a just a special thing.
 So here is how we're going to frame the problem.
 And if you've watched the decision transformer video,
 this should be familiar to you.
 We model a episode as a sequence in the following manner.
 This is almost as we've seen it, except the rewards right here.
 They are not individual rewards, but they are this thing right here,
 the sum of all the rewards at this end, the next steps,
 which they call the returns to go.
 So this, for example, says from here until the end of the episode,
 I'm going to gather 50 reward.
 Now, maybe you're in this state and you made an action that gave you a reward of one.
 So then this here would be 49.
 So you'd say, well, from here on out, I'm going to make 49 reward and so on.
 So the benefit of this is that at inference time,
 you can just put like a really high reward right here.
 So at inference time, you would always you would model these things you would get
 from the environment.
 So you'd start out with like just a big reward right here,
 just whatever the maximum you've observed plus 10% or something
 to just encourage your model to go very high.
 And you plug the state in here that the environment has given you
 and you let the model produce this one.
 So it's important that at training time, we do sequence modeling,
 really model the sequence of returns and state and action as a GPT,
 like next token prediction.
 However, at inference time, we obviously only predict the action
 and the environment is going to give us these two things,
 or the environment is going to give us the reward.
 And then we simply subtract the reward from the previous returns to go.
 And we plug that in here.
 And then we plug in the state we got from the environment.
 We let the model predict the next action right here and so on.
 So this is very cool because much like something like upside down
 reinforcement learning, this is conditioned on it like a desired reward.
 This also has advantages and disadvantages,
 but the advantage is we can control the reward we want at inference time.
 So we don't always have to go for a high, super high reward, but we can.
 Yeah, so this is the setup.
 You don't actually need to understand much more.
 But what we're going to do is we're going to model this as a sequence
 in our data set, and then at inference time, we just put some high returns to go.
 And that's it.
 We're going to use a transformer for that, for the sequence model.
 And they're going to use a bunch of different models right here.
 For example, GPT-2 small, which is a pre-trained model.
 They also pre-train their own that they call chibi-t, which is the same size.
 So that is the same parameter count as the original decision transformer
 to make it comparable to them.
 So the decision transformer is the one that introduced this transformer
 as sequence model for reinforcement learning.
 And they are going to see this chibi-t model has the exact same amount
 of parameters as the decision transformer, so they can directly
 compare what the language pre-training is going to gain them in the same model.
 They also use CLIP.
 However, they only, as far as I am aware, they only use the text encoder part of CLIP
 because that's an autoregressive model, which can do the sequence modeling.
 And they use image GPT, which is an autoregressive model
 that goes via image tokens.
 So an image GPT, it would split up the image into, no, not pixels, but chunks,
 I believe, either chunks or pixels.
 I don't even remember.
 And it would do the sequence model, essentially go through the image
 like this, and then like this, and then like this.
 So it framed the image as a sequence of either patches or pixels
 and go through it as a sequence model.
 So that's a sequence model too.
 We can pre-train it, and then we can apply it to this space.
 They do various things right here, other than just language modeling,
 sorry, other than just language or sequence prediction.
 Let's call that sequence prediction right here.
 Other than just sequence prediction for the reinforcement learning data,
 they do two more things.
 First of all, they want to align the input representations.
 So they have a set of language embeddings, which comes from the pre-training data set.
 Now, obviously, the pre-training data set has a tokenizer.
 That tokenizer generates tokens from the text, and every one of these tokens
 will have one of these embeddings associated with it.
 So V is the vocabulary size.
 However, obviously, in the reinforcement learning settings there,
 we don't have the same tokens.
 We don't have the same input modality even.
 And therefore, we don't need a tokenizer because it's already tokenized, right?
 Each of these things right here is a token.
 However, what we do need is now a new vocabulary, not a new vocabulary,
 but a new embedding matrix, so to say.
 So we have a different amount of tokens, so from one to the 3N tokens.
 And what we're going to want to do is, what they say at least,
 we want to have a set of linear projections that will map the return
 embeddings, the action embeddings, and the state embeddings
 to be very close in their cosine similarity to some embedding vector
 in the original setting.
 So that means they want to force, not force,
 they want to encourage the model to sort of reuse the embeddings
 that it used during the language model training.
 So for each of the input embeddings, they're
 going to find the maximum, the closest, nearest neighbor in cosine space
 of the embeddings of the original vocabulary.
 And then they're going to encourage the input embedding,
 the new input embedding, to be closer to that.
 So that is just a loss that they add during training.
 So you can see right here, this is the loss for the language
 or the sequence modeling decision transformer objective.
 This is the loss that encourages the embeddings
 to be close to the original language embeddings
 or to one of the original language embeddings.
 And this loss right here is the continuation of language modeling.
 So during training of the sequence prediction
 for reinforcement learning, they additionally also do,
 that's what they call language model co-training,
 continuing to train jointly on language modeling
 and trajectory modeling.
 This allows us to encourage, this allows us to encouraging,
 it probably should be encourage, the model's transformer backbone
 to be able to handle both language and trajectory simultaneously.
 OK, maybe it helps.
 This seems either like an idea that had been had at some point
 or something they had to put in after the fact
 just to make it even a bit better,
 or because maybe it didn't work, though they ablated it
 at some point.
 And it also works without.
 So that's almost it.
 Yeah, they describe a little bit their baselines
 and their setup.
 I was a bit confused here.
 It says it's a batch size of 65,000 tokens, which I don't,
 like, I don't, is that, I don't, batch size is usually not
 in tokens, like the sequence length would be in tokens.
 But in any case, they say for our additional objectives,
 we decay lambda 1 and lambda 2 to reach 0 after 5,000 steps.
 We tuned the initial values of lambda 1 and lambda 2.
 And, you know, these seem, they seem reasonable.
 But the fact that you have to, like,
 decay the additional losses after x many steps and so on,
 it points to a little bit of brittleness in them.
 And I'm not sure always how brittle these things are,
 because reinforcement learning is traditionally
 kind of a very brittle field.
 So the main, the main results we have right here,
 the top one is four games in Atari.
 The bottom one is, I believe, three environments
 in the, in the OpenAI gym that are, oh, sorry,
 the, this is a data set, the D4RL data set.
 All of this is offline reinforcement learning.
 On top, you also have the 1% DQN replay Atari data set.
 So as you can see, in many cases, the,
 both the Chibi-T and the GPT-2, by the way,
 GPT-2 is a lot larger than, so this is a lot larger
 in parameters than the Chibi-T model,
 and therefore also than the decision transformer model.
 So just, just saying that.
 So here, the pre-trained models outperform the other ones
 in quite a few tasks.
 However, there is also Qbert, where they still
 do outperform the decision transformer, as you can see.
 But the, they're, one of the baselines
 is just a lot stronger.
 The other baselines are just useless.
 That's kind of what I mean when I complain about,
 when I complain about reinforcement learning
 is that it is just weird.
 Like a bit of a different environment
 can make a large difference.
 But as you can see, the pre-language pre-trained
 models consistently outperform the decision transformer
 models.
 Also something to note right here,
 this is mean and variance across three seeds.
 So this is variance, I'm going to guess they
 mean standard deviation.
 And that is like a large number.
 So if that's the standard deviation,
 then the differences to the decision transformer,
 they are well, well within that.
 And that means, I mean, it is visible that across experiments,
 we see the same trend, right?
 That gives it credence.
 But also, this just seems extremely noisy.
 And yeah, I'm not going to say I'm
 going to sound like reviewer 2 when I say,
 well, you should make more experiments to estimate
 or to get smaller error bars.
 But it just seems like, I don't know,
 it seems like results that you can't really put a lot of weight
 on because they're very noisy.
 However, a bit, like a little bit less noisy
 are the experiments here on the bottom.
 You can see that the standard deviations here
 are quite a bit smaller than on top.
 That's also three seeds.
 I like how they wrote the number three here
 and the word three right here.
 That is just something that you never
 see until someone points it out.
 You can also see right here that the decision transformer,
 for example, is rather consistently outperformed.
 What's also interesting is that image GPT just sucks.
 Like you can see right here, like it just,
 it doesn't get anywhere on any of these tasks.
 Also clip very often underperforms.
 You can see, for example, here clip underperforms.
 And they do have some hypotheses on that.
 That being said, there are still a lot of times
 where the baselines here are quite a bit better
 or just better than all of these transformer-based models.
 So just pointing that out.
 Yeah.
 They do also analyze, and this I find really interesting,
 the attention pattern between the GPT-2 pre-trained model,
 the image GPT pre-trained model, and what I understand
 is a randomly initialized model that has just been fine-tuned.
 Yeah, randomly initialized model that has just been fine-tuned.
 So there's no pre-training.
 So all of these models are fine-tuned,
 but the random one hasn't been pre-trained.
 Interestingly, if you look at GPT-2,
 you can see these bands right here.
 And the bands are always in the distance of 3.
 So there's always 3 distance.
 Now, 3 should be an interesting number
 if you remember how the sequence is made right here.
 So there is always going to be 1, 2, 3.
 These tokens come in packets of 3, right?
 Their next return would be here.
 The next state would be here.
 The next action would be here.
 So every token in this attention pattern
 is most focused on multiples of 3 behind it
 in order to predict the next token.
 So there's always a lag of attention to multiples of 3,
 which means that essentially, if I
 want to predict the next return, probably the last return
 ends are the most important.
 If I want to predict the next action,
 maybe the last actions are important.
 This might also be a property of the environment.
 This is on Hopper.
 So on these continuous control tasks,
 I guess it's very often the case that I'm just
 going to repeat an action for a while
 if I want to achieve some goal.
 I don't know the frame rate exactly of these things.
 However, that seems to be something
 that is rather maybe viable to do.
 And therefore, looking at the last action
 can give me a lot of clues about the next action.
 Looking at the last state can give me a lot of clues
 about the next state.
 I would wonder how this changes if it's something like, well,
 I don't even know, anywhere where I don't naturally
 repeat my last action often.
 You can see this is the early layer.
 Then in the middle layer, the GPT-2,
 it seems to sort of focus on particular states that
 seem to be important, as you can see right here.
 So this is where the attention comes from.
 This is where it goes to.
 And you can see that it kind of decides
 that particular states are important,
 and it kind of remains at that.
 So it selects a few states that, or a few tokens
 that it chooses to attend particularly to.
 In contrast to that, our image GPT
 seems to have a large recency bias.
 So if you see this right here, there's really
 this band right here, which essentially means
 that every token attends to kind of the few tokens behind it
 in order to predict it.
 Then, well, the question is, is it even
 worth looking at stuff further down?
 Because this model clearly doesn't learn at all.
 So I would consider this and this
 just to be kind of random noise.
 The early layers might be interesting,
 though, because there is kind of a pattern.
 And maybe that is influenced by the pre-training.
 So in image GPT, since you have your image,
 and maybe it's in chunks, maybe it's in pixels,
 but I can imagine that if I want to predict
 a particular chunk, that maybe the last few that I've
 predicted, unless I cross a boundary right here
 and go one line down, the last few that I predicted
 are or might be particularly worth looking at.
 And rather distant chunks might be not worth looking at very
 much, other than in language modeling,
 where I often have to go a little bit more
 across the distance and the exact neighboring words might
 not be as important.
 So that might explain why image GPT has
 this particular recency bias pattern in its attention.
 What's also interesting is that the randomly initialized model,
 look at that.
 This is another interesting pattern.
 And you can see that it's very much the same as in the GPT
 example happens, except much more extreme.
 So you have these rows.
 For example, this row right here,
 you can see there is a hard attention for three back.
 It's really hard attention.
 Then there are rows where you can see right here,
 there is always these two, and then these two,
 and then these two, with particular attention
 on the first one and then also slight attention
 on the second one.
 And it's a special pattern.
 So no, I'm one off, sorry, in the one above.
 So this is the hard three.
 Then the one below is the, I'm going to call it the soft three.
 So there is one strong one and one weak one.
 And then the one even below that,
 there is one semi-strong, one weak, and one really weak.
 So what's happening?
 I'm not exactly, so what I don't know here
 is which of these tokens is returns, which ones is state,
 and which one is action.
 But I'm going to just guess, and I might be totally wrong
 right here, that the very strong bias here, that
 is going to be the returns to go, which would only
 focus on the last returns to go.
 And then after that would be the state tokens.
 So what the state tokens would do is, and you can see this,
 I'm just going to, so let's say this is the returns to go,
 the bright ones.
 And you can see that in the state tokens, there is,
 actually there is one missing here on the diagonal.
 So this diagonal one here is just completely blank,
 which means that it just kind of ignores the token behind it,
 which is the reward, right?
 So what it cares about is the last state,
 and it also cares about the last action maybe.
 I don't know how to interpret that very much otherwise.
 So if I want to predict the next state,
 I'm going to care about the last state,
 and the action after that, maybe that makes sense.
 If I want to predict the next action,
 then I might be able to care about all of the stuff
 beforehand a little bit.
 Again, I don't know if I'm interpreting this correctly.
 However, what I am able to say is
 that there is a very, very structured attention
 right here.
 There is this pattern of three is very prevalent,
 and it is in general very, very structured.
 So this seems to be actually the best kind of attention, right?
 It is very structured in the way it looks at the information.
 It learns exactly, aha, there is a structure to it.
 I'm going to attend to the different parts
 in this different structure.
 However, my hypothesis is, and that is not super duper
 discussed in the paper.
 I mean, it is discussed, but my hypothesis
 is that this bias here, it might be almost too strong.
 It might learn the exact structure of this stuff,
 but it might be too strong, and it might miss information.
 Because it, for example, says, well, I
 don't need to know anything in between here,
 because the most relevant thing for predicting the return
 is the last return, and therefore, I'm not
 even going to look at other stuff.
 Whereas the language model pre-training just kind of
 acts as a regularizer that says, well, you should maybe
 look at all of the stuff, even though you don't find it
 super useful in this particular data.
 Now, one thing that I didn't point out in the video
 that I wanted to point out right now
 is that if you look at GPT-2 at the very left column, what it
 does is it focuses particularly on the returns to go steps.
 It doesn't matter which step it is at.
 It always kind of looks back at the very first token, which
 is the returns to go of the whole episode,
 and among other things, also at the second and the third
 returns to go token.
 And this is important, because the returns to go
 is kind of an indicator of how the episode's going to go along.
 If the returns to go are low, it means
 that entirely different episode paths should be chosen in order
 to achieve that reward.
 Whereas if the returns to go is high,
 then I would have to do different actions
 to get that returns to go.
 So it makes a lot of sense to look at the returns
 to go tokens.
 And rather than, whereas you can see in the right hand
 column, the randomly initialized thing,
 it only really focuses on the returns
 to go in these middle layers whenever it needs
 to predict the next return.
 And so it's much more diffuse, and it doesn't condition
 all of what it does a lot on these returns,
 where it makes total sense to do that.
 Because in one instance, the language modeling
 is just sampling any sort of high likelihood trajectory.
 However, additionally in the GPT-2 case,
 it is almost like conditioning that sampling
 on the most relevant information that distinguishes
 between the different futures.
 I hope that makes sense.
 Why a model that would learn to focus in particular
 on this information would be better
 at sampling appropriate trajectories
 for the current episode.
 All right, back to my comments in the past.
 We know that language models retain large parts
 of their pre-training even during fine tuning.
 So the language modeling thing might just
 be like a very good prior.
 And I wonder if we could build these types of priors
 into the decision transformers if we didn't do
 language model pre-training, but just as sort of like
 a bias or a regularizer or something like this.
 Yeah, you can see that through the random attention
 at the end, you do not get this focus as you get
 with the language model thing that it focuses
 on particularly interesting last states,
 but you'd rather you do get like an attention matrix
 in the last layer that is kind of diffuse
 and sort of similar to the image GPT
 that just doesn't work at all.
 So yeah, that would be my maybe postulation
 that maybe it is possible to achieve the same effect
 by introducing the correct regularizers.
 However, I don't know.
 So they look at a few other things
 which I just quickly wanna go through.
 Because they have pre-trained, they can demonstrate
 that their model converges much more quickly.
 So instead of like three hours, their models
 of the same size needs 43 minutes and their model
 that is a lot larger, I believe GPT-2 is 144 times larger.
 It only uses an hour and 27 minutes.
 So still half of the time than this decision transformer.
 Now, I also wonder whether they have based their code base
 on the decision transformer or whether some
 of this difference is also due to just kind
 of like a better implementation.
 So yeah, that is that.
 They have some analysis right here.
 For example, they say they hypothesize
 that a generative training objective is useful.
 That's how they explain why CLIP might not be as effective
 because CLIP is ultimately a discriminative objective
 or a contrastive objective.
 They also say that there are underlying similarities
 between language modeling and trajectory modeling
 where there is a large difference between image modeling
 and trajectory modeling, which is it's a hypothesis.
 They say, yeah, there is the language modeling
 has a natural sequential nature.
 The versus image modeling is kind
 of a forced autoregressive task.
 I agree with that, but I'm not sure
 if there's really due to like language being
 particularly similar or whether, as I said
 it might just be a good prior.
 This would be an interesting question to investigate.
 And it might ultimately turn out to be the same thing.
 So, you know, interestingly
 the context size doesn't really matter.
 You can see right here, if they increase the context size
 they do get worse actually.
 So yeah, that's worse.
 It's just more noisy, which is special
 which actually means that these models aren't appropriate yet
 or we haven't really figured out
 how to appropriately use them yet, right?
 More information shouldn't necessarily give you
 less of a reward unless I guess
 maybe you have a fixed size data set
 and therefore you have less training data points.
 So maybe that's an effect of that.
 Interestingly, the pre-trained models, they do scale better
 which I guess you might have expected
 if you've been in deep learning the last few years
 but if you just take a decision transformer
 it will overfit after a while if you scale it up.
 So these are millions of parameters.
 You scale it up, it actually gets worse.
 Actually not sure if that's overfitting or just, you know
 it gets too big and then the average reward decreases.
 However, if you pre-train first, then it can handle
 and it will actually increase with more data.
 Interesting would be to see if that at some point
 actually declines again or if that sort of holds up
 if the language model pre-training
 for which there is like infinite data, right?
 In language model pre-training, you can get infinite data
 and therefore it could be that this just kind of
 gets you diminishing returns but not ever come down again.
 Yeah.
 They also experiment with freezing parameters
 and they say that this drastically reduces performance.
 So if they only train, if they only train, what do you say?
 Only action state and return projections being trained.
 So only this alignment of this projection of the,
 the projection of the token embeddings are being trained.
 That doesn't work much, which is also surprising
 because there is a lot of work that kind of shows that
 you don't have to train many parameters
 of these transformer models to effectively transform
 or transfer them from one task to the other.
 They say that this might be, this might be the case
 that this might be due to the task of generative modeling
 being harder as opposed to discriminative classification
 where this was previously applied.
 They have a lot of, yeah, they pose a lot of hypotheses here
 of why things might be and I feel each one of them
 could be its own research paper.
 Yeah, I'm gonna leave it at that for the paper explanation.
 I hope you got a little bit an intuition.
 I still find it very, very special and very cool
 that this even works and I think it's an,
 it's an like a sign of the times of our models
 just becoming the same models for all modalities.
 This would not even have been possible a few years ago
 where every modality would use very different models
 like CNN for images and RNNs for language and so on.
 Although RNNs were used for RL already,
 but given that our models converge and we're getting,
 we're learning so much more,
 this type of research is really cool.
 Yeah, let me know what you think is,
 have we overlooked something right here,
 like something that could easily explain why this works
 and gives good results that just no one kinda sees
 or are there more applications for this?
 Let us know what you think and bye bye.