World Models, Explained

Why do even our best AI models need tens of thousands of examples to learn skills that a human picks up in a handful of tries? Solving this problem is one of the great open challenges in modern AI. World models, which give AI an internal simulation of its environment, are one of the most promising paths forward. In this episode of Decoded, YC's Ankit Gupta and Francois Chaubard discuss the intuition and math behind world models, new research, and current applications in self-driving, robotics, and more. Full Transcript: https://ycrootaccess.substack.com/p/world-models-an-intuitive-introduction

• 0:00Speaker 00

  One of the biggest open problems in AI right now is how to solve sample efficiency. That is, how do you get models to quickly learn new tasks or skills from relatively small amounts of training data?

• 0:08Speaker 01

  Humans do this incredibly well. We can learn new games, concepts, and skills, often after just a handful of tries. Our best models, on the other hand, often need tens of thousands of data points just to learn.

• 0:18Speaker 00

  So today, we're going to discuss what many top researchers believe is the most promising path to closing that gap: world models.

• 0:24Speaker 01

  We're going to discuss the motivation and math behind world models, current applications, and why this approach might be the key to unlocking AGI.

• 0:39Speaker 00

  You and I have talked a lot about the various ways people are training models and the sample efficiency of them. Why don't we start by just defining sample efficiency and how we intuitively think about it as humans?

• 0:49Speaker 01

  Yeah, so I think from my perspective, the two major problems that we haven't left to solve is intelligence per watt and intelligence per sample. Intelligence per watt is like how many valve perplexity points we get per watt of spend. And then intelligence per sample is basically if I have one additional sample in my dataset, how much more intelligent am I getting? And so if I imagine I have a new tasks, like RAGI, for example, I think like really Francois Challet has been on the forefront of this thinking and talking about intelligence as a rate of skill acquisition versus skill acquisition. And that's very different. And so how fast do we get smarter with more and more samples? And these things are incredibly poor at getting smarter with fewer and fewer samples.

• 1:33Speaker 00

  And for context, you know, the RAGI test sets are a really good example of cases where humans are intuitively very good at them. Most humans can intuitively solve those puzzles with some amount of thinking and effort. But our current state-of-the-art AI systems, what people consider frontier intelligence, basically can't do them.

• 1:51Speaker 01

  Right. I mean, we come into new problems with such inductive bias from K through 12, like all these math in school that we've had, that, you know, these models are kind of getting from compressing the entire internet. And so when we come in, we're not coming in tabula rasa, just like bare bones. But even so, that they have, you know, I don't know what percent of the internet you've read. I've read very little percent of the internet. But despite that, and having read the entire internet, it still can't really do well in generalizing to these new tasks.

• 2:23Speaker 00

  So now let's think about this in like the extreme cases. In the extreme case where, let's say, we were perfectly sample efficient. You know, we were as sample efficient as possible. What would that mean in terms of a model that is taking a set of actions in the world?

• 2:37Speaker 01

  Well, I guess the perfect sample efficiency would be zero samples. And like, there are examples of this. And that sounds absurd to say, but and the example, the hypothetical I'll give on this is imagine I had a perfect world model. Then I should never go to the environment to go and collect samples to train on. And well, that can't possibly happen. First of all, no, it actually can happen. We do it all the time. It's called Newton's second law of motion. It's like Newton mechanics, like we basically know how to like get an object from point A to point B with a rocket quite easily, just by following like Newton's laws of motion.

• 3:14Speaker 00

  Yeah, like when NASA plans to intercept an asteroid and is planning it, you know, years in advance and can set it off in a trajectory where it just glides to the right thing and intersects to the right point. That is an example of a perfect world model we've built, where we're then just letting that world model act. And that that system does not need to intelligently collect new samples from the environment to decide which direction to go next. It can already, it's already been pre-programmed and it can perfectly do it.

• 3:39Speaker 01

  Yeah. Can you imagine if like we needed to collect one million training examples of like us shooting spaceships to the moon to like know how to do it? It was like this completely, it would be, we definitely wouldn't have the Apollo missions, right? But we do have that that ability because the real world is differentiable and we can do something called model predictive control that we're going to talk about in a little bit. But even in our own brain, I was just, you know, thinking about this on the drive up. But like, there's so many ways that like I can basically think about the things that you are going to say, or what a VC is going to say when I'm, when I was pitching them, or what a customer might say. And even product being having taste, what is taste is like predicting that other people are going to like this thing.

• 4:21Speaker 01

  And so we've built this world model over years of entrepreneurship, 10 years of like getting it wrong, right? That maybe Bill Gates, Steve Jobs, and Jensen have 50 years of, you know, world modeling experience to know what people want. And, and basically, this is actually proven in the 1967 CogSci study by Richardson. They basically showed that if you take a cohort of three different people, three groups of people, and you have one go practice layups in basketball, and they go and they shoot, they improve, for one hour, they improve by like, I think it was like 24% or something like that. And then if you take the other one and they just blindfold them and they imagine laying

• 5:05Speaker 01

  up a basketball, they improve at 23%.

• 5:07Speaker 00

  Interesting.

• 5:08Speaker 01

  Against the control. I mean, that's insane. It means that we have this crazy good world model. And there's this neuroscientist at Stanford named Sean Druckmann, who basically is of the view that the entire point of the growing neocortex during the Great Cortical Expansion 10 million years ago was to get better and better and better and better world modeling. And having just like my little VLA, which we'll define, of doing the predicting next action is not as good as having a world model to lean on, either for training, for training purposes, or for test time adaptation.

• 5:39Speaker 00

  Yeah. What it fundamentally comes down to is, you know, we as humans, we think about our intuitive ability to think as coming from some implicit world model we have in our heads, encoded by genetics and our ability to learn and whatever else. It seems like models can do surprisingly intelligent things despite not having an explicit world model when it comes to natural language. When they're just talking, it seems like, you know, maybe under the hood, deep inside the weight somewhere, there's some kind of implicit understanding of the world. But there isn't an explicit representation of that. But it seems like in certain domains, especially in robotics and self-driving, as we'll talk about, that sort of breaks down. And, you know, maybe it would be helpful now to just think a little bit about and just sort of define some of the pieces of what makes it challenging in these different

• 6:24Speaker 00

  domains. And then we can use that to kind of build up to why it's particularly hard in things like self-driving and robotics to get these types of predictive models to work.

• 6:33Speaker 01

  Yeah, let's do it. So let's actually like take a step back and just talk about like control, reinforcement learning, and define some, define some common terms. So typically, in, we teach a course called Decision Making Under Uncertainty, which is like the main reinforcement learning course at Stanford. I like to show like a specific example of, let's say, I have some drone, and this is my poor little drone here, and it has some mass M. And we know that that gravity G is pulling down on it. And it's currently at position T with velocity T, which we will collectively call the state.

• 7:10Speaker 01

  And to be really clear, this is going to be P, X, P, Y, P, Z, T, T, T, and V, X, V, Y, Z, V, Z.

• 7:24Speaker 00

  It's like a six-dimensional state vector.

• 7:26Speaker 01

  Yep. And we have some thrust vector U that we control. And we're trying to get to some point P star and V star, which is, V star is typically zero. And so you have some platform that I want this thing, this drone to land on.

• 7:43Speaker 00

  Yep.

• 7:43Speaker 01

  This is this control problem, right? And so, let's say this is like, and we'll go through optical or optimal.

• 7:52Speaker 00

  Optimal, yeah.

• 7:52Speaker 01

  Optimal control. So how would I actually solve this? So the first thing I need to know is my transition function. And so this is my state transition function, which is ST plus one.

• 8:04Speaker 00

  Given the previous.

• 8:05Speaker 01

  Given ST, and my action, which I control, is UT. And so this is my state transition or dynamics function or a world model. This is a world model.

• 8:15Speaker 00

  This is like a very fundamental for, for context, you know, this equivalent to transition function you would think about in RL in general.

• 8:21Speaker 01

  Exactly. And so, and then what I'm trying to learn is something called a policy, which is like what UT should I emit, given?

• 8:34Speaker 00

  Some ST.

• 8:35Speaker 01

  Yep. And so this is the ultimate question. What should I do? What's, what action should I take, given some state ST? And so the way that we'll solve this, luckily, we have a world model that is perfect. It is called.

• 8:49Speaker 00

  Newtonian physics.

• 8:50Speaker 01

  Newtonian physics. This is like Newton's second law of motion, which is F equals ma. And so we know that the position PT plus one is going to equal PT plus delta T VT plus one-half delta T squared. So everyone's taking high school, high school physics.

• 9:11Speaker 00

  Yep.

• 9:12Speaker 01

  And the same thing for the velocity, blah, blah, blah.

• 9:15Speaker 00

  Delta TA.

• 9:16Speaker 01

  And then my acceleration is the sum of, some of the, some of the forces, which is going to be my UT, I think I divide by the mass, and G. And so that's it. And now I have my transition function. Now, how do I get to a policy? And I'm going to apply something called model predictive control or real-time model predictive control, which is like the way that SpaceX lands the rocket on, on some platform in the ocean. And what you're going to do is you're going to set up your loss function. You're going to minimize sum over all T. You have UT to infinity. And I'm going to minimize my P star minus PT.

