la leWorldModel
Simultaneous latent action and world model learning from passive video for controllable prediction.
Interactive Demo
Try the model in your browser.
Step the latent dynamics yourself and watch the decoder render the rollout live.
I began this project with three questions:
Can a world model become controllable without ever seeing environment actions?
Can we recover those actions directly from passive video?
Can we do it all in one training run, without a clunky multi-stage pipeline?
leWorldModel is already a very nice idea, but it still assumes you have access to the environment actions. That is fine if you live inside a benchmark, but much less fine if you want to train on the actual internet, gameplay footage, robotics videos, or whatever other suspicious collection of pixels you have lying around.
I call the resulting model la leWorldModel, short for latent action leWorldModel.
This post is about a very simple extension: instead of conditioning on environment actions, we learn a discrete latent action space directly from video and train that jointly with the world model. The result is a controllable world model trained from passive video alone.
The fun part is that this is not a separate latent-action-pretrain-then-world-model recipe. The latent actions and the world model are learned simultaneously, end to end, in one run. Here I focus on the PushT results.
The Problem
Most controllable world models quietly rely on one very expensive luxury: labelled actions. If your data comes from a simulator or an instrumented environment, that is easy. If your data is just video, you suddenly have dynamics but no control signal.
That is the bottleneck. Video is abundant. Action-labelled video is not. So if we want world models that scale beyond neat little RL datasets, they need a way to infer control from observation alone.
The Idea
The core trick is very simple. We replace the observed action with a learned latent action.
- An inverse-dynamics-style module looks at neighbouring observations and predicts a discrete latent action.
- A VQ-style codebook forces those actions into a compact reusable set.
- The world model conditions on those learned actions instead of environment-provided ones.
- Everything is trained jointly, so the action space is shaped directly by what helps prediction and control.
In other words: rather than asking the dataset what action happened, we ask the model to invent an action vocabulary that makes the video controllable.
Why This Is Different
Compared with leWorldModel
leWorldModel assumes the action is available and then learns to predict the future under that action. Here the action is no longer given. It has to be discovered.
Compared with Genie-style Pipelines
Genie introduces the term LAM, for latent action modelling, and the basic philosophy is close to what I want here: recover actions from passive video rather than asking the environment to hand them to you. The main difference is where that action model lives and how it is trained.
In Genie, the LAM is trained in pixel space. Effectively, it behaves like a VQ-VAE: an encoder takes \(x_{1:t}, x_{t+1}\) and produces a discrete action \(a_t\), then a decoder takes \(x_{1:t}, a_t\) and reconstructs \(x_{t+1}\). The main world model is then trained separately, using a different VAE and the predicted latent actions. It works, but it is naturally a staged pipeline.
In shorthand, the latent action model in Genie looks like this:
Here I avoid that split. The latent action model operates directly in the learned latent space of the world model, which is I think the key reason joint training becomes natural rather than awkward.
leWorldModel
Needs environment actions.
Genie-style
Video-only, but often multi-stage.
la leWorldModel
Video-only and jointly learns latent actions plus the world model in one run.
How It Works
The Objective
At a high level, the model observes video, compresses it into a learned latent space, predicts a discrete latent action between frames, and uses that action to drive future prediction. The action codebook is trained at the same time as the world model, so the two parts co-adapt.
Concretely, I use an inverse dynamics model that peeks one step ahead. It takes \(z_{1:t}, z_{t+1}\), where the \(z\)'s are latents produced by the encoder rather than raw pixels, and outputs a discrete latent action \(a_t\). That action is then what the world model conditions on.
In my setup, the same idea moves into latent space:
This is similar in spirit to Genie’s LAM, but importantly it is not trained with its own pixel reconstruction objective. There is no separate decoder trying to rebuild the next frame from pixels. Instead, the latent action model is trained entirely through the overall world-model objective. In other words, the action space only has permission to exist insofar as it helps prediction.
I think this latent-space formulation is also why the whole thing is so stable. Jointly learning actions and dynamics would be much more brittle in pixel space, where the model is constantly entangled with reconstruction details. In latent space the optimisation problem is cleaner, and in practice the dynamics model turns out to be easier to train than the visual decoder.
I still use a small codebook, with 8 primitives, to stop the latent actions from quietly becoming arbitrary information carriers. That bottleneck is a good inductive bias here: if you want something that behaves like character control, you should force the model to express itself through a small reusable action vocabulary rather than a secret side-channel.
Written out more explicitly, the training objective is:
In implementation terms:
| Term | Role |
|---|---|
| \(\mathcal{L}_{\mathrm{pred}}\) | The main world-model objective. It makes the action-conditioned predictor match the target next embedding. |
| \(w_{\mathrm{sigreg}} \, \mathcal{L}_{\mathrm{sigreg}}\) | A regulariser on the learned embeddings that prevents trivial collapse. |
| \(\mathcal{L}_{\mathrm{codebook}}^{*}\) | Moves the codebook toward the inverse-dynamics outputs. |
| \(\mathcal{L}_{\mathrm{commit}}^{*}\) | Moves the inverse-dynamics outputs toward the selected codebook entry. |
* Quantisation is kept differentiable with the straight-through approximation.
Visual Decoder
Because the model works directly in embedding space, unlike Genie, there is no free visual simulator attached to the latent action model. That is part of what makes the setup elegant: the dynamics model only has to represent what matters for dynamics, rather than wasting capacity on pixel reconstruction.
The price is that if I want to actually visualise imagined rollouts, I need to train a separate visual decoder afterwards to map latent states back into frames. Fortunately this part is not too hard once the latent space is already structured.
For that decoder, I use a reconstruction loss of the form:
The top-k term is the useful trick here. I originally used that idea in my older post on sample efficient AR diffusion, and it works well again here: it lets the decoder focus on the hardest regions rather than averaging all reconstruction errors into harmless blur. This is especially helpful in environments with very unbalanced pixels, for example a small moving object on a plain coloured background, where a reconstruction loss can otherwise be dominated by the already-correct easy pixels.
Action Decoder for PushT
For PushT I also trained a simple decoder that translates learned latent actions back into the environment's accepted actions. This makes it possible to plan in latent action space, then execute those plans in the actual environment.
This translator is needed because PushT is a continuous-control environment, whereas my method deliberately forces control through a discrete latent action space. So even if the world model learns a good latent action vocabulary, there is no obvious direct mapping back to the real-valued action that PushT expects.
There is also a second complication. Following the original leWorldModel setup, each state is spaced 5 frames apart. So a single latent action is effectively standing in for 5 frames of continuous control, during which both direction and magnitude can change. The translator therefore has to infer a plausible continuous action from the current state together with the chosen latent action, rather than merely looking up a one-to-one correspondence.
Experiments
PushT
PushT is a nice testbed because controllability matters immediately. It is not enough to merely model the video distribution; the learned actions need to support actual useful behaviour.
Here I train on video, infer latent actions, learn the world model jointly, and then decode latent actions back into environment actions for control and planning.
Training
These are the training curves from the PushT run.
The nice part is that learning latent actions does not seem to disrupt training dynamics much at all. These curves look very similar to the ones in the supplementary material of leWorldModel, which suggests the latent-action extension slots into the original optimisation story fairly cleanly.
Planning
Under the PushT planning challenge, following the evaluation style of the original leWorldModel paper, success hovered around 30% over 50 trials. That is much lower than the number reported in the paper, and my guess is that the extra latent-to-real-action translation layer is the main reason. Even so, the qualitative behaviour is still fairly reasonable: the model often does the right kind of thing, even when it fails to complete the task reliably enough to score as a full success.
In the above videos, the imagined trajectories are first found inside the latent world model by CEM search over the learned latent action primitives, and only then rendered into frames by the visual decoder. The enacted trajectory is what happens when those latent decisions are translated back into real environment actions and executed in PushT. The ideal trajectory is the ground-truth rollout that reaches the shown final T position.
Comparison
For reference, here is the fixed-FLOPs comparison against the baselines.
la leWorldModel comes in well below the fully action-conditioned leWorldModel baseline, but I do not think this is the core method failing. In this setup the weak point is the discretisation-and-translation stack: PushT is a continuous-control problem, yet the planner operates over a small discrete latent action vocabulary, and each latent action has to stand in for five frames of underlying control rather than a single clean environment step. Those latent choices then still have to be translated back into executable continuous actions. That makes this comparison harsher than the original leWorldModel setting, where the true actions are already available.
What Actually Matters
The point is not just that the model can simulate video. The point is that the learned latent actions become operationally useful for control and planning.
- The model can be trained on passive video rather than action-labelled trajectories.
- The latent actions are learned jointly with the dynamics rather than in a separate pipeline stage.
- The latent dynamics model and latent action model train in a surprisingly stable way; in practice, the visual decoder is harder to train than the dynamics.
- The resulting latent actions are useful enough to support real planning and execution in PushT.
Limitations
The latent actions are useful, but they are not guaranteed to line up with human-interpretable actions. Their semantics are whatever the model finds convenient for prediction and control. That is probably the right inductive bias, but it does mean the action vocabulary may be strange.
The other obvious limitation is coverage. If the video data does not contain enough variation, the model cannot recover control modes that never appear. Passive video is abundant, but it still only teaches the possibilities it actually shows.
Takeaway
leWorldModel becomes much more interesting once you remove the requirement for environment actions. A small latent-action extension is enough to train on arbitrary video and recover a controllable world model.
The main thing I like about this result is its simplicity. It feels like the natural generalisation: if actions are the missing ingredient in passive video, then learn them. Better yet, learn them jointly with the world model so the whole system has to agree on what an action is for.