World Models, Explained

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? 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. So today we're going to discuss what many top researchers believe is the most promising path to closing that gap. World models. 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.

You and I have talked a lot about the various ways people are training models and sample efficiency of them. Why don't we start by just defining sample efficiency and how we intuitively think about it as humans? Yeah. So I think from my perspective, the two major problems that we haven't left us off in 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 data set, how much more intelligent am I getting? And so if I imagine I have a new task like our AGI, for example, I think like really for sure, like it's 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. And in context, the ArcGI test tests 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 system is what people consider frontier intelligence basically can't do them. Right. I mean, we come into new problems with such inductive bias from K through 12, like all these math and school that we've had, that these models are kind of getting from the entire depressing the entire internet. And so when we come in, we're not coming in tabular also just like bare bones, but even so that they have, I don't know what percent of the internet you've read.

I read very little percent of the internet. But despite that and having read the entire internet, it still can't really do well and generalizing to these new tasks. 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? 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 the example of 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. So I know it actually can happen. We do a lot of time.

It's called new and second level motion. It's like new 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 Newton's laws of motion. Yeah, like when NASA plans to intercept an asteroid and planning it, you know, years in advance and can set it off in a trajectory where it just glides the right thing and intersect to the right point. That is an example of a perfect world model we built where we're then just letting the world model act.

And that system does not need to intelligently collect new samples from the environment to decide which direction to go next. It can already be pre programming can perfectly do it. Yeah, can you imagine if like we need to collect one million training examples of like us shooting spaceships to the moon to like know how to do it? It's like this completely would definitely have the publishers, right?

But we do have that ability because the real world is differentiable and we can do something called model control that we're going to talk about in a little bit. But even in our own brain, I was just 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 were going to say or what a VC is going to say when I was pitching them or what customer might say. And even products being having taste.

What is taste is like predicting that other people are going to like this thing. And so we 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 basically this is actually proven in the 1967 Kogsei study by Richardson.

They basically showed that if you take a cohort of three different groups of people and you have one go practice layups and basketball and they go on the shoot. They improve it 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 up basketball, they improve at 23% against the control.

I mean, that's insane. That means that we have this crazy good world model and there's this neuroscience at Stanford named Shawl Drachman who basically is of the view that the entire point of the growing neocortex during the great critical 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 next prediction. 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. Yeah, what a fundamental comes down to is we as humans we think about our intuitive ability to think as coming from some implicit world model we have in our heads and code it by genetics and ability to learn 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's 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 maybe it would be helpful now to just think a little bit about and sort of define some of the pieces of what makes it challenging in the future. And then we can use that to kind of build up to why it's particularly hard in things like self-driving robotics to get these types of predictive models to work.

Yeah, let's do it. So let's actually take a step back and just talk about like control, reinforcement learning and define some common terms. So typically in we teach a course of decision making under uncertainty, which is like the main reinforcement course at Stanford. I like to show a specific example of let's say I have some drone and this is my poor little drone here and 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. And to be really clear, this is going to be P X, P Y, P Z, T T T and V X, V Z, V Z. It's like a six dimensional state vector. 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 I want this thing, this is a control problem. Right. And so this is like, I'm going through optical or optimal, 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 S T plus one given ST and my action, which I control is UT. And so this is my state transition or dynamic function or a world model. This is like a very fundamental for context.

Equivalent transition function you think about in RL in general. Exactly. And so, and what I'm trying to learn is something called a policy, which is like what UT should I emit given some ST. And so this is the ultimate question.

What should I do? What action should I take given some state ST? And so the way that will solve this, and luckily we have a world model that is perfect. This is new, new, new physics is like new in second law of motion, which is F equals M a.

And so we know that the position PT plus one is going to equal PT plus delta T V T plus one half of the T squared. Everyone's taking high school, high school physics. And the same thing for the velocity. And then my acceleration is some of the forces, which is going to be my UT, 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 object 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 some over all T. You have UT to infinity. And I'm going to minimize my P star minus PT plus V star minus VT.

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 U max thrust. Yeah, I can be achieved.

And so this is easily solvable with comic optimization. And so this is convex. This is convex. This is convex.