• 9:57Speaker 01

  Plus v star minus vt.

• 10:03Speaker 01

  And usually you add this little lambda ut, which is like how much energy you're exerting. And you can't have infinite thrust. So you typically will have to say ut, you max thrust.

• 10:20Speaker 00

  Yeah.

• 10:21Speaker 01

  That can be achieved. And so this is easily solvable with comics optimization. And so this is convex. This is convex. This is convex. The sum of convex functions is convex. This is a convex constraint. And so I DCP disciplined convex programming means that I can put this into a CVX pi and it will just give me out my policy, which will be the solution will be the optimal

• 10:47Speaker 01

  ut plus one all the way to infinity.

• 10:50Speaker 00

  So we can solve this in closed form, basically. Like, you know, we can, because we have this world model of Newtonian physics, we can say at every step exactly how this drone should fly so that it lands on the appropriate.

• 11:01Speaker 01

  Exactly.

• 11:02Speaker 00

  Under a set of constraints like max thrust.

• 11:04Speaker 01

  Exactly. You'll, you'll run your log barrier or interior point, whatever, to some solver on this and it will give me my optimal. And this will be literally the optimal path that this thing can take to get to this state. And that will minimize. And then, and I can, I can do, increase this if I, if I want it to do the least energy path or if, and I make that zero if I want it to be the fastest. And so that's typically the way that you would do what I would call like deterministic

• 11:36Speaker 01

  differentiable control.

• 11:38Speaker 00

  Right.

• 11:39Speaker 01

  And why differentiable? Because I can take the, I can form the Lagrangian by, by taking this minus this constraint and, and take the gradient of it. And I can do Monrobin's.

• 11:50Speaker 00

  You use the fact that it's differentiable to, to do the optimization.

• 11:54Speaker 01

  Exactly. If this is, if this is non-differentiable, you cannot do comics optimization and you cannot do SGD, even if it's non-convex, you could still solve and get, and get a pretty good solution as we do in deep learning. But I, I, if it's non-differentiable, you kind of can't.

• 12:07Speaker 00

  There's nothing you can do. So yeah, let's have an example then of how you could make this non-differentiable. Like, well, what's a, what's a scenario? I guess even if like this drone scenario where it now becomes non-differentiable.

• 12:16Speaker 01

  Yeah. So I'll put this adversary named Anket.

• 12:19Speaker 00

  Okay.

• 12:20Speaker 01

  And, and your job is to, you have another drone, let's say. Anket's drone is to try to hit me and stop me from getting there.

• 12:28Speaker 00

  Now, from the position of your drone, you don't know what actions I'm going to take.

• 12:32Speaker 01

  Right. And so now, let's just call this the, uh, this would be now, we're definitely not, not deterministic. We're stochastic.

• 12:42Speaker 01

  And stochastic and non-differentiable.

• 12:45Speaker 00

  Yeah.

• 12:46Speaker 01

  And in this case, my state transition, what is st plus one? It's going to be my state I'm in now, my thrust and what Anket's going to do.

• 12:59Speaker 00

  Right. And these, it was all differentiable until this new variable.

• 13:03Speaker 01

  Yeah. And I can't like backprop through your brain. To tell it, say what you're going to do with your little drone controller. Right? It's completely, uh, non-differentiable now. And I am resorting and I have to resort to this awful area called reinforcement learning, which is just super brutal and it's sprawling and there's so many different things. And you'll hear things like when you study initial, uh, reinforcement learning called value iteration or policy iteration.

• 13:32Speaker 01

  And there's DQN or deep Q learning or just Q learning. There's actor critic. There's all this bag of stuff.

• 13:40Speaker 00

  All of this stuff ultimately comes down to ways to estimate, to, to model this non-differentiable stochastic process.

• 13:49Speaker 01

  Exactly. Yeah. And so like, that's basically the, the main thing is, is you're going to start talking about, uh, this as a model where I'm going to introduce this psi to say that this is going to be some model that's going to take in these things and then output this. And that we're going to train it over many, many instances, instantiations of this. And that's to get a better and better world model. And then I need to train some policy, ATST. And then typically you also need a value function. And that is the value of some state. And to discern between the value of, of different states. And like, in this case, I don't know what a valid state is, but like, let's just say I was doing, um, uh, like the SpaceX with, um, launching rockets and landing rockets

• 14:34Speaker 01

  in Florida. Let's just say that like, there's different, if I have my launch pad here and I have a whole bunch of houses here, let's just say the path going from here to here. I may think that doing this and then coming across here and burning all these houses alive.

• 14:53Speaker 00

  Right.

• 14:54Speaker 01

  Maybe not highly valued. So I might say, as an example, they typically call this like some type of cone here. And I might say like, it's low value to be here or, and it's very high value to be, to be in this cone or something, right?

• 15:07Speaker 00

  In a sense of value gives you some expectation of future rewards, like the sum of future rewards you're getting. And so if you, if you're in a bad space, you would set the value to zero or negative, negative infinity or something.

• 15:17Speaker 01

  Yeah. So, so we can, we can, we should introduce RT as well. And so typically like if you're playing Go or chess, like winning the game, uh, you can say winning the game is plus one, minus one for losing, draw zero. That's what's done in AlphaGo. In chess, we have these heuristics, like a, a pawn is worth one point, a rook is worth five, et cetera, et cetera. So you can like already have reward is the difference in, in, in board, board state. And then this, yes, will be the sum of my discount.

• 15:51Speaker 01

  Should just do T of RT.

• 15:54Speaker 00

  Yeah.

• 15:55Speaker 01

  Given, and it, and it's important also to, to, to use this nomenclature, um, v pi. And the reason why that's important is because what it, what's actually happening here is this is the discounted reward following policy pi.

• 16:10Speaker 00

  Correct.

• 16:11Speaker 01

  And that means that when I'm in this state, I will take this action and then I'll end up in this to st plus one, and then I'll take this action and it's, it, it, taking it greedy. And so that's the value with respect to pi.

• 16:20Speaker 00

  Yeah. And so ultimately what it comes down to is we are trying to still find a new policy pi and along the way we will use the learning models in various capacities. This is standard RL to estimate the value function, given the rewards we're receiving.

• 16:35Speaker 01

  Right.

• 16:36Speaker 00

  And then where world models come in is a way of incorporating all of those into some sort of joint modeling of the state and action distribution so that we can make more intelligent policies off of it.

• 16:47Speaker 01

  Right. And so your standard kind of setup for this is what I'm always trying to get to at the end of the day is some joint distribution, which would be st plus one, given, uh, where I'm at now, uh, where I'm at now. And then this factorizes with chain rule simply to my pi, my policy, AT given st.

• 17:11Speaker 00

  Yep.

• 17:11Speaker 01

  And my world model. I'll give this, this is usually represented with theta. And this is my, uh, world model, which would be, um, st plus one, given, uh, st and AT.

• 17:24Speaker 00

  Yeah.

• 17:25Speaker 01

  And so, and these are typically learned, uh, separately. And like, and like, you can imagine, in fact, actually you can actually learn this. This is a video generation model and I have the frame st and I predict the next frame st plus one.

• 17:39Speaker 00

  Right.

• 17:40Speaker 01

  And then, and we'll get into this.

• 17:41Speaker 00

  Yeah. For those of us who kind of saw our, our diffusion model series, often people these days use video diffusion for exactly this.

• 17:47Speaker 01

  Yeah. And then what you can do, and this is like the in vogue thing to do since Danajar and, and, um, the Dreamer paper series from V1 to V4 is do action conditioning later, like similar to clip where we will inject this like input head, uh, input tail to come into the model to, uh, influence and, and, and enable the world model to have embodiment. What does that mean? It means that not only can I predict like as a plant or a tree on the, on, growing on the, the side of the, of the building, I can like see the world going, passing by, but I can, I can actually influence it. And I can change the, the world and I can, I can learn that with AT. And it's far fewer samples, uh, to do this post-action conditioning.

• 18:30Speaker 01

  If I already have a really good, uh, st to st plus one world model.

• 18:35Speaker 00

  And so here you're saying, you know, what's also in vogue now is jointly training these versus separately training them.

• 18:41Speaker 01

  Exactly. And so this is called, that is called the world action model where the, some of the issues here is one, there's all these training dynamics. If these things are just disparate, trained on different sets and things like that. The other issue is plainly obvious, what I have to do to actually do test time planning is I'll have to sample my, with model one, invoke theta and then pass that sampled action into here. And then roll it out to S T plus one and it's very expensive and it's a very not real time. Two major issues and why, like, why can't we just scale up AlphaGo to like solve all the problems? Is because it's because of this property. If I have one invocation to the model and it gives me both.

• 19:25Speaker 01

  Here's the action I should take and here's the S T plus one that'll end up much, much cheaper and much, much faster.

