From Diffusion Policies to World-Action Models: Teaching Robots to Predict, Plan, and Act

DISCLAIMER: THIS IS NOT SUPER BEGINNER FRIENDLY

These are a series of mental notes on generative robot policies, joint-embedding predictive architectures, latent-space planning, and the emerging connection between world models and control . More iterations of these would include even more agentic methods that are getting super popular

Robotics sits at an interesting intersection of perception, prediction, and decision-making.

A robot observes the world, proposes an action, and then has to deal with what actually happens. Its next observation may differ from everything it has seen during training. Objects move unexpectedly, contacts are imprecise, and small errors can accumulate over a long sequence of actions.

This makes robot learning more than a problem of mapping observations to actions. It is also a problem of learning how the world responds to those actions.

Recently, I've been exploring several ideas that approach this problem from different directions:

• Diffusion and flow-matching policies: generating structured action sequences.
• Multistep Inverse Problems (MIP): refining actions through iterative computation.
• Joint-Embedding Predictive Architectures (JEPAs): predicting in learned representation spaces rather than reconstructing pixels.
• World-Action Models (WAMs): combining predictive world representations with action generation.
• Differentiable model-predictive control: optimizing candidate actions using gradients through a learned model.

The interesting part is how these ideas fit together.

1. The problem with simply imitating demonstrations

A conventional imitation-learning policy learns a mapping from observations to actions:

$$a_t={\pi }_{\theta }(o_t)a\_t\; =\; \backslash pi\_\backslash theta(o\_t)$$

Here, $o_t$ is the observation at time $t$, and $a_t$ is the predicted action.

In practice, many robot policies predict an action chunk rather than a single action. Given the current observation, the policy produces a sequence of future controls:

$$A_t=(a_t,a_{t+1},\dots ,a_{t+H-1})A\_t\; =\; (a\_t,a\_\{t+1\},\backslash ldots,a\_\{t+H-1\})$$

Action chunking can improve temporal consistency and reduce the frequency of policy inference.

However, behavioral cloning has a familiar weakness: the policy is trained on demonstrations, but deployed on trajectories partly determined by its own mistakes.

An action that is slightly wrong can move the robot into a state that was rare or absent in the training data. Subsequent predictions then have to operate outside the distribution the policy learned to imitate.

This is one reason generative policies have attracted attention. Instead of predicting one action sequence through a direct regression, they learn a distribution over plausible action sequences.

2. Why generative action policies work

Diffusion policies generate actions by iteratively transforming noisy samples into structured predictions. Flow-matching policies learn a vector field that transports samples toward the target action distribution.

For a simple linear flow-matching construction, let $A$ be a demonstrated action chunk and let $\epsilon$ be Gaussian noise. An intermediate sample is

$$A(s)=(1-s)ϵ+sA,ϵ\sim \mathcal{N}(0,I),s\sim \mathcal{U}([0,1])\mathrm{.}A(s)\; =\; (1-s)\backslash epsilon\; +\; sA,\; \backslash qquad\; \backslash epsilon\; \backslash sim\; \backslash mathcal\{N\}(0,I),\; \backslash qquad\; s\; \backslash sim\; \backslash mathcal\{U\}([0,1]).$$

The policy learns a conditional velocity field using

$${\mathcal{L}}_{\mathrm{f}\mathrm{l}\mathrm{o}\mathrm{w}}(\theta )={\mathbb{E}}_{(o,A),ϵ,s}\left[{\parallel v_{\theta }(A(s),s\mid o)-(A-ϵ)\parallel }_2^2\right]\mathrm{.}\backslash mathcal\{L\}\_\{\backslash mathrm\{flow\}\}(\backslash theta)\; =\; \backslash mathbb\{E\}\_\{(o,A),\backslash ,\backslash epsilon,\backslash ,s\}\; \backslash left[\; \backslash left\backslash |\; v\_\backslash theta\backslash bigl(A(s),s\backslash mid\; o\backslash bigr)\; -\; (A-\backslash epsilon)\; \backslash right\backslash |\_2^2\; \backslash right].$$

At inference time, it generates an action chunk by integrating