The sum of convex functions is convex. This is a convex constraint inside DCP, discipline, convex programming means I can put this into CVX pi and it will just give me out my policy, which will be the solution will be the optimal UT plus one all the way to infinity. So we can solve this in closed form. We can because we have this world model of Newtonian physics, we can say every step exactly how this drone should fly.

So that lands on the perperic. Exactly. Under a set of constraints like Maxwell's. Exactly.

You'll run your log barrier, enter 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 right to get to this state and that will minimize. And then I can do increase this if I want to do the least energy path for if I make that zero if I want to be the fastest. That's typically the way that you would do what I would call deterministic.

Differentiable control. And why differentiable? Because I can take the I can form the Lagrangian by taking this minus this constraint and take the gradient of it. I can do the Robbins.

You use the fact that it's differentiable to do the optimization. Exactly. If there's a non-differentiable you cannot do convex optimization. You cannot do SGD.

Even if it's non-convex, you can still solve and get a pretty good solution as we do in deep learning. But if it's non-differentiable you can't do. So let's have an example of how you could make this non-differentiable. What's a scenario?

Even if it's like this drone scenario where it now becomes non-differentiable. Yeah. So I'll put this adversary named on it. And your job is to you have another drone.

Let's say on kids drone is to try to hit me and stop me from getting there. Now from the position of your drone you don't know what actions I'm going to take. Right. And so now let's just call this the this would be now we're definitely not deterministic.

We're stochastic. And stochastic and non-differentiable. And in this case my state transition what is ST plus one. It's going to be my say I mean now my thrust and what I'm going to do.

Right. And these it was all differentiable until variable. Yeah. I can't like back prop through your brain.

Yeah. To say what you're going to do with your little drunk controller. Right. It's completely non-differentiable now.

And I'm resorting and I have to resort to this awful area called reinforcement learning. It's super brutal and it's sprawling and there's so many different things. And you'll hear things like when you study initial reinforcement learning called value iteration or policy iteration. And there's DQN or DQ learning or just Q learning.

There's after critic. There's all this. All of this stuff ultimately comes down to ways to estimate to model this non-differentiable stochastic process. Exactly.

Yeah. And so like that's basically the main thing is you're going to start talking about this as a model where I'm going to introduce this. This is going to be some model is going to take in these things and then output this and that we're going to train it over many, many instant instantiations of this. And that's so you get a better and better world model.

And then I need to train some policy. At ST. And then typically you also need a value function. And that is the value of some state and to discern between the value 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 like the SpaceX with launching rockets and landing rockets 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, maybe not highly valued.

So I might say as an example, they typically call this like some kind of a cone here and I might say like it's low value to be here and it's very high value to be in this column or something. In a sense of value gives you some expectation of future rewards like the sum of future rewards you're getting. So if you're in a bad space, you would set the value to zero or negative from negative to negative. Yeah, so we can we should introduce our R.T.

as well. And so typically like if you're playing go or chess like winning the game, you can say winning the game is plus one minus one for losing draw zero. That's what's done enough ago in chess. We have these heuristics like a pond is worth one point a rook is worth five or set or set or so you can like already have reward is the difference in board board state.

And then this yes will be the sum of my discount. Well, of our T given and it's important also to use this nature. Be pie and the reason why that's important is because what it what's actually happening here is this is the discounted reward following policy pie. Correct.

And that means that when I'm in the state, I will take a section and I'll end up in this as you plus one and I'll take this action and it's taking it greedy. And so that's the value with respect to pie. And so ultimately what it comes down to is we are trying to still find a new policy pie and along the way we will use the learning models in various capacities. To estimate the value function given the words we're receiving.

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 it became a more intelligent policy. Right. And so your standard kind of set up for this is what I'm always trying to get to at the end of the day is some joint distribution, which would be s t plus one given where I'm at now.

Where I'm at now. And then this factorizes with chain rule simply to my pie, my policy, a t given s t and my world. And I'll give this is usually represented with data. And this is my world model, which would be s t plus one given s t and a t.

And so, and these are typically learned 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 s t and I predict the right frame s t plus one. And then, and we'll get into this.