• 19:30Speaker 00

  Okay. So I think that's a really good segue. I think, why don't we now motivate everything we just described? Through a series of increasingly complex environments. So I'll contend that I think the right set of environments for us to consider is chess followed by Go followed by self-driving followed by robotics.

• 19:49Speaker 01

  All right. So let's go through a couple of examples of problems that we want to apply reinforce and learn to. So chess is a pretty easy one. There's an eight by eight grid.

• 19:59Speaker 01

  And so typically when you, when you approach any RL problem, you're going to look at star. And so this, the size of the state, uh, the number of states I can be in. So if I have these eight here and these eight, so this is eight, 16, 32. So it'd be 32 to the 64.

• 20:20Speaker 00

  Yes, quite large.

• 20:22Speaker 01

  Quite large. Then my transition function is stochastic and non-differentiable because you can.

• 20:28Speaker 00

  You don't know what the other players are gonna do.

• 20:30Speaker 01

  So if I'm like in playing chess.com at my house, I move and then something happens and it comes back and then now you moved and the board has changed. So I can't really differentiate through what the other player is doing. The cardinal action space is actually quite small. Even though there's 32 pieces and all that stuff, there's only eight possible moves in expectation that you can actually, that are legit moves. So like.

• 20:55Speaker 00

  In any given state, there's only eight-ish moves you could do.

• 20:59Speaker 01

  Let's just say in the beginning, I can move all my pawns. I can move my horses. So that's 10.

• 21:03Speaker 00

  Yeah.

• 21:03Speaker 01

  That's like, not that much. So this is extremely small. And then my reward, we can use the heuristic based approach, or we can just say, you know, plus one, zero or minus one if I lose, plus one if I win. And so this is very tractable.

• 21:17Speaker 00

  You say it's tractable, even though there's a really big state space here.

• 21:21Speaker 01

  Yeah.

• 21:21Speaker 00

  But why don't we talk about that for just a second? I think this is a really important point. And when you say it's tractable, you're specifically referring to the action space being small because it affects the kind of like combinatorial expansion here. Should we talk about that for just a second?

• 21:34Speaker 01

  Yeah.

• 21:34Speaker 00

  Or maybe we can add Go and then kind of contrast the two.

• 21:36Speaker 01

  Yeah. So why don't we do that? Because it's, because I want to get to the AlphaGo, uh, the way that they solve this. And you're right. So if I were to do this naively and I just took my S T plus one and I want to do look aheads. What I would do is I would take all of the actions I can take. So there's eight. So I would do action one, action two, action eight.

• 22:00Speaker 01

  And then each one of these, I need to expand it for all possible states. And so now I need to do cardinality S, which we just said is this huge frigging number. And so I have to do that eight times and I have to do it again. I have to do it again. So just doing, looking forward, one move is like quite intractable.

• 22:18Speaker 00

  Although at the same time, you know, the, everyone starts at the same starting position. And while it is a really large space, you know, it, there isn't an infinity number of potential. There's actually a really, really small number of game boards, even four moves into the game.

• 22:35Speaker 01

  Right.

• 22:35Speaker 00

  As opposed to a game where you could start any permutation, for example, of initial game state and what a few states down.

• 22:42Speaker 01

  Yeah. So this is like definitely over, uh, done because there's, there's, it's, it's much, much less than this in practice.

• 22:49Speaker 00

  Yes.

• 22:50Speaker 01

  But just naively, like looking at, you know, uh, what possible game states could be, uh, as a rough math here. But this is roughly the idea. And then each one of these leaves, I need to invoke my value function.

• 23:02Speaker 00

  Right.

• 23:03Speaker 01

  Which is the value of that state T plus one. And so I have to do that all many times and we'll get this while we go, but like this ends up being estimating the leaf node, uh, because at the end of the day, my policy, A T S T. I want to pick on the arg max of like the value of the, the following.

• 23:25Speaker 00

  The argmax action, I guess it'd be an A here.

• 23:27Speaker 01

  A, exactly. Yeah. The argmax over A of the value of the state of the, of the end state, S T plus N, let's say. It's like, that's the, the main goal here. Um, and so for me to do that, I need to like roll all this out, estimate the value and then pick the, the best one. And so this, this quickly grows. Um, however, and we'll see this with AlphaGo, which is actually, actually has an even bigger state space. So I think it's 19 by 19.

• 23:57Speaker 00

  I think it's about right now.

• 23:58Speaker 01

  So you have this 19 by 19 grid. You can, in each one, it can be black, white or, or nothing there. So I have three, uh, so let's do our star again. So the cardinality of the state, I think, is going to be S, uh, two or three, my ternary thing here.

• 24:17Speaker 00

  19 squared, I guess.

• 24:18Speaker 01

  The 19 squared. I think it's 361. Yeah, something like that. 361. My transition, same issue. I don't know. Uh, my action space is going to be 361, let's say.

• 24:31Speaker 00

  So it's a good amount bigger than chess.

• 24:32Speaker 01

  Much bigger.

• 24:33Speaker 00

  But it's still not, uh, enormous.

• 24:36Speaker 01

  Yeah.

• 24:37Speaker 00

  As we'll see in a second.

• 24:38Speaker 01

  Yeah. And so basically what they do, they call this Z, which is kind of annoying, but let's call it R. And it's the terminal, it's the terminal when they won the game. And they basically, you know, you have your trajectory, which is S zero, A zero, R zero.

• 24:59Speaker 01

  Then all the way to the end of the game.

• 25:01Speaker 00

  Yep.

• 25:02Speaker 01

  S N A N R N. And if you won, then all of these, uh, all the moves that black, if black won, all the moves that black did get plus. All the moves that white did were minus one. And they just, that's how they create their, um, their rollouts.

• 25:19Speaker 00

  Rollout refers to a taking N steps of play of all players, one after another.

• 25:27Speaker 01

  Yeah.

• 25:27Speaker 00

  Of moves under a specific policy at the, at the particular instantiation of it.

• 25:32Speaker 01

  Right. So let's just, let's probably under this policy, P theta T. And we're going to overload T, but like this is that instantiation. We froze that model. We froze that model and we play, I think it's like 70 games. And we like treat all of those and we said, we're going to sub-sample a bunch of, um, of these state action results, state action results to train our, to update our policy in our, um, in our world model, our transition model. And what it's actually doing is we, we take in an S T, we give it to some theta, then it wants to output, um,

• 26:07Speaker 01

  the probability of S T plus one being played, uh, which is our transition function and, uh, the, uh, value of the current state.

• 26:19Speaker 00

  And how do we get the value?

• 26:21Speaker 01

  And so the value of the current state, uh, well, both of them are coming out of, out of the model, but basically the loss function, L. Theta is going to equal, and it's going to be eerily close to this, uh, control problem one, is we have some V theta minus this Z, which we'll just call it R here, um, squared. And then plus, uh, actually, sorry, it's minus this pi, which I'll explain in a second, log P theta.

• 26:55Speaker 01

  And I think they, everyone includes this, but they include it in the paper. So I'll include it there as well, which is the, um, weight decay.

• 27:03Speaker 00

  Yep.

• 27:03Speaker 01

  And so, um, so this is basically what, uh, our loss function is. Then we'll play a bunch of these games. And let's try to be a little bit organized here. And, uh, and so this is our setup. This is our architecture. And now the most, once we train this thing, we do an insane, insanely expensive task of, uh, of test time planning. And so this trend in RL is just called test time planning.

• 27:33Speaker 00

  And the specific algorithm they use here for this is Monte Carlo research.

• 27:37Speaker 01

  MCTS. And so this is one of the possible things that you could do. Uh, it ends up working extremely well if you have small action spaces.

• 27:45Speaker 00

  Yeah. So let's, let's just like very intuitively talk about what MCTS does. And a lot of people have heard about Monte Carlo research because AlphaGo was such a, you know, big moment.

• 27:54Speaker 01

  Yeah.

• 27:54Speaker 00

  How exactly does that math into our star and value function and policy?

• 27:58Speaker 01

  Yep. So I'll take this S T. This will give me, uh, 361, uh, uh, numbers that sum to one. And so I'll have some probability of, uh, of where these things are going to go for the, of where my, my opponent will play, um, here.

• 28:15Speaker 00

  So these are like the sets of actions.

• 28:17Speaker 01

  Yeah. So I'm here.

• 28:20Speaker 01

  So that I have all my S T plus ones, I have 361 of these things. Um, and then.

• 28:25Speaker 00

  To be clear, this is like action one, action two, all the way to action 361.

• 28:30Speaker 01

  Exactly. Yeah. And the, um, we have to estimate the value of each one of these. And so then we have to invoke the model all 361 times to give me values for each one of these things. And then I will select. I'll select it based on the, the U C B, the upper confidence bound, which is this equation that is roughly something like, um, balancing, uh, my value function of S T plus one, which they're going to, in, in the literature, it'll be called the Q value because it's actually the difference between a value function and a Q value is just that I have the action as well.

• 29:11Speaker 00

  Yep.