$$\frac{d{X}_s}{ds}=v_{\theta }(X_s,s\mid o),X_0\sim \mathcal{N}(0,I),\widehat{A}=X_1\mathrm{.}\backslash frac\{dX\_s\}\{ds\}\; =\; v\_\backslash theta(X\_s,s\backslash mid\; o),\; \backslash qquad\; X\_0\; \backslash sim\; \backslash mathcal\{N\}(0,I),\; \backslash qquad\; \backslash hat\; A\; =\; X\_1.$$

Both approaches can model multimodal behavior: when several actions are plausible, the policy does not necessarily have to average them into a single, potentially undesirable action.

But an important question remains:

Does the improvement come from the generative formulation itself, or from the iterative computation that comes with it?

The paper Much Ado About Noising: Dispelling the Myths of Generative Robotic Control investigates this question by comparing generative and regression-based approaches under controlled conditions.

One particularly interesting takeaway is that the advantage of diffusion or flow-based policies is not universal. Under matched architectures and training setups, simpler regression approaches can perform competitively, while iterative computation and stochasticity can matter substantially in particular settings.

This shifts the question from whether generative policies are inherently superior to a more useful one:

Which properties of the learning and inference procedure actually improve control?

That is where iterative refinement becomes interesting.

3. MIP: can a policy improve its own prediction?

Minimal Iterative Policy (MIP) isolates supervised iterative computation and training-time stochasticity without requiring a conventional generative sampling loop.

Using action-chunk notation, its two deterministic inference steps can be written as

$${\widehat{A}}_t^{(0)}=f_{\theta }(o_t,0,0),\backslash hat\; A\_t^\{(0)\}\; =\; f\_\backslash theta(o\_t,0,0),$$

$${\widehat{A}}_t^{(1)}=f_{\theta }(o_t,s^{\ast }{\widehat{A}}_t^{(0)},s^{\ast }),s^{\ast }=\mathrm{0.9.}\backslash hat\; A\_t^\{(1)\}\; =\; f\_\backslash theta\backslash left(o\_t,\backslash ,s^*\backslash hat\; A\_t^\{(0)\},\backslash ,s^*\backslash right),\; \backslash qquad\; s^*=0.9.$$

Here, $s^*=0.9$ in the paper’s principal formulation. The superscripts label computational passes, not environment timesteps.

The broader idea is worth considering: iterative computation can be a resource for improving action predictions, even without running a conventional diffusion sampling loop.

This suggests a useful design question for robotics:

Can we obtain some of the benefits of iterative generative inference through a simpler refinement mechanism?

The answer depends on the task, the training procedure, and the policy architecture. It should be measured rather than assumed.

4. From predicting actions to predicting the world

Action generation alone does not tell us whether an action will produce a desirable outcome.

A policy can learn which actions tend to accompany particular observations without explicitly learning a predictive representation of the consequences of those actions.

World models approach the problem differently. They learn representations of the environment and how those representations evolve over time.

Let the current observation be encoded into a latent state:

$$z_t=E(o_t)\mathrm{.}z\_t\; =\; E(o\_t).$$

A learned dynamics model predicts a future latent representation conditioned on the current state and an action:

$${\widehat{z}}_{t+1}=F_{\varphi }(z_t,a_t)\mathrm{.}\backslash hat\; z\_\{t+1\}\; =\; F\_\backslash phi(z\_t,a\_t).$$

This notation assumes that 𝑧 𝑡 z t ​

contains enough information for predicting the next state. When a single observation is insufficient, the encoder or predictor can instead condition on a history of observations and past actions.

Rather than reconstructing every pixel of a future image, the model predicts a representation that should preserve information useful for understanding the future.

This distinction matters. Pixel-level prediction can spend substantial modeling capacity on appearance details that are irrelevant to a particular control objective. A good latent representation can instead prioritize structure that matters for predicting what happens next.

Joint-Embedding Predictive Architectures, or JEPAs, provide one influential approach to representation learning through prediction in latent space.

The key idea is not simply to compress an image. It is to learn representations in which predicting another state becomes a useful learning signal.