For those of us who kind of saw our diffusion model series often people these use video diffusion for exactly this. Yeah. And then what you can do, and this is like the in vogue thing to do since denage our and dreamer paper series from v one to before is. Do action conditioning later, like similar to clip where we will inject this like input head input tail to come into the model to influence 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 tree on the run on the side of the building, I can like see the world go passing by, but I can actually influence it and I can change the world and I can learn that with a t. And it's far fewer samples to do this post action conditioning. If I already had a really good s t s e bus on the world model.

And so here you're saying, you know, what's also in vogue now is jointly training these verses separately training them exactly. So this is called that is called the world action model where some of the issues here is one, there's all these train dynamics of these things are disparate train 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 sample action into here and then roll it out to s t plus one and it's very expensive and it's a very simple time to major issues and why, like, why can't we just scale up alpha go 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 here's the action I should take and here's the s t plus one to end up much much cheaper and much faster.

Okay, so I think it's really good. So 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. So let's go through a couple examples of problems that we want to apply reinforcement to so chess is a pretty easy one.

There's an eight by a grid. And so typically when you approach any problem, you're going to look at star. And so this size of the state, the number of states I can be in. So if I have these eight here and these eight, so it's eight, 1632, so it'd be 32 to the 64 quite large quite large.

And then my transition function is stochastic and not differentable because you can go with the other. So if I'm like in playing just like on my house, I move and then something happens and it comes back and then now you moved and the board has changed. I can't really differentiate through what the other player is doing. The current action space is actually quite small.

Even though there's 32 pieces and all that stuff that is only eight possible moves in expectation that you can actually that are legit moves. So like in a given state, but only if it moves, you could do it. Like I just say in the beginning, I can move on my pawns. I can move my horses.

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

You say it's tractable, even though there's a really big space here. Yeah, but when 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. So we talk about it for just a second.

Yeah, maybe we can add go and contrast the two. Yeah. So why don't we do that because it's because I want to get to the AlphaGo the way that they saw this. And you're right.

So if I were to do this naively, and I just took an SD plus one and I want to do look ahead. 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 to action eight.

And then each one of these, I need to expand it for all possible states. And so now I need to do carnality S, which we just said is a huge for a 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 looking forward, one move is like white and track. Although at the same time, everyone starts at the same starting position. And while it is a really large space, there isn't an infinity number of potential.

There's actually a really small number of game boards, even four moves into the game, as opposed to a game where you can start any permutation, for example, a racial game state and what if you say, still? Yeah. So this is like definitely over done because it's much, much less in this practice. But just naively, like looking at, you know, what possible game states could be 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, 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 off the go. But like, this ends up being estimated in the leaf node, because at the end of the day, my policy, at ST, I want to pick on the argmax of like the value of the following action.

Yeah, exactly. Yeah. The argmax over A of the value of the state of the end state ST plus n, let's say, is like, that's the main goal here. And so for me to do that, I need to roll all this out, estimate the value, and then pick the best one.

And so this quickly grows. However, and we'll see this about we go, which actually has an even bigger state space. So I think it's 19 by 19. I'm very fun.

So you have this 19 by 19 grid, you can in each one, it can be black, white, or nothing there. So I have three. So let's do our star again. So the currently, the state, I think, is going to S two or three, my turn everything here.

The 19 squared, I think it's a one. You see one my transition, same issue. I don't know. My action space is going to be three 61, let's say.

So it's a good amount bigger than chess much bigger, but it's still not enormous. Yeah. And so basically what they do, they call this Z, which is kind of annoying, but let's go 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, then all the way to the end of the game. And if you won, then all of these, all the moves that black, if black one, all the moves that black did get plus all the moves that white did remind us one. And they just, that's how they create their, their role else. Roll out refers to a taking end steps of play of all players, one after another.

Yeah. Of moves under a specific policy of the particular instance. Right. Right.

So it's probably under this policy p theta t and we're an 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're going to sub sample a bunch of these state action results, state action results to train our to update our policy and our world model transition model.

And what it's actually doing is we take an ST, we give it to some theta and it wants to output the probability of S D plus one being played, which is our transition function and the value of the current state. And how do we get the velocity? And so the value of the current state, well, both of them are coming out of the model, but basically the loss function L theta is going to equal. And it's going to be really close to this control problem.