• 29:12Speaker 01

  So it's be S T then A T. Um, so we'll just call it Q value, which is my, um, exploitation term. And then my exploration term will be something like, uh, it's this funky square root of N.

• 29:32Speaker 01

  So it's the argmax of A of my Q. And then I have this, which is the probability of this, this move being played, which we have from here of, of S, let's just call it S T plus one. And then I have this term, which is this sum over, uh, N S B divided by N S A. And.

• 29:58Speaker 00

  Yeah. What's, what's the intuition behind this term?

• 30:01Speaker 01

  So these ends is, is the, the visit count during my MCTS process. So this whole tree, I'm going to.

• 30:10Speaker 00

  So this tree can get really big, right? It's 361 per thing.

• 30:13Speaker 01

  And it's depth of 30.

• 30:15Speaker 00

  So you can't visit every single leaf node.

• 30:17Speaker 01

  Exactly. And so you want to keep track of which uh which state did you end up in? And what action did you take when you were in that state? And you want to make sure that you you have good exploration, right? And so the way you keep track, the way you ensure that you have good exploration is you want to not just be greedy and always pick the highest value one, because that could be very myopic. And so what you'll do is during this MCT process, you'll start this dictionary, which will be all zeros of the visit count of being in this state and taking this action.

• 30:52Speaker 00

  Yep.

• 30:52Speaker 01

  And then, once you go through your first rollout, you'll do, you'll go here, you'll, all these things will be added to zero. You'll have some probability, we're going to bias it towards the higher probability of places to go. And then we'll go, we'll expand those trees. And then we will update the counts that we visited this. And that will basically reduce the amount of probability that we're going to select it again, because this this will reduce my my exploration term. And if it's highly valued, then we're going to increase the queue on this, because this is the expected value of going down this this this path.

• 31:29Speaker 00

  So the gist of it is fundamentally like you want to take the optimal-ish path, but have enough exploration in this really expensive step you're doing here, so that you are making sure you're getting a decent chunk of the other potential leaf nodes you could traverse to

• 31:49Speaker 01

  Right.

• 31:49Speaker 00

  in these 30 step rollouts.

• 31:51Speaker 01

  And so I'm going to do this this MCTS simulation 800 times here. And then for all 800, I have to go through this whole process. And I have to invoke the model like at least 30 times to get through all here. And so that's, you know, 27,000.

• 32:08Speaker 00

  800 times 30.

• 32:09Speaker 01

  Yeah. So 24,000 invocations of the model to to develop this tree. And then, once I have it.

• 32:15Speaker 00

  That's per step.

• 32:16Speaker 01

  Per step, just to do one action into the game. A lot of people don't understand that this is like, you don't like store this MCTS tree. You like, you throw it away after you you make the move. But once, it's very expensive to develop this MCTS tree. And once you have it, the probabilities of traversal are actually extremely useful for training. And then you end up biasing it and you train it with the MCTS tree, which is like a little bit seems like circular motion or something like that, like, but you end up treating that as as the pi that you'll train in your loss function. So we have the R of did we win or lose? We have the the pi of of what was the end result of this whole expensive process.

• 33:01Speaker 01

  And then at test time, we are going to do these 24,000 steps every single

• 33:08Speaker 01

  every single move to pick the arg max that gives that that satisfies both exploration and exploration and exploitation.

• 33:18Speaker 00

  In this case, you know, this still feels somewhat tractable, though, because the action space is small enough where this like kind of works.

• 33:26Speaker 01

  Exactly.

• 33:26Speaker 00

  But now, like, let's say hypothetically, maybe we can draw like an imaginary Go, a game of Go, where it's like, you know, let's let's say this game of Go was like 1,000 by 1,000. And so now you have A equals, you know, more or less a million.

• 33:46Speaker 00

  And now this this tree we're drawing here that has to take here, this has cardinal or like, you know, width, I guess, 1 million.

• 33:56Speaker 01

  Right.

• 33:57Speaker 00

  And there's like S zero through S 1 million. And the number of, you know, steps you would have to take here presumably will have to be way more than 800 in order to get any reasonable kind of sampling of this. And so you're probably multiplying the test time cost of doing a rollout or of doing a next step prediction astronomically. If the game was even, let's say, you know, this is only 100x bigger than the current game, not even 50x bigger than the current game.

• 34:27Speaker 01

  Everyone was very excited about AlphaGo and at the time, and what was it, 2017, 2016? Everyone's very excited about this. And the important thing to pick up is that we did 800 MCTS simulations and to cover 361 possible actions on average. So that gives us about two samples roughly on an expectation for every single action.

• 34:51Speaker 00

  So here you need like 2 million of them for a similar depth.

• 34:54Speaker 01

  For for a similar depth. And then that's still to do a depth of 30, I would still have to do this times 30. This would be 60 million invocations of the model. So that better be a small model, right? That's a lot. So, yeah.

• 35:05Speaker 00

  So that's to do a single action, to be clear.

• 35:07Speaker 01

  Yeah, so exactly, to do one action. So just imagine, so why AlphaGo doesn't scale?

• 35:19Speaker 01

  To me, there's one, the cardinality of the action space must be extremely small. If it's big, sad.

• 35:29Speaker 01

  Two, the I need a perfect deterministic environment, right? Like this, this, this doesn't change. The rules of this game don't change, but like the rules of the stock market change all the time. The rules to like venture change all the time, like the real world changes quite often. So like homo schadistic.

• 35:52Speaker 01

  And real time. If you saw the the movie, the documentary, it was such such such an amazing documentary. I'd highly recommend it to anyone that wants to watch it.

• 36:03Speaker 01

  The guy is sitting there for like 60 seconds, maybe five minutes, waiting for the computer to like decide. And and it's kind of like, imagine that we were driving a car and like you took like 60 seconds to like turn the steering wheel, everyone's dead, like the whole car is dead.

• 36:16Speaker 00

  Right.

• 36:16Speaker 01

  And so like, you know, now let's talk about robotics and self-driving car. And why this, why that approach kind of can't scale?

• 36:26Speaker 00

  Yeah, I think the really good contrast here, because intuitively, I think in thinking through this exact star layout, it actually really changed how I think about the kind of problem space of both of these two. So like, and let's take self-driving car as an example. This is one, you know, many people have started to experience for the first time, because we have some self-driving cars that actually work. You have Waymo and Tesla, FSD and whatnot, that seem like they kind of work. So like, let's maybe apply your same star framing here. I would contend that the state space of self-driving car is enormous. And it's actually not intuitive to me whether it's more or less large than this one, right? I mean, in a sense, the chess in AlphaGo state space is already like more than the number of atoms in the universe or something to that effect.

• 37:10Speaker 01

  Right.

• 37:10Speaker 00

  But like, just to emphasize that here, you know, you are considering, you know, surroundings, vehicle state.

• 37:19Speaker 01

  Yep.

• 37:20Speaker 00

  Like, you know, camera, like

• 37:24Speaker 01

  Weather.

• 37:24Speaker 00

  Weather. I guess the point is like

• 37:27Speaker 01

  Road conditions.

• 37:28Speaker 00

  It's like massive. This is massive.

• 37:30Speaker 01

  For all intents and purposes, infinite.

• 37:32Speaker 00

  For all intents and purposes, it is infinite. Correct.

• 37:33Speaker 01

  Yeah. And and so is the space of pixels, like, you know, like, what can I put in an image? I can take an image of anything.

• 37:43Speaker 00

  Yes, true.

• 37:44Speaker 01

  And so we're able to handle it. And the same thing here, where we compress from the board state. We don't represent the board state. We compress it with a ConvNet. And so they have some deep, some some some deep ConvNet that actually takes this state and converts it into a latent.

• 37:59Speaker 00

  Right.

• 37:59Speaker 01

  And that latent compression is sufficient to kind of like do pattern matching, do do some type of like symmetric symmetric equivariance kind of things. And same thing with this. And even better with JPA, which we can talk about at the end there, which is like basically taking some type of state space and doing all of our optimization in the latent space, which stable diffusion did, that worked extremely well, which reduces our state space dramatically, because I'm in some latent high dimensional space.

• 38:27Speaker 00

  So like, the key thing there is that, yeah, despite this state space being effectively infinite, we've actually gotten really good at compressing this.

• 38:35Speaker 01

  Yeah.

• 38:35Speaker 00

  And we'll talk more about some of the tricks for how we actually do this in practice here. But the TLDR is, you know, where there's like 10 years of deep learning work that basically makes us extremely good at compressing that very fast.

• 38:48Speaker 01

  Exactly right. Exactly.

• 38:49Speaker 00

  T seems to have a similar problem as before.

• 38:51Speaker 01

  Right.

• 38:51Speaker 00

  In fact, maybe even more extreme. There's like infinity other variables around you.

• 38:55Speaker 01

  Right. In some ways, you'd think that it's, this is physics. Newton's laws, laws of motion should apply. If I turn the steering wheel like this and I hit the gas, I should be able to really easily model this. But what is non-differentiable is that I have, if I'm going into a circle, right? It's like the most, the biggest issue that that we faced when I was doing self-driving car is like you're imposing your will onto, maybe driving in India.