For robotics, this creates a natural opportunity: learn representations of the world, learn how actions change those representations, and use the resulting model to make better decisions.

5. LeWAM: combining world prediction and action generation

One paper I found especially interesting is LeWAM: End-to-End World and Action Modeling with JEPAs.

LeWAM brings together several objectives that are often considered separately:

1. Learning representations of observations.
2. Predicting future representations.
3. Generating actions.
4. Regularizing the learned representation space.

Its approach combines latent prediction with flow-matching-based action modeling and SIGReg, a regularization objective for the learned representations.

Conceptually, the system has to learn both how to represent the world and how to generate actions in that world.

A simplified objective illustrates the idea:

$$\mathcal{L}={\lambda }_{\mathrm{p}\mathrm{r}\mathrm{e}\mathrm{d}}{\mathcal{L}}_{\mathrm{p}\mathrm{r}\mathrm{e}\mathrm{d}}+{\lambda }_{\mathrm{a}\mathrm{c}\mathrm{t}\mathrm{i}\mathrm{o}\mathrm{n}}{\mathcal{L}}_{\mathrm{f}\mathrm{l}\mathrm{o}\mathrm{w}}+{\lambda }_{\mathrm{r}\mathrm{e}\mathrm{g}}{\mathcal{L}}_{\mathrm{S}\mathrm{I}\mathrm{G}\mathrm{R}\mathrm{e}\mathrm{g}}\mathrm{.}\backslash mathcal\{L\}\; =\; \backslash lambda\_\{\backslash mathrm\{pred\}\}\backslash mathcal\{L\}\_\{\backslash mathrm\{pred\}\}\; +\; \backslash lambda\_\{\backslash mathrm\{action\}\}\backslash mathcal\{L\}\_\{\backslash mathrm\{flow\}\}\; +\; \backslash lambda\_\{\backslash mathrm\{reg\}\}\backslash mathcal\{L\}\_\{\backslash mathrm\{SIGReg\}\}.$$

In the state-prediction branch, $K$ denotes the prediction offset. The dynamics predictor conditions on observation-history latents and the intervening clean actions. This is a compact schematic, not a complete specification of the architecture.

The important design principle is the joint learning problem: the representation is trained to support prediction, while the policy learns to produce useful actions.

Why might this help?

A representation optimized only for visual similarity may not be the best representation for control. Conversely, a policy trained only to imitate actions may fail to capture the structure of the environment that determines whether those actions will work.

Jointly learning predictive representations and action generation provides a way to align these two requirements.

The paper reports improvements on its evaluated tasks, but those results should not be interpreted as evidence that the architecture will outperform every alternative across robotics benchmarks.

The interesting contribution is the approach to combining world modeling and action generation in a single learning framework.

6. Attention masks are part of the model, not just an implementation detail

A subtle but important issue in joint world-and-action modeling is information flow during training.

Suppose a model receives historical observations, clean action targets, state queries, and action-generation queries. If the attention pattern allows a query to access information it would not have at inference time, the training objective can become misleading.

The model may learn to exploit access to target information instead of learning the intended prediction.

LeWAM uses a structured attention design associated with its Mixture-of-Transformers architecture. Different token groups have different information-access patterns.

This is not merely a computational optimization. The mask determines which information is available to each prediction.

The general principle applies well beyond robotics:

A model's attention pattern defines part of its learning problem.

When building predictive or generative systems, it is worth asking:

• Which tokens represent observed information?
• Which tokens represent targets?
• Which predictions are allowed to condition on which inputs?
• Does training-time information access match inference-time availability?

A carefully designed attention mask can help enforce the distinction between conditioning information and the target being predicted.

7. Differentiable latent MPC: using the model to improve actions

Once we have a differentiable dynamics model, another possibility emerges: instead of asking the policy to produce the final action immediately, we can optimize candidate actions against predicted outcomes.

This connects world models to model-predictive control (MPC).

Suppose we have a current latent state $z_t$, a learned dynamics model $F_\phi$, and a sequence of candidate actions $A$.

We can roll the dynamics model forward:

At time $t$, initialize a latent rollout from the current encoded observation:

$${\widehat{z}}_{t\mid t}=E_{\psi }(o_t)\mathrm{.}\backslash hat\; z\_\{t\backslash mid\; t\}\; =\; E\_\backslash psi(o\_t).$$

For a candidate action sequence

$$A=(a_{t\mid t},\dots ,a_{t+H-1\mid t}),A\; =\; (a\_\{t\backslash mid\; t\},\backslash ldots,a\_\{t+H-1\backslash mid\; t\}),$$

predict its consequences recursively:

$${\widehat{z}}_{t+k+1\mid t}=F_{\varphi }({\widehat{z}}_{t+k\mid t},a_{t+k\mid t}),k=0,\dots ,H-1.\backslash hat\; z\_\{t+k+1\backslash mid\; t\}\; =\; F\_\backslash phi\backslash left(\backslash hat\; z\_\{t+k\backslash mid\; t\},a\_\{t+k\backslash mid\; t\}\backslash right),\; \backslash qquad\; k=0,\backslash ldots,H-1.$$

Score the predicted trajectory using

$$J_t(A)=\sum _{k=0}^{H-1}\mathrm{\ell }({\widehat{z}}_{t+k\mid t},a_{t+k\mid t})+{\mathrm{\ell }}_T({\widehat{z}}_{t+H\mid t},z_g),J\_t(A)\; =\; \backslash sum\_\{k=0\}^\{H-1\}\; \backslash ell\backslash left(\backslash hat\; z\_\{t+k\backslash mid\; t\},a\_\{t+k\backslash mid\; t\}\backslash right)\; +\; \backslash ell\_T\backslash left(\backslash hat\; z\_\{t+H\backslash mid\; t\},z\_g\backslash right),$$

where $z_g=E_\psi(o_g)$ is the representation of a goal observation.

Planning then solves

$$A_t^{\ast }inargmi{n}_{Ainmathcal{U}^H}{J}_t(A),A\_t^*\; \backslash \backslash in\; \backslash \backslash arg\backslash \backslash min\_\{A\backslash \backslash in\backslash \backslash mathcal\{U\}^H\}\; J\_t(A),$$

where $\mathcal{U}$ is the admissible action set.

A simple projected-gradient update is

$$A^{(j+1)}=P{i}_{mathcal{U}^H}left[A^{(j)}-et{a}_{j}nabl{a}_A{J}_{t}left(A^{(j)}right)right]\mathrm{.}A^\{(j+1)\}\; =\; \backslash \backslash Pi\_\{\backslash \backslash mathcal\{U\}^H\}\; \backslash \backslash left[\; A^\{(j)\}-\backslash \backslash eta\_j\backslash \backslash nabla\_A\; J\_t\backslash \backslash left(A^\{(j)\}\backslash \backslash right)\; \backslash \backslash right].$$

Here, $j$ indexes optimization iterations, $\eta_j$ is the step size, and $\Pi_{\mathcal{U}^H}$ projects the updated sequence onto the admissible action set. The encoder and dynamics parameters remain fixed: gradients pass through the model, but the optimizer updates only the candidate actions.

Execute the first optimized action—or a short prefix—then incorporate the next observation and replan.

If the learned dynamics and objective are differentiable, gradients can flow through the predicted trajectory and back to the action variables.

An important detail: we do not necessarily need to update the world model's parameters during planning.

We can freeze the learned model and optimize only the candidate actions. The model supplies gradients describing how changes in those actions affect the predicted objective.

This is different from training the model. It is inference-time optimization using a learned model.

It also introduces real challenges. Prediction errors can accumulate over long horizons, gradients can be poorly conditioned, and the optimizer may exploit inaccuracies in the learned dynamics. A low predicted cost does not guarantee a good real-world outcome.

For practical systems, planning therefore needs appropriate constraints, careful horizon selection, and feedback from the actual environment.

8. Why temporal context matters

A single observation can be insufficient to infer the state of a dynamic environment.

Velocity, contact, object motion, and the consequences of previous actions may be ambiguous from one frame. A history of observations and actions can provide information that a single image cannot.