One is we have some theta minus this Z, which we'll just call it R here. And then plus actually, so X minus this pi, which also in a second log p theta. And I think they, everyone includes this, but they include it in the paper so I'm clear there's well, which is the way to get it. And so, so this is basically what our loss function is, then we'll play a bunch of these games, let's try to be a little bit organized here.

And so this is our setup, this architecture. And now once we train this thing, we do an insane, insanely expensive task of of test time planning. And so this trend in RL is just called test time planning. And the specific algorithm that you use here for this is much of our research.

And so this is one of the possible things that you could do and ends up working extremely well if you have small action spaces. Yeah, so let's like really talk about what MCTS does. And a lot of people have heard about once in college research because AlphaGo was such a big moment. How exactly does that map into our starting value function policy?

Yeah, so I'll take this ST, this will give me 361 numbers that sum to one. And so I'll have some probability of where these things are going to go for where my opponent will play here. So these are like the sets of actions. Yeah.

So I'm here. So I have all my S, C plus ones that have 361 of these things. And then this is like action one action to all the way to action. Exactly.

And the, 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 UCB, the upper confidence balance, which is this equation that is roughly something like balancing my value function of S, T plus one, which they're going to be in the literature, they call the Q value because it's actually the difference between value function and Q value is just that I have the action as well. So it's be ST and AT.

So it's called a Q value, which is my exploitation term. And then my exploration term will be something like this funky square root of N. So it's the argmax of A of my Q. And then I have this, which is the probability of this moving blade, which we have from here of S, let's just call it ST plus one.

And then I have this term, which is this sum over NSP divided by NSA. And what's the intuition on this term? So these ends is the visit count during my NCTS process. So this whole tree, I'm going to.

So this tree can get really big, it's 361 per thing. So we have this depth of 30. So this is every single week. Exactly.

And so you want to keep track of which, 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 have good exploration. Right. And so the way you keep track of 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 local.

Very myopic. And so what you'll do is during this MCTS process, you'll start this dictionary, which will be all zeros of the visit count of being in this state and taking this action. And then once you go through your first rollout, you'll do, you'll go here, you'll all these things will be in it to zero. You'll have some probability what we're going to bias it towards the higher probability of places to go and then we'll expand those trees.

And we will update the counts that we visited this and that will basically reduce the amount of probably that we're going to select it again because this will reduce my exploration term. And if it's highly valued, then we're going to increase the key on this because this is the expected value of going down this this map. 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 node you could traverse to right in these 30 step rollouts. 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. Yeah, 24,000 invocations of the model to develop this tree and then what's I have 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 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. By saying 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 a pie that you'll train in your loss.

So you have the R of the winner lose, we have the pie of what was the end result of this whole expensive process. And that test time, we are going to do these 24,000 steps every single every single move to pick the argmax that gives that satisfies those exploration exploration exploration and exploitation. In this case, you know, this still feels somewhat tractable though because the action space is small enough for this like kind of works exactly now let's say hypothetically maybe we can draw like an imaginary go a game of go where it's like, you know, let's say this game of go was like 1000 by 1000. And so now you have a equals, you know, more or less, a million.

And now this this tree we're drawing here that has to take here this has cardinal or like, you know, with a guess 1 million. And there's like S zero through S 1 million. And the number of steps you would have to take here presumably 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 doing a next up prediction astronomically if the game was even let's say, you know, this is only 100% of the time.

And so this is only 100 X bigger than the current game running 50 X bigger than the current game. Everyone was very excited about off go and at the time in 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 three and 61 possible actions. So that gives us about two samples of replication for every single action.

So here you need like 2 million of them. For a similar depth and then that's still to do a depth of 30. I would still have to do this times 30 to 60 million invocations of the model. That better be a small model.

So yeah, single action to be clear. Yeah, so exactly to do one action. So just imagine. So why off go doesn't scale?

No. To me, there's one. The carnality of the action space must be a shimmy small. If it's big, sad.

Two. I need a perfect deterministic environment. Right? Like 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 almost good. And if you saw the movie, the documentary is such an amazing documentary. I highly recommend it to anyone that watches it. The guy is sitting there for like 60 seconds, maybe 5 minutes waiting for the computer to like decide.

And it's kind of like imagine we were driving a car and like you look like 60 seconds to like turn the steering wheel. Everyone's dead. The whole car is dead. And so like, you know, now let's talk about robotics and stuff driving car and why this why that approach kind of can scale.