• 39:22Speaker 01

  You're imposing your will onto the environment. And like, people just kind of adapt naturally. Like, if you were doing Newton's law of motion, you were going to collide. And so that the optimal policy, if you were doing strict Newtonians here, would be like, don't move, because anything you do, you're going to crash.

• 39:36Speaker 00

  Yes.

• 39:37Speaker 01

  But it's not true like that. Then we wouldn't function, like cars wouldn't go down the road. And so you have to model the the envi, you have to include other people in the environment and understand the embodiment of like how your action will change other people's actions.

• 39:51Speaker 02

  YC's next batch is now taking applications. Got a startup in you? Apply at ycombinator.com/apply. It's never too early and filling out the app will level up your idea. Okay, back to the video.

• 40:05Speaker 00

  Now let's talk about the action space. You know, like, one way to look at the action space is that it seems relatively small. Seems like, well, you know, you turn the steering wheel left or right, you hit the brake, you hit the, you hit the gas. Doesn't seem that big, but like, how big is it actually? Like, how do we actually represent these action spaces when it comes to a realistic self-driving car scenario?

• 40:23Speaker 01

  Yeah, I don't know how they how they do this nowadays. They're doing a whole bunch of like bird's eye view, different things like that.

• 40:30Speaker 00

  I consider even just like a very simplified.

• 40:32Speaker 01

  But what do you have? You have a steering wheel that you can turn left, right. You have a brake pad. And you have the gas.

• 40:39Speaker 00

  Yeah.

• 40:40Speaker 01

  And so

• 40:41Speaker 00

  This thing is like 365 degrees.

• 40:43Speaker 01

  Yeah.

• 40:44Speaker 00

  So it's like a 1 to 365, let's say.

• 40:46Speaker 01

  Yeah.

• 40:46Speaker 00

  Or 0 to 365.

• 40:48Speaker 01

  Yep. And you, let's just say you break this up into 10 different severities.

• 40:53Speaker 00

  You're already, even with just this oversimplified model, your action space cardinality

• 40:57Speaker 01

  Right.

• 40:58Speaker 00

  is 365,000. So that's like 100x bigger than AlphaGo. It's in fact, it's about the size of the example, or it's a decent amount smaller than the size we said.

• 41:09Speaker 01

  Exactly.

• 41:10Speaker 00

  Which is break MCTS.

• 41:11Speaker 01

  And so, yeah, so 36,000 action space is very large. And then even worse, unless you're Tesla. We have a bunch of video of people driving cars. We don't have video of like dash cams like that. Like you actually don't have, again, only Tesla has this of the action as well. And so the things that you have access to, that your trajectories are just like ST, ST plus one.

• 41:33Speaker 00

  Yes.

• 41:33Speaker 01

  ST plus two.

• 41:34Speaker 00

  So there's a, you're saying there's a decent number of these that's from like dash cam footage on YouTube or something, but not really that many either.

• 41:41Speaker 01

  Yeah.

• 41:41Speaker 00

  Relative.

• 41:42Speaker 01

  And so if you wanted to do a self-driving car and you didn't want to go spend a million dollars, trillion dollars on going, collecting all this data, then you want to leverage this data somehow. And this is going to be really applicable for robotics, because we have a lot of

• 41:56Speaker 01

  videos of people doing things.

• 41:57Speaker 00

  Yeah.

• 41:58Speaker 01

  Right? Especially with egocentric, like we, we have those videos, but we, what we don't have is.

• 42:04Speaker 00

  The actions they take.

• 42:06Speaker 01

  Yeah.

• 42:06Speaker 00

  No. So this is like,

• 42:10Speaker 00

  this is, this is a sequence of what you're showing here.

• 42:12Speaker 01

  Unless you're Tesla.

• 42:13Speaker 00

  Unless you're Tesla.

• 42:13Speaker 01

  And Tesla has this. So this is a huge competitive mode of like, what do people do in that state? And then so you can behavior clone to go from here to here, from here to here, go here to here, et cetera. But even then, it's still very, very difficult. You have to, it's, it's not sufficient. People think that like, okay, I have this, we have a self-driving car, right? I mean, the amount of work that they're doing at FSD is like incredible and it's, it's not generally available. Like you can't, you know, it's not Waymo level yet.

• 42:39Speaker 00

  Would this be a good moment to briefly talk about model free versus model based RL? Cause I think that's an important distinction that's going to be relevant when we talk about more world models.

• 42:47Speaker 01

  Yeah. So this is a perfect point. So model free just means that my, my policy pi of A T given S T. I have no world model involved. It's literally, it's literally doing what I said. I grab a bunch of these and I train, go from S to A, S to A.

• 43:05Speaker 00

  Just predict the next day.

• 43:06Speaker 01

  That's it. And that's, and this is logic called VLA. You know, this is like giving us pretty good results. It's behavior cloning. It's all the, the, the stuff that it's not getting us the Rosie the robot just yet, but.

• 43:19Speaker 00

  In many ways, it's the closest thing that just looks like the next token prediction from LLMs that seems to scale pretty well with natural language. I mean, it's not exactly the same thing because there's no action exactly, but picking a token is not exactly the same thing, but it's very analogous to that.

• 43:32Speaker 01

  Yep.

• 43:32Speaker 00

  Like basic thing.

• 43:33Speaker 01

  I basically take away the tokenizer head and I give it an action space and I collect a bunch of teleops data, you know, like this as, as the self-driving car does in Tesla. And I just taken the, the state, which is some image and/or maybe sequence of images. And then I'll output some action and that's it. And this is, let's say, model three. Cause I don't have a model for the environment. And then now if I do model based

• 44:02Speaker 01

  RL, I have not just some pi, but I have also my psi as well here. And so by,

• 44:14Speaker 01

  by including this, I can have a much stronger policy, but it would take a lot more time to perform inference, because I have to do this full test time planning.

• 44:22Speaker 00

  Just to remind us that psi is referring to this specific transition function, right? It's referring to this. You're saying this is specifically referring to a function of ST plus one given ST and action T.

• 44:37Speaker 01

  Yes.

• 44:38Speaker 00

  So it's like your ability to predict the next state you'll be in is, is the crux of it.

• 44:42Speaker 01

  Yep.

• 44:43Speaker 00

  As opposed to just directly predicting the actions.

• 44:45Speaker 01

  Yeah. And the main thing that I believe is that this is required for AGI. This is what the hu, the human brain is, is.

• 44:53Speaker 00

  At least in the way the human brain does it.

• 44:54Speaker 01

  Yeah. And let me go further in saying that, like, if you look at the billions of years of evolution, basically there's this thing called 10 million, 10 million years ago called the great cortical expansion, which you see the size of a brain just explode, get bigger, bigger, bigger exponentially up until us. And it basically stops. And if the entire point of the neocortex is world modeling, what happened is we started from VLA's. This would be like ants.

• 45:21Speaker 00

  Fish or whatever.

• 45:22Speaker 01

  And fish. Yeah. Right. Just like very like, you know, lizard brain, whatever you want to call it. And then we developed this neocortex to like, you know, go from our, our motor cortex to actually simulate what's going to happen. And that makes us just so much smarter. And then we, once we get those samples, we can compress it when we sleep or otherwise with this hippocampal, short ray ripple, whatever you want to call it. And then that helps us develop a better policy and that marriage between the two is, is not only helps us train on hallucinated examples, but it also allows us to test time plan.

• 45:57Speaker 00

  Right. I, I guess the, the kind of extreme case then of self-driving car is kind of general robotics.

• 46:03Speaker 01

  Yes.

• 46:04Speaker 00

  Right? So if you're, if you're like a humanoid company, like figure or pie or whatever, again, same ST AR setup.

• 46:11Speaker 01

  Yep.

• 46:12Speaker 00

  I, I guess the gist of it is that A is now even bigger.

• 46:14Speaker 01

  Yeah.

• 46:15Speaker 00

  Right? It is like, I guess a very simple robot would be, how would you, how would you parameterize the action space? Like, let's take, let's take a very basic one.

• 46:21Speaker 01

  If I take like my six-axis arm.

• 46:24Speaker 00

  Yeah.

• 46:25Speaker 01

  As your, your standard here that we're actually working on right now in Stanford Robotics Center. You have two degrees of freedom, two degrees of freedom, two degrees of freedom. Uh, and then you have another two for the end effector.

• 46:35Speaker 00

  Right.

• 46:36Speaker 01

  And so.

• 46:36Speaker 00

  That's a simple end effector, not even like a, not even like a.

• 46:38Speaker 01

  It's literally a one axis, like, you know, you can rotate, but you have the, the, the, the one axis, you me style thing. So this is eight. So you have 16 degrees of freedom. And let's just say that you do the three sixty five, the one by 10 or whatever, you know, kind of thing. I mean.

• 46:54Speaker 00

  It's like 10 to the 16.