This is one motivation behind longer-context world-action models.

NVIDIA's Long-WAM work explores the use of longer temporal context for world-action modeling. Its reported experiments show that context length can materially affect robot-task performance in the evaluated settings.

The broader lesson is that a world model's predictive capability depends not only on its architecture, but also on the information it receives.

Longer context is not automatically better: it introduces additional computational costs and may require an architecture and training setup that can make effective use of that information.

But when earlier observations reveal how the world is moving, what the robot has already attempted, or how an object responds to contact, discarding that history can be a significant limitation.

9. Putting the ideas together

These approaches can be viewed as different pieces of a larger control pipeline.

Represent. Encode observations into a latent state that preserves information relevant to prediction and control.

Propose. Generate candidate actions using a learned policy, potentially through flow matching or another generative procedure.

Predict. Use a learned world model to estimate how the environment will evolve under those actions.

Refine. Improve candidate actions through iterative inference or gradient-based optimization.

Act. Execute an action or action chunk in the real environment.

Observe. Incorporate the resulting state and repeat the process.

The resulting loop is:

$$\text{Observe}\to \text{Represent}\to \text{Propose}\to \text{Predict}\to \text{Refine}\to \text{Act}\backslash text\{Observe\}\; \backslash rightarrow\; \backslash text\{Represent\}\; \backslash rightarrow\; \backslash text\{Propose\}\; \backslash rightarrow\; \backslash text\{Predict\}\; \backslash rightarrow\; \backslash text\{Refine\}\; \backslash rightarrow\; \backslash text\{Act\}$$

This is a conceptual synthesis of the ideas discussed above, not a claim that all the cited systems implement this exact pipeline.

The pieces can be combined in different ways, and each introduces its own trade-offs. A generative policy may provide diverse action proposals; a latent dynamics model may estimate their consequences; and an optimizer may refine those proposals when the model and objective support reliable gradients.

The key research question becomes how to combine these capabilities without making the system too expensive, brittle, or difficult to train.

10. What I want to investigate next

A few questions stand out to me:

• Generative inference versus iterative refinement: Under matched conditions, when does a diffusion or flow-matching policy outperform a simpler policy with additional refinement steps?
• Representation quality for control: Which latent objectives produce representations that are useful for predicting action consequences, rather than merely encoding visual structure?
• Planning horizon versus model error: How quickly does the benefit of longer-horizon latent planning disappear as prediction uncertainty accumulates?
• Action proposal and optimization: Can a learned policy generate strong initial proposals while differentiable planning performs a small number of targeted corrections?
• Evaluation beyond success rate: How should we measure robustness to distribution shift, inference latency, action smoothness, and sensitivity to model error?

These questions need controlled experiments. Comparing architectures across different datasets, compute budgets, action spaces, and training procedures can make it difficult to identify what actually caused an improvement.

A useful experimental setup would keep the environment and evaluation protocol fixed, compare regression, flow-matching, and iterative-refinement baselines, and then introduce latent prediction and planning as separate components.

That would make it possible to measure not just whether a larger system works better, but which component contributes the improvement.

Closing thoughts

What I find compelling about this direction is the shift from treating robotics as a single observation-to-action mapping toward treating it as a predictive decision-making problem.

Generative policies offer ways to model distributions over actions. JEPAs offer a framework for learning predictive representations. World-action models bring prediction and action generation closer together. Differentiable planning offers a way to use a learned model to improve candidate actions at inference time.

None of these ideas removes the need for robust feedback, careful evaluation, or accurate-enough predictions. Nor does combining them guarantee better performance.

But together, they suggest a productive research direction: learn representations that predict the consequences of actions, generate promising action sequences, and use the model to decide which actions are worth taking.

That is the connection I want to explore further.

Papers and resources

• Much Ado About Noising: Dispelling the Myths of Generative Robotic Control
• LeWAM: End-to-End World and Action Modeling with JEPAs
• Long-WAM

This article is a synthesis of the cited research and my own questions about how the methods fit together. The proposed combined pipeline and future experiments are research directions, not claims that a single published system has already demonstrated all of these capabilities.