Yeah, I think a really good contrast here because intuitively, I think in thinking through this exact star layout, it actually really is how I think about the kind of problem space about the computer. So like 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 the way well and that's left is the only thing that I work. So like let me 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. I mean, in a sense, the chess in AlphaGo space is already like more than a number of atoms in the universe or something. But like just to emphasize that here, you know, you are considering, you know, surroundings, vehicle state, like, you know, camera, like weather. I guess the point is like, it's like massive.

It's like massive. It's all from the purpose of it. It's important to us. Correct.

And so is the space of pixels like us. You know, like what can I put in an image? I can put in an image of anything. Yes, true.

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 can press it with a common that because they have some deep, some deep common that actually takes this thing and converts it into a latent.

Right. And that latent compression is sufficient to kind of like do pattern matching, do some type of like symmetric, symmetric, uh, 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, uh, that works extremely well, which reduces our state space dramatically because I'm in some latent, high dimensional space. So like the key thing there is that, yeah, despite this space being effectively infinite, we've actually done really good at compressing this.

Yeah. 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. Exactly. Exactly.

T seems to have a similar problem as before. In fact, maybe even more extreme. There's like infinity other variables around you. Right.

In some ways you'd think that it's, this is physics. This is laws of Russian should apply. If I turn this thing, we're like this. 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 we face in, when I was doing a self-driving car is like you were imposing your will onto, maybe driving in India. Exactly. You're imposing your will onto the environment. And like people just kind of adapt naturally.

Like if you were doing, there's a lot of motion, you were going to collide. And so 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. But it's not true. Like that, then we wouldn't function.

The cars wouldn't go down the road. And so you have to model the environment. You have to include other people in the environment and understand the embodiment of like how your action will change other people's actions. YC's next batch is now taking applications.

Got a startup in you. Apply at ycombinator.com.com. It's never too early and filling out the app will level up your idea. Okay.

Back to the video. 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. It seems like, well, you know, you turn this feeling, we're left to right.

You hit the break, you hit the gas. It doesn't seem that big. But like how big is it? Like how do we actually represent these action spaces?

Like a realistic self-driving car scenario? Yeah. I don't know how they do this nowadays. They're doing a whole bunch of like bird's-eye view, different things like that.

Things are even just like a very simplified case. 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.

And so this thing is like 365 degrees. Yeah. So it's like a one to three 65, let's say zero to three 65. Yep.

And let's just say you break this up into 10 different, uh, uh, severities. You're already a, even with this over simplified model, your action space cardinality, right? It's 365,000. So that's like a hundred X bigger than awful.

It's in fact, it's about the size of the example. Or in fact, like decent amount smaller than the size we said. Exactly. Yeah.

And so yeah, so 36,000 action spaces very large and even worse, unless you're Tesla, we have a bunch of video of people driving cars. We don't have video of like dash Gams like that. You actually don't have, again, only Tesla has this of the action as well. So the things that you have access to the trajectories are just like ST, ST plus one, ST plus two.

So there's a decent number of these that's like dash cam footage on YouTube or something. Not really that many either. Yeah. Relatively.

So if you wanted to do a subject car and you didn't want to go spend a million dollars, $20 on going collecting all this data, then you want to leverage this data somehow. And this is going to be really applicable for, uh, robotics because we have a lot of, uh, videos of people doing things, right? Especially with egocentric. Like we have those videos, but we don't have is.

The actions they take. Yeah. So this is like, this is a sequence. Yeah.

Unless Tesla has this. So this is a huge competitive mode of like, what do people do in that state? And 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.

It's not sufficient. People think that like, okay, I have this. We have such a 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 wing level. Um, yeah. Would this be a good moment to briefly talk about model three versus model based RL?

Yeah. I think that's an important distinction that's going to be relevant. We talked about world models. Yeah.

So this is a perfect point. So model free just means that my policy pie of a T given ST, 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 a.

Just put it the next day. And that's it. And this is what I just called the LA. You know, this is like, giving us pretty good results.

It's behavior cloning. It's all the stuff that is not getting us the rosy robot just yet. But in many ways, it's the closest thing that just looks like the next token prediction from LMS that seems to scale pretty well with natural language. Exactly.