• 46:55Speaker 01

  It's like insane. It's like, yeah, it's an insane number. And so much bigger than self-driving car. And even worse, like getting teleops data is extremely painful and expensive. It's not just like, oh, we'll just get some people in the Philippines. We'll give them like some, you know, things or whatever. It's like totally, totally doesn't work.

• 47:15Speaker 00

  And nor is there yet something like Tesla's fleet where there are cars deployed that people are just using. And they're not even necessarily realizing that every time they turn the steering wheel, they're providing this, this data set for Tesla.

• 47:29Speaker 01

  And then even worse, you have this like what's called cross embodiment gap. And so if I were to like, train this policy on Tesla Model X, and I were to like, put it on a Tesla Model Three, it wouldn't work.

• 47:43Speaker 00

  No.

• 47:44Speaker 01

  Like it totally wouldn't work. Like all the, so much, so much of this. The, the way that if, if I were to break on a Model Three versus a Model X, the Model X weighs more. It has different dynamics, aerodynamics, anything like that. And so what's actually going to happen is very different. Like the, the degradation you have across, across, across embodiments is very, very, very strong.

• 48:04Speaker 00

  And clearly Tesla's figured various ways to get around that. I mean, they, they have these that roll out, but actually even with Tesla's new FSD today, they don't roll out in all the cars at the same time, probably for more or less that reason. And in this case, it's even harder now. I mean, you have bigger differences between embodiments than a Model Three versus Y.

• 48:19Speaker 01

  Yeah.

• 48:20Speaker 00

  And you have way bigger action spaces. You have to have a model.

• 48:23Speaker 01

  Yeah. Lane McIntosh, I played hockey with at Stanford, who now runs Tesla FSD. I can ask him, but I would bet money that they shard the data per model, per car type.

• 48:35Speaker 00

  Yeah. Wouldn't be surprised.

• 48:36Speaker 01

  I just, because that's what I would do. There's no way that like, you know, I, I would trust, you know, data that was collected on a Model X on a Model Three. I just wouldn't, no way I would trust it.

• 48:46Speaker 00

  Okay. So now that we understand the basic setup here and why the action space problem is so big, why don't we talk a little bit about how world models actually fit into this? You know, maybe first, you know, I guess what didn't work about the naive world models and how do we fix those? And then let's kind of talk about some of the newest world modeling techniques.

• 49:01Speaker 01

  Cool. So like in robotics in particular, it's very hard to get these, this kind of trajectories that you want, that you kind of need to train for your VLA's. And people spend up, you know, with a whole bunch of teleops data. It's very expensive, very expensive. Ideally, what we would do is take like data like this from someone who is just like, puts a camera on them and just like making sushi. Okay. Like I want to make a sushi robot. Um, how do I do it? Give it to all the sushi chefs. Don't put anything in their hands and just have them start cutting up sushi, making sushi.

• 49:30Speaker 00

  And ideally we would train it in that way you were describing of like somehow we would train a model just on these two. And then later add this.

• 49:37Speaker 01

  Yes.

• 49:37Speaker 00

  Afterwards.

• 49:38Speaker 01

  And so the first real person that, um, you know, went after this was Jürgen Smidhuber, please. So, so he doesn't yell at us. We have to, we have to make sure we say that.

• 49:49Speaker 01

  But he has this really cool paper called world models, very aptly named. And it's basically, he took these like open AI gym, classic games, car racing, and I think Doom as well. And then just like trained a model at that time was like an RNN. He had some funky zero order stuff in there, whatever. But basically, the key premise was I can take an environment. I can extract a whole bunch of this type of data off of it. I think he actually does actually this data, but we'll get into Dreamer where he does it in this paper, in this way. And then trains a policy on only the synthetic data, the imaginated rollouts.

• 50:34Speaker 01

  And it actually performs well in the environment. This is the first time in my understanding that that actually happened and it actually works really well.

• 50:41Speaker 00

  And then So the key thing there is you can basically use this. If you have some predictive model of this in that case, and eventually of this, you can use that as basically a synthetic training set to train your policy model and then basically fine tune it on real data later.

• 50:54Speaker 01

  Exactly. And which is just like a really powerful idea, especially since in robotics, the limiting step is access to large amounts of state action data. And so now the Dreamer series, so basically this publish publishes in May of 2018. Danajar Hafner publishes Dreamer one, I think in November of 2018. And then now he's been on this rampage for the last seven years, publishing these papers. And Dreamer V4, I think is the capstone of it, where he basically does the same thing and he focuses on Minecraft. And he trains these, a world, a world model on this type of data.

• 51:35Speaker 01

  And then injects action conditioning on a very small amount of data to get to this type of world model that can, that has the action conditioning as well. And then samples a lot from it. And then trains a policy on those synthetic, uh, imaginated rollouts. And it's, the policy is so good that it's the first paper to mine diamonds in Minecraft. I'm not a good Minecraft player, but apparently that's extremely difficult. That's like next level difficulty and it did it all on synthetic data, which is kind of crazy.

• 52:05Speaker 00

  And the key unlock there, yeah, use synthetic data specifically on a model trained on just this sort of state transition type of thing.

• 52:13Speaker 01

  Yes.

• 52:13Speaker 00

  And this ends up being very convenient because it turns out we as a society have a lot of this.

• 52:18Speaker 01

  Exactly. Yeah. All of YouTube, right? He does do a very small amount of data from, to enable the action conditioning. And that, that allows you to do this full simulated rollout. But yeah, it's true. So we have, we have YouTube. We have like Flickr, we have all these data sets online of like, you know, people doing things. We'd like to use it and no one has really gotten that to work. And then now that with this, um, these like video generated generation models, we can take that data, create a world model out of it. Add action conditioning, post train it with action conditioning for some new task that is, we want it to do chopping down wood or, uh, you know, making sushi or folding my bed or whatever it is, only a few amount of examples.

• 53:02Speaker 01

  And then we can train a policy on this, in this neural simulation.

• 53:08Speaker 00

  Yeah. And you know, we put out a video about diffusion models very recently and flow matching. I imagine that now ties very closely to this, right? Ultimately, the, the kind of current state of the art, best way to do this on basically infinity data that we have available and can keep generating is using, say, our video diffusion slash flow matching.

• 53:26Speaker 01

  Exactly. Yeah. So like, if you have your, your C dance or your Sora or

• 53:30Speaker 00

  One or whatever.

• 53:31Speaker 01

  Exactly. All those models, like basically the idea is now we have them and they're already trained and they're great. Let's do a small amount of action conditioning on them to get to this, uh, this world model. And then we can sample from it a bunch and then train. And this is exactly what we've, uh, did with Gaia. And Gaia, I think they raised 1.5 billion dollars to basically run with this idea for self-driving car. Um, I think a bunch of companies, um, Nvidia, uh, this, this paper here, uh, is basically talking about doing exactly the same, this dream zero for robotics. And

• 54:07Speaker 00

  What I thought was really cool about this paper is that they, they do exactly this process where they have this, um, joint model of, um, state transitions and actions. They train it by first instantiating it with the open source one video diffusion model. And then it only takes them about 500 hours of teleop data, which is basically exactly this.

• 54:27Speaker 01

  Right.

• 54:27Speaker 00

  To get it to be pretty good. And they have a lot of clever tricks that allow it to be cross embodiment and work on unseen task with relatively small amounts of data.

• 54:34Speaker 01

  Right.

• 54:34Speaker 00

  And it really is taking basically the exact concept, I believe, from the Dreamer paper and applying it specifically to these robot embodiments.

• 54:41Speaker 01

  Exactly.

• 54:42Speaker 00

  Um, and it turns out it actually works, uh, actually better than I would have anticipated it working.

• 54:46Speaker 01

  Right. So yeah, so I think that this is basically the, the, the path to, it was the path, I believe it was the path to get humans, uh, to be as good as we are genetically over the last 10, 20 million years of evolution, a bigger world model helps, uh, for training and for test time planning. And I think it'll be the same thing as true as for robotics.

• 55:09Speaker 00

  What's also cool is there's a bunch of applications of this, so things outside of robotics too. I mean, there was a weather planning paper, for example, we were reading this GenCast paper, which I think applies a relatively similar concept in terms of how they model, you know, literally the world, the world's weather with something like this.

• 55:26Speaker 01

  Yeah, we have to talk about the world model for the world. Yeah. So basically they do this exact same thing where, you know, the key unlocks for this whole thing was getting diffusion to work in very high dimensional state spaces, like we talked about in the last lecture, and then learning to, to use that to can action condition in the way that he's done. But they did this for the entire world with this exact same diffusion steps, which go from some, and they go back to, uh, two time steps lag of, of order two, AR two for the statisticians there. And then basically predict the next, uh, state of the world based on the, those things with this Lingiv and diffusion rollouts.

• 56:07Speaker 01

  My, my big assertion is that, um, it was necessary for the human brain to develop world modeling. I actually just saw this paper that I wanted to make sure to call out because I thought it was so great. Uh, out of, uh, University of Washington, where they say explicitly in the, in the abstract, each cortical area estimates both latent sensory states and actions. And the cortex as a whole predicts the consequences of those actions. That sounds like a world model to me.

