From scratch, single file implementation of a scaled-down version of MuZero for CartPole/Atari(WIP)
Learn to reimplement MuZero from scratch in a single file. This talk demonstrates core reinforcement learning concepts and shows how to build complex algorithms step-by-step.
Overview
I’m reimplementing reinforcement learning algorithms from scratch, most recently MuZero. I want to understand what comes after LLMs. Models learning purely from experience of interacting with the environment seem interesting, see for example David Silver’s recent start-up announcement. I want to show that it is possible to reimplement core parts of these algorithms in a single file.
Video
Transcript
Generated 3 months ago
Summary
Generating a talk summary...
View full transcript
I'll stop the time. So so I have a rather long title. So it's called from scratch, single file implementation of a scaled down version of MuZero for Martí file in the title and work in progress. So what this means is there's 1 take home message. Basically, it's that, kind of, the beginning, I was a bit afraid of these famous AI papers, but, it shouldn't be like this.
So it's actually possible to implement the key parts of algorithms in, basically, a single file, couple of 100 lines of code. I actually have 2 files for this version, and, it's not the full scope, obviously. So but kind of I got some of ends of life for a very simple environment, which is CartPole, and I'm working on the Atari part. So we'll give a bit of background. So we have a couple of talks on reinforcement planning.
But just to say Aswin, you are in this setup where you observe a certain state, which may be the board or then take an action, which may be turning the knobs. You get a reward, which may be potentially be minus 1 when you lose the ball, and you are in a new stage when you, execute the action. So here, the CartPole environment is this, poll that's balanced on a card that you can move left and right. And if it tilts over too much, you get, no rewards. And for every time step that it stays upright, you get 1 point.
And you want to maximize the total expected rewards. Then this guy here is basically so the inventor of reinforcement learning, which Sutton, and, he also coined the term of a Scheerer lesson. So this means that every time you, ingest human knowledge into an AI system, it will be better now. But in the long run, methods that scale with computation, will win. And the question for me is, why I'm also motivated about about enforcement planning is that, it's a question for me whether this holds for LLM free training.
The news, like, a screenshot of the paper that I'm reimplementing and, basically, some news that's also interesting is that kind of the main deep mind person, working on reinforcement planning, David Silva, he left, and he raised 1000000000 dollars for his own startup doing reinforcement learning. This now I leave the presentation so you can find this paper that I'm talking about on archive. It looks like this, basically, and I want to speak about the part here. So it's also Müller based reinforcement learning. So you learn a model of the environment that you used to plan.
So you are in a certain state, which may be represented by the Robert here. You roll out your model, and then you take your best Adrian according to your model, and you train it based on experience that you have, and then map the model back to the experience. Then I want to show some code. So I have, some GitHub, how this looks like for so there's 2 parts. There's the Monte Carlo research.
This is the planning algorithm, and then there's the actual training code. And this is what it can look like like if you roll out the policies. So for example, Martí the roots, which is the current state. You go out in the first Müller. You add 1 node, take action 2, in this case.
In the second simulation, you take a different action and so on. Then you stop, and you take, the most promising action. Then in terms of code, this is the tree search. I don't wanna show it too Müller. But there is something called MCTX, semantic college research Aswin JAX, which is released also by DeepMind, and I debugged against that.
This is actually the codes, I won't show. So it's it's maybe 500 lines, has no dependency except PyTorch. In some sense, it's not very I mean, so, like, there are elements that are repeated across different approaches. So you collect some data, and then you have a training loop. What's maybe a bit specialist that they treat, not values as scalars, but as, on a discrete grid, basically.
So they, 10.5, they would see as 5 times as 0.5 times 10 and 0.5 times 11. And you have some rollouts. Here you see basically, you have the tree search, and then you take Adrian from that. And here you have the environment. In training, I want to show, like, an earlier version, which is maybe more spagety, but does the same thing.
So you have so you can roll out in 2 dimensions. So 1 is at a certain state, how much do you look in the future, and the other 1, how many steps do you wanna train your Müller? And you need to be very careful about the environment when it ends. And so this is what a typical training 1 looks like. It's not perfect.
The perfect score would be 500, and, it's not exactly stable, but I get some signs of life. Thank you. Any questions from the audience? All right. So let's keep rolling with the demos, and welcome Dani to the stage.