But it's very analogous to that. Yeah. 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 the subject card does in Tesla. And I just taken 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. They don't have a model for the environment. And then now if I do model based RL, I have not just some pie, but I have also by size while here.

So by 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 planning. Just remind us that size referring to this specific transition function, right? It's referring to this using this is specifically referring to a function of ST plus one given ST and action T. Yes.

Because like your ability to predict the next state you'll be in is the core of it. Yep. As opposed to just directly predicting reactions. Yeah.

And the main thing that I believe is that this is required for AGI. This is what the human brain is. At least the way the human brain does it. And let me go further and say that like if you look at the billions of years of evolution, basically there's this thing called 10 million years ago called the Great Chortical Expansion, which you see the size of a brain just explode 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 and fish. Yeah, right? That's like very like lizard brain, whatever.

And then we developed this neocortex to like, you know, go from our motor cortex to actually simulate what's going to happen. And that makes us just so much smarter. And then once we get those samples, we can compress it when we sleep or otherwise with this hippocampal, short-river, whatever you want to call it. And then that helps us develop a better policy.

And that marriage between the two is not only helps us train on hallucinated examples, but it also allows us to test time plan. I guess the kind of extreme case that self-driving car is general robotics. Yes. Right.

So if you're like a humanly company, like, bigger or pie or whatever, again, same STAR setup. I guess the gist of it is that A is now even bigger. Yeah. Right.

So if I take like my six axis arm as your standard year that we're actually working on right now in the robotic center, you have two degrees of freedom, two degrees of freedom, two degrees of freedom. And then you have an O2 for the end effector. Right. And so that's a simple effect.

They're not even like a... It's literally one axis. Like, you know, you can rotate, but you have the one axis, you mean style thing. So this is eight, so you have 16 degrees of freedom.

And let's just say that you do the 365 to whatever, you know, kind of thing. It's like 10 to the 16. It's like 10 to the 16. 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, just get some Philippines. We'll give them like some, you know, things or whatever.

It's like totally, totally doesn't work. 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 thing is that for Tesla. Right. 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 3, it wouldn't work. No. Like it totally wouldn't work. Like all the so much of this, the way that if I were to break on a Model 3 versus a Model X, a Model X, it weighs more, it has different dynamics and dynamics and like that.

And so what's actually going to happen is very different. Like the degradation you have across across and body mens is very, very strong. And clearly Tesla's figured various ways to get around that. I mean, they have these that roll up and actually even with Tesla's new FSP 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 a body mens than a Model 3 versus Y. Yeah. And you have way bigger action spaces.

You have some on that one. And then I play a talk with it as they ever now runs Tesla FSD. I can ask him, but I would bet money that they shard the data per Model, per car type. Yeah.

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

So now that we understand the basic setup here and why the actual space problem is so big, why 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 new world models and how do we fix those and then let's talk about some of the newest world modeling techniques. Cool. So like in robotics in particular, it's very hard to get these, this kind of trajectories that you use your VLAs and people spend up, you know, with a whole bunch of teleops data, it's a very expensive idea. Ideally, what we would do is take like data like this from someone who is just like, put the camera on them and just like making sushi.

Okay. Like I want to make a sushi robot. 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. 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 later on this afterwards. Yes. And so the first real person that you know, when after this was, you're going to smith you were at, please, so he doesn't yell at us.

We have to make sure we say that. make sure we say that. But he has a 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 way.

And then trains a policy on only the, the synthetic data, the imagination of rollouts. And it actually performs well in the environment. It's the first time in my understanding that that actually happened and it actually works really well. And then so the key thing there, 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 in on a real deal here.

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. And so now the Dreamer series, basically this published publishes in May of 2018. Danjaer Haffner publishes Dreamer one, I think in November, and then now he's been on this rampage for the last seven years publishing these papers.

And Dreamer before, I think, is the capstone of it, where he basically does the same thing and he focuses on Minecraft. And he trains a world model on this type of data, and then injects action conditioning on a very small amount of data to get to this type of world model that has the action conditioning as well. And then samples a lot from it, and then trains a policy on those synthetic, imaginative 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 in synthetic data, which is crazy. And the key on that there, yeah, uses synthetic data specifically on a model train on just this sort of state transition type of thing. And this ends up being very convenient because it turns out we as a society have a lot of this.