• 56:37Speaker 00

  Yeah.

• 56:37Speaker 01

  Right?

• 56:38Speaker 00

  It's actually describing exactly these two equations here.

• 56:41Speaker 01

  Exactly.

• 56:42Speaker 00

  Where we're estimating both the sensory latent states and actions. I mean, I guess it's really the joint model that we showed earlier.

• 56:47Speaker 01

  Right.

• 56:47Speaker 00

  Is what he's describing here. It's exactly this, this equation he's showing now.

• 56:52Speaker 01

  Exactly right. And so, uh, if it works in us, it should work in robotics. Um, and I think that that takes us the rest of the distance.

• 56:59Speaker 00

  Why don't we talk briefly about latent world models, especially the JEPa concept? Cause I think there's been a number of papers that use JEPa as an element of their, I guess, architecture. Why don't we just briefly introduce JEPa and how it fits into the current landscape of world modeling?

• 57:15Speaker 01

  Yeah. In classic RL, you'll have like, you know, if you do study Q learning, for example, you basically keep this matrix called the Q matrix.

• 57:23Speaker 00

  Yep.

• 57:23Speaker 01

  And it's going to be, uh, S by A. And so I have this, um, S by A, states and actions. And each one, I need, you know, some amount of counts of being in this state, action. And I take the average value of being, of taking that action in this state.

• 57:44Speaker 00

  Yes.

• 57:44Speaker 01

  And that's my key value there. And it's a little bit more complicated than that. There's Bellman equation, all this backup, all this stuff like that. But so this scales horribly because as the cardinality of my state space gets bigger and my kind of action space gets bigger, stuff, I don't have enough. And I become less and less sample efficient, right?

• 58:01Speaker 00

  And in the case of like robots or whatever, state is like, yeah, it's this whole thing we described earlier, right? It's absolutely massive because it has all of these elements in it. You couldn't really enumerate a huge grid.

• 58:09Speaker 01

  And so the classic trick, I mean, since I took, you know, uh, C229 with Andrew Ng in 2012 is you do this.

• 58:17Speaker 00

  Stick a neural network on it.

• 58:18Speaker 01

  Exactly. And you basically are just going to compress that state into some lower dimensional state space. This is actually predates deep learning. Uh, we were doing stuff like this. I think my first paper was basically doing something like this, uh, basically turning like a grid into like, uh, a bunch of like pyramids. And like, and, and the state was how much I'm in pyramid one or pyramid two or whatever. But anyway, the neural network, you can just do this. And so basically what, uh, the, the key idea in JEPa, if I have, um, an image one, and I have image two, and I have image three.

• 58:54Speaker 01

  I can do my, my world modeling, uh, my, my world modeling of ST plus one, uh, given ST and AT in pixel space and have, this is, uh, let's say at time T, T plus one, T plus two, et cetera, et cetera. And I have to actually predict now the full, uh, image. That's extremely expensive from a computation standpoint, and also from like a sample efficiency standpoint. What I can do instead is put this through some ComNet.

• 59:26Speaker 00

  Some encoder.

• 59:27Speaker 01

  Some encoder. And then I'll get a latent for T. And I'll have a latent for T plus one, I'll have a latent for Z T plus two. And then I'll have from this, from Z T, I want to predict Z T plus one hat. And then my goal is to make this and this, uh, make, and my loss function will be something very simple, like.

• 1:00:00Speaker 01

  Want to minimize this. That's it. Now, this doesn't work. This collapses hard. And so what happens, like, it's basically just if you, if it just predicts zero.

• 1:00:09Speaker 00

  Yeah.

• 1:00:09Speaker 01

  Done. Just zeros, which the model will learn to do. And I'm actually incorporating this into my current research right now. And so what you need to do is something called Sigreg. Or, uh, this is one technique, Vicreg is another. Where basically I add this another term that basically says, uh, I want the, um, over a large enough batch size. I want the, the distribution of Z. T plus one to follow a Gaussian, you know?

• 1:00:39Speaker 00

  It's kind of like a normalize it like a, like a batch norm type of.

• 1:00:42Speaker 01

  Yeah.

• 1:00:42Speaker 00

  Trick. I mean, not in the same.

• 1:00:43Speaker 01

  And, and if, if it's zero, it can't be this.

• 1:00:46Speaker 00

  Yeah.

• 1:00:46Speaker 01

  Right? Cause then this is non-zero. And so maybe I think that there's probably this or something like that. But basically this prevents it from modal collapse and it makes it do something good. And this is the most recent paper for the audience is LEWM, LE World Model, which is super, super great. However, to be completely frank, the, this, this is self-supervised learning, super great. It doesn't work that well. If you were to not do, uh, these techniques and there's, there's a bunch of other techniques that you can do. It will actually outperform much better and that are, let's say, for example, um, if I'm going to do an LLM. And you have like, you know, Francois. Likes sushi, which is definitely true.

• 1:01:31Speaker 01

  And I tokenize this into a bunch of different tokens here. And this is token ID six, nineteen, twenty-eight, whatever. And I look up the encoding into this and it's going to be, uh, E one.

• 1:01:46Speaker 00

  Yes.

• 1:01:47Speaker 01

  E two.

• 1:01:50Speaker 01

  E three, et cetera. What you can actually do is have the LLM output, uh, with the LLM will take in, take in these things and will output.

• 1:02:05Speaker 01

  The, the next token. And so it would be like, let's call it H. Uh, this would be the low jets coming out of it. T plus one. And what you can do is actually have this be close to ET plus one.

• 1:02:19Speaker 01

  And a lot of people are playing with this idea and getting rid of the cross that should be lost entirely. And so if you were to do this, it actually is a proxy for the cross entropy loss. And there is no cross entropy loss. And the, the cross entropy head is actually very expensive. And so this is very cheap and like this lady just grabbing it. So people are playing around with this idea. And as a, as a, basically as a, as a cheaper proxy for the cross entropy loss. So there's a lots of different ideas on basically, uh, taking this Jpy idea to not just pixels, but to L, LNs as well.

• 1:02:52Speaker 00

  Yeah. Interesting. Yeah.

• 1:02:54Speaker 01

  So just to define what Jpy is, it's joint embedding predictive architecture.

• 1:02:58Speaker 00

  I, I think one of the things I find, uh, cool about this Jpy idea is it feels like an idea we see over and over in deep learning. And there's a version of this idea that's basically the stable diffusion idea.

• 1:03:08Speaker 01

  Yep.

• 1:03:08Speaker 00

  There's a version of this idea that in my company, training graph convolutional neural networks to, um, design drugs. We use to do, you know, latent variable generation, for example. And it's like, it's an idea that comes back over and over and then has this, yeah, various tricks that it actually takes to get it to work in practice. Okay. Now we have a pretty good sense for how world models work. We have a pretty good sense for what the state of the art looks like. If we trust this paper, and it seems like these kind of work on robots too. This paper is only from the end, uh, end of last year, this year. And it seems like they have various methods that allow you to train on relatively small amounts of data. That's tractable and pre-train on diffusion models. So are we good?

• 1:03:48Speaker 01

  We're done.

• 1:03:48Speaker 00

  It all work?

• 1:03:49Speaker 01

  Yeah. This is 2016, it's 26 will be the year of the robot. We're gonna have Rosie the robot in your house, you know? Yeah, no, I don't think so.

• 1:03:58Speaker 00

  What are, what are one or two, you know, because there's lots of open problems remaining. What are like a few open problems? Maybe we can emphasize here that the community can go emphasize working on.

• 1:04:06Speaker 01

  Yeah. So, um, I think the first one is that, uh, PINS doesn't really work. What is PINS? Physics-informed neural networks. So PINS.

• 1:04:18Speaker 01

  Doesn't really work. This is physics-informed neural networks. And so basically if like almost all of the self-driving car data looks like this, the car is driving down the road. And let's just say, for example, uh, I have, you know, a, a house here.

• 1:04:36Speaker 01

  And I want to train the model on you not driving into the house. And so let's say I put, I put it into a state right here to drive into the house. What's going to happen is because almost all the data is like, looks like this, driving down the road. This would just turn magically into like a highway. And I'm just like, boo, just don't worry.

• 1:04:57Speaker 00

  It basically is like a ton of data not to do that, either from simulation for that to not happen.

• 1:05:02Speaker 01

  In fact, I actually don't even know if because of the data distribution, there's no data here. There's almost all the data here. And like, when you're training a neural network, it has a tendency to collapse. If you don't keep the mini, mini batch composition, uh, like very even over the, you know, over the class space or whatever you want, want to call it. But like, you'd have to train on, uh, you have, you have to be very careful about your data mixing to make sure you get this right. Just solve this problem that no one really has. But even then, the, if you take just a simple thing like this, this is like the, the iconic example. And I have some sine wave and I want, and I have these as my X.

• 1:05:47Speaker 01

  Ba, ba, ba, ba, ba. And I have these as my Y. So this is complete interpolation.

• 1:05:53Speaker 00

  No.

• 1:05:54Speaker 01

  Maybe mess this up, but Y like this.

• 1:05:56Speaker 00

  No.

• 1:05:57Speaker 01

  We can't get to like machine precision. We can't, what is, I don't know, what is it? One eight minus sixteen, whatever it is. We can't, we, we, the SGD will not get to zero, effectively zero. So we'll always have some residual. And for us to be like a really good world model, to simulate body interactions, like to, to simulate this, what's going to happen when I do this? And like, let's say that I'm trying to be LeBron James. Like there's like, I saw this one video of, um, Steph Curry dribbling about a basketball on a court. And he just felt that there was a dead spot in the court. And he, cause he's so good and he knows exactly the physics of what's going to happen. If I hit this, you know, the ball with this force, like the ball is going to come back exactly to this spot. And it just didn't. And he knew it wasn't him. It was the, the court.

• 1:06:40Speaker 01

  And he found a dead spot in the court. Like that's how good the, the, the human brain is at world modeling. In my opinion, I think it's an SGD issue. I think it's probably an architecture issue. I think Sam Altman just kind of came and just said that he thinks that there's definitely an architecture that's going to be more performant than the transformer. I think he's right. I think the, the transformer doesn't do compression, uh, uh, in the time domain at all. It just keeps around everything. Um, so anyway, so I think that the getting higher fidelity in the world model is extremely important. One, I think two.

• 1:07:14Speaker 00

  Seems like test time probably is going to be a big thing, like adaptation.

• 1:07:17Speaker 01

  Exactly. Test time planning, we, the, how quickly the human brain can, you know, i- i- in times of, in sports and things like that, when you're playing tennis, I think you're a tennis player. Like how quickly we can adapt to what a player is doing and things like that. We're not going to sleep and like retraining. We're, we're very quick to adapt to a new, new environment.

• 1:07:39Speaker 00

  Like the out of distribution prediction.

• 1:07:41Speaker 01

  Exactly.

• 1:07:41Speaker 00

  Really challenging.

• 1:07:41Speaker 01

  And like one little data point, we can like quickly adapt to that new thing and change. I think there's been a lot of papers, uh, uh, on like basically estimating the, the friction coefficients. And so like those can change over time. If you go to a human environment or not, for example, like this, this friction might change. And that's important in control.

• 1:08:00Speaker 00

  No.

• 1:08:00Speaker 01

  Um, and so you need to estimate that very quickly and adapt. And that these models just kind of don't have a mechanism to do it.

• 1:08:06Speaker 00

  Yeah. And then I guess there's like the practical s- speed elements of these, right? A lot of these are doing some sort of expensive planning step. And we're doing some sort of like, uh, we're, we're kind of hacking around it with this retraining process and synthetic data. But even so, like to really get maximum performance right now, you'd want to do something that's closer to like the AlphaGo style rollout. And that's extremely slow.

• 1:08:31Speaker 01

  Right. The MCTS process can't happen. The other thing that, that is pretty crazy about the way that the brain works is that like, everything is kind of running autonomously. And so like you'll, you might be like in the middle of saying sentence one and be like, oh, actually no, something else. And so like, what does happen there? It's like type one and type two thinking are happening at the same time in some way. And so like, there's definitely, uh, you know, some, um, really cool mix of these like heterogeneous models and like some are overriding others and like taking control of the motor cortex and like commanding the body to do a thing, you know?

• 1:09:07Speaker 00

  Okay. But on the flip side now, we, um, talked in the past video about the squint test and how we felt that autoregressive LLMs maybe don't pass the squint test. Why don't we reintroduce what the squint test was for a second? And then maybe let's think about whether this passes the squint test despite all those limitations.

• 1:09:23Speaker 01

  Yeah. And the squint test for me, I think is like, um, this comes from the Yann Le Cun. We didn't need, uh, flapping wings to achieve flight. Um, and to that I say, well, we did need two wings. And like, if I squint and I look at a bird and I squint and I look at a plane, I'm like, mm, yeah.

• 1:09:39Speaker 00

  It's kind of similar.

• 1:09:39Speaker 01

  It looks right. Um, similarly, if I squint and I look at the human brain and I squint and I look at all these, these world models. We have like this VLA, this action policy, and that they're doing test time planning together and things like that. It's getting really close. It's much, much closer.

• 1:09:55Speaker 00

  Seems closer than an autoregressive LLM and that like this concept of a world model of, you know? Implicitly predicting future states and actions feels intuitively like what our brains do, and it seems like there's some neuroscience evidence to support that.

• 1:10:09Speaker 01

  Yeah, I mean, I'm getting to the conclusion that I think that the brain is the optimizer, not the model, and that the brain emits, like, has models that it invokes, but the brain is somehow also the optimizer itself. And so in that way, it doesn't pass the squint, because like, you know, something magical is happening when you're sleeping. There's no intelligent species that we're aware of that have any amount of intelligence that don't sleep. And so like octopuses, dolphins, all those elephants, they all sleep. There's some reason for that. And that seems like a really thing about like the evolutionary recourse of sleeping, like you get eaten when you sleep. So like, for the benefit of sleeping should be so, so much better to outperform that. So I think we don't have

• 1:10:52Speaker 01

  this idea of awake sleep in our current architecture, but I can imagine, I'm like simulating, you know, compressed from the hippocampus, some like experience in the day, I'm like training on more of those examples, right?

• 1:11:04Speaker 00

  You're like collecting a whole bunch of these experience rollouts and then you're updating your policy function overnight or something.

• 1:11:11Speaker 01

  There's got to be something like, there's this thing called short wave ripple, where like the hippocampus when you're sleeping like emits these spike trains that are actually reversed from when they actually happen back in through the both, both the hemispheres and for like seven times and then it like stops. So like, there's something happening there that's very training something. Yeah. And if you don't sleep, then you don't, you don't have long-term memory.

• 1:11:33Speaker 00

  Right.

• 1:11:33Speaker 01

  Right? And so like, there's definitely a reason why we're training things that happened into our brain.

• 1:11:40Speaker 00

  So where does that put us now? We have all this work happening with world models. How should we think about what's coming ahead in these next few years in the research community?

• 1:11:47Speaker 01

  Yeah, I think that like, we're going to see a lot more of these world models in robotic policies. I think that's going to unlock probably full self-driving would be like a, one of those examples that they can get the real time-ness of it.

• 1:12:00Speaker 00

  Seems like that's coming.

• 1:12:01Speaker 01

  They can probably solve it with more compute to like have parallel things. And you probably don't need it for like most standard things, maybe like, you know, getting out of weird parking jams and like, things like that would take us some time, similar to the Rosie the Robot, which we've always wanted to have a Rosie the Robot to like, you know, clean up my room for me. I think that like, this feels like we're getting to good enough that we can pay up for data and compute to get to Rosie the Robot. It does feel like that. It'll be expensive to collect the data and do the dreamer sequence of going from state to state and then getting the action conditioning to work. But like, I feel like it should work.

• 1:12:38Speaker 00

  Yeah. I mean, what's pretty cool is we see a lot of companies at YC working at every step of this, from the collecting egocentric data, collecting the teleop data, training their own world models and action models, building new embodiments and then making ways of adapting those embodiments. And it feels like this is the first year where you see demos where you're like, okay, this actually like, kind of is starting to look like it's going somewhere.

• 1:13:02Speaker 01

  Yeah.

• 1:13:02Speaker 00

  And it seems like a very exciting year.

• 1:13:04Speaker 01

  Yeah. So anyway, I think that there are real AI problems to solve still. We talked about pins, we talked about the real time issues. And then on the robotic side, there's a real issue, like it's amazing how effective our epidermis is in terms of we can detect tactile.

• 1:13:20Speaker 00

  Oh, epidermis.

• 1:13:21Speaker 01

  Yeah, epidermis. Our tactile, we can detect shear force, we can detect temperature, and it's everywhere. And so like versus, you know, like the, we get like one little sensor that only does tactile. We don't have the friction component. We don't have temperature. We don't have all these, the feeling. We can't estimate coefficient of friction very quickly. I can touch something and say, oh, this is smooth, this is rough. It doesn't, we don't have any of that. And if I numb your hands, I actually had this experience just recently, if I numb your hands, like you actually can't tie your shoes.

• 1:13:50Speaker 00

  Yeah.

• 1:13:51Speaker 01

  So you can't perform control. And so like, yeah, if you like, you know, if you train enough, on enough human data, tying your laces, do I think you can do it with no feedback? Maybe, maybe. But like, how much would you need if you did actually have the human like touch? Like, I think it'd be so much easier.

• 1:14:10Speaker 00

  Yeah. Well, there's a lot of more research to do then.

• 1:14:13Speaker 01

  Yeah, yeah.

• 1:14:14Speaker 00

  Francois, thanks so much for joining us. Thanks so much for watching, everyone. We'll be back with the next episode of Decoded.