Exactly. Yeah, all of YouTube, right? He does do a very small amount of data from to enable the action conditioning, and that allows you to do this full simulated rollout. But yeah, it's true.

So we have YouTube, we have Flickr, we have all these data sets online of people doing things, we like to use it. And no one has really gotten that to work. And now with this 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 tasks that we wanted to do, chopping down wood or making sushi or folding my bed, or whatever it is, only a few amount of examples. And then we can train a policy on this, in this neural simulation.

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

Yeah, so like if you have your C-dance or your Sora or one, 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 world model and then we can sample from it a bunch. And then this is exactly what we've did with Gaia. And Gaia, I think they've raised $1.5 billion to basically run with SIDI for such an hour.

I think a bunch of companies, NVIDIA, this paper here is basically talking about doing exactly the same, the stream zero for robotics. And what I thought was really cool about this paper is that the exact thing was processed where they have this joint model of 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 tele-out data which is basically exactly this.

To get it to be pretty good and they have a lot of clever tricks that allowed it to be cross embodiment and working on scene tasks with relatively small amounts of data. And it really is taking basically the exact concept I believe from the Dreamer paper and applying it specifically to these robot embodiment. Exactly. And it turns out it actually works actually better than I would anticipate it.

Yeah, so I think this is basically the path to, it was the path, I believe that it'll path to get humans to be as good as we are genetically over the last 10, 20 million years of evolution. A bigger world model helps for training and for test time planning. And I think it'll be the same thing as true as for robotics. What's also cool is there's a bunch of applications of this sort of things outside of robotics too.

There was a weather planning paper, for example, with Gencast paper which I think applies a relatively similar concept in terms of how they model literally the world, the world, the world, the world, the world, the weather, or something like this. Yeah, we have to talk about the world model for the world. Yeah, so basically this exact same thing where the key unlocks for this whole thing was getting diffusion to work in very high dimensional states like we talked about in the last lecture, and then learning to use that to action condition in the way that he's done. But they did this for the entire world with this exact same diffusion steps, where to go from some, and they go back to two time steps, lag of order two, AR two for the statisticians there, and then basically predict the next state of the world based on those things with this ling of diffusion rollouts.

My biggest question is that 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 that was so great at a university of Washington where they say explicitly 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. Yeah, right.

It's actually describing exactly these two regions here. Exactly, where we're estimating both the sensory latent states and actions, and I guess it's really the joint model that we showed earlier. It's exactly this equation. Yeah, exactly.

And so if it works in us, it should work in robotics, and I think that that takes us the rest of the distance. Why don't we talk briefly about latent world models, especially the Java concept, because I think there's been a number of papers that use Java as an element of their architecture or was briefly introduced by Java and how it fits into the current landscape of world modeling. Yeah, in classic RL, if you do study Q-line, for example, you basically keep this matrix called the Q-matrix, and it's going to be S by A. And so I have this S by A, state and actions, and each one, I need some amount of counts being in this state action, and I take the average value of taking that action in this state, and that's my Q value there.

And it's a little bit more complicated in that. There's a little bit more complicated in this backup, all this stuff like that. So this scales horribly because as the cardinality in my space space gets bigger and actually it's bigger, stuff I don't have enough to, I become less and less simplification. And the classic thing is like, real, less or whatever, it's like, yeah, it's this whole thing we described earlier, it's absolutely massive, but that's all of these elements in it, it's been really enumerated, huge grip.

And so the classic trick, I mean, since I took, C-229 with Andrew Rong in 2012, is you do this. Pick a neural network on it. Exactly, and you basically are just going to compress that state into some lower dimensional state space. This actually predates deep learning.

We were doing stuff like this. I think my first paper was basically doing something like this, basically turning like a grid into like a bunch of like pyramids. And like, and the state was how much I'm in pyramid one or pyramid two or whatever. But anyway, the neural network can just do this.

And so basically what the key idea in JPAI, if I have an image one and I have image two, I have image three, I can do my world modeling, my world modeling of S, D plus one, given ST and AT in pixel space and have this is, let's say at time T, T plus one, T plus two, et cetera, and I have to actually predict now the full 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 comment. Some encoder.

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 my goal is to make this and this, make my loss function will be something very simple, like on a minute. And then I'll have this. Now this doesn't work, this collapses hard.

And so what happens, if you just predict zero, that's zero. So 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 SIGRAG, or this is one technique of SIGRAG, is another.

Where basically I add this another term that basically says, I want the over a large enough batch size. I want the distribution of Z, T plus one, to follow a Gaussian. It's kind of like a normalize it, like a batch norm type of a correct. And if it's zero, it can't be this.

Right? Because then this is non-zero. And so maybe I think there's probably this or something like that. But basically this prevents it from local apps.

And it makes it do something good. And the most recent paper for the audience is LEWM. LEWM, LEWM model is super, super great. However, to be completely frank, this is self-supervised learning.

Super great. It doesn't work that well. If you were to not do these techniques, and there's a bunch of other techniques that you can do, it will actually outperform much better. Or let's say for example, if I'm going to do an LOM, and you have like, you know, Francois likes sushi, which is definitely true.

And I tokenize this into a bunch of different tokens here. And this is token ID 6, 19, 28, whatever. And I look up the encoding into this. And that's going to be E1, E2, E3, et cetera.

What you can actually do is have the LOM output, with the LOM, we'll take in these things, and we'll output the next token. And so it would be like, let's call it H. This would be the load jets coming out of it, plus one. And what you can do is actually have this be close to E, T plus 1.

And a lot of people are playing with this idea, and getting rid of the cross entropy loss 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 cross entropy head is actually very expensive. And so this is very cheap.

And this is like just grabbing it. So people are playing around with this idea, and as a, basically, as a cheaper proxy for the cross entropy loss, so there's lots of different ideas on basically taking this JPA idea to not just pixels, but to LMs as well. Yeah, interesting. So just to find what JPA is, it's joint embedding predictive architecture.

I think one of the things I find cool about this JPA 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 staple diffusion idea. There's a version of this idea that in my company training, Graph Convolutional Neural Networks to design drugs. We use to do latent variable generation, for example.

And it's like an idea that comes back over and over, and then has this various tricks that it actually takes to get it to work. OK, now we have a pretty good sense of how world models work. We have a pretty good sense of what the state of the art looks like. If we trust this paper, it seems like it's kind of work on robots, too.

The paper's only from the 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-trained on fusion models. So are we good? Does it all work?

Yeah, it's 26. It's 26 with a year of the robot. We're going to have a roset robot in your house. Yeah, no.

What are one or two? There's lots of open problems remaining. What are a few open problems? Maybe we can emphasize here that the community can go emphasize what we're doing.

Yeah, so I think the first one is that pins doesn't really work. What does pins? Physics and form neural networks. So pins doesn't really work.

There's physics and form neural networks. And so basically, if 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, I have a house here. And I want to train the model on not driving into the house.

And so let's say 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, trying down the road. This will just turn magically into a highway. And you're just like, ooh, don't worry.

You might be crash off. It basically means like a ton of data not to do that. Different simulation for that to not happen. 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 when you're training on network, it has a tendency to collapse if you don't keep the mini-batch composition, like very even over the class base or whatever you want to call it. But you have to train on, you have to be very careful about your data mixing to make sure you get this right to solve this problem that no one really has. But even then, if you take just a simple thing like this, this is like the conic example.

And I have some sine wave. And I want, and I have these as my x, bop, bop, bop, bop. And I have these as my y. So this complete interpolation.

No, that's maybe messed this up, but why? Like this. No. We can't get to like machine precision.

What is it? What is it? What is it, what is it? One of the 16 or whatever it is.

We can't, the SGD will not get to 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 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, I saw this one video of Steph Curry tripling about basketball on a court. And he just felt that there was a dead spot in the court. And he, because he's so good.

And he knows exactly the physics of what's going to happen. If I hit this, you know, the vault with this force, like the vault was going to come back exactly the spot. And it just didn't, and he knew it wasn't him. It was the court.

And he found a dead spot in the court. Like that's how good 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 in and just said that he thinks that there's definitely architecture that's going to be more performing than the transformer. I think he's right. I think the transformer doesn't do compression in the time domain at all. It just keeps right on everything.