Building a First World Model › Planning with Imagined Trajectories

This section covers the payoff: using the trained model to choose actions without consulting the real environment.

Given the current filtered latent $h_0$ and a candidate action sequence $(a_0, \dots, a_{H-1})$, the model can unroll a purely imaginary trajectory:

$$\bar{h}_{t+1} = w_\theta(\bar{h}_t, a_t), \qquad \hat{r}_t = g_\omega([\bar{h}_{t+1}; a_t])$$

and summarize it with a predicted return:

$$J(a_{0:H-1}) \;=\; \sum_{t=0}^{H-1} \gamma^t \, \hat{r}_t$$

where $\gamma \in (0, 1]$ is a discount factor. Planning is now a search problem over action sequences: find the sequence maximizing $J$. Note what has been achieved. The search touches the real environment zero times. All evaluation of candidate futures happens inside the model, at the speed of matrix multiplication.

The discount factor $\gamma$ plays a quiet but important role here. Setting $\gamma = 1$ treats all steps within the horizon equally, which suits our task since reaching the goal at any point within ten steps is equally valuable. Values below one emphasize near-term rewards, which has a side benefit for imperfect models: distant imagined rewards are computed from deeply compounded predictions, the least trustworthy numbers the model produces, so discounting them is a crude form of epistemic humility baked into the objective. In general, shorter horizons and stronger discounts both reduce exposure to compounding error, at the cost of myopia about long-term consequences. The choice of horizon and discount is thus another place where model reliability and task requirements must be balanced explicitly.

With a differentiable model we could even backpropagate the return gradient to the actions themselves, a technique explored in Part VII: Planning and Agency. Here we use a simpler and very robust derivative-free method.

The cross-entropy method (CEM) maintains a Gaussian distribution over action sequences and iteratively sharpens it around high scorers:

1. Sample a population of action sequences from the current Gaussian.
2. Roll out each candidate in imagination and compute its predicted return.
3. Keep the top few sequences, the elites.
4. Refit the Gaussian's mean and standard deviation to the elites.
5. Repeat for a handful of iterations, then return the mean's first action.

CEM is not clever, and that is its virtue. It makes no assumptions about the return landscape beyond local smoothness, it parallelizes trivially, and it degrades gracefully when the model is imperfect. For a first world model it is exactly the right partner.

A few implementation details deserve explanation, since they matter in practice. Sampling the initial Gaussian with a standard deviation equal to the full action range means the first iteration explores broadly across the space of possible plans, and each subsequent refitting concentrates the distribution around whatever worked. The small constant added to the refitted standard deviation prevents premature collapse, the situation where the population becomes so concentrated that all candidates are effectively identical and the search stalls. Taking only the first action of the best plan, rather than executing the whole plan, is deferred to the next subsection. And because every candidate rollout consists of the same number of matrix multiplications through the same networks, the entire population can be scored in a single batched pass, which is why CEM feels instantaneous even on a laptop despite evaluating hundreds of imagined futures per decision.

Executing a whole planned sequence open-loop would be reckless: any model error compounds over the horizon. Instead we use receding-horizon control, also known as model-predictive control (MPC):

• Every environment step, filter the newest observation into the latent state.
• Plan a fresh $H$-step action sequence by CEM in imagination.
• Execute only the first action.
• Observe what actually happened, and repeat.

Frequent replanning converts the model from an oracle into a advisor whose suggestions are constantly fact-checked by reality. Short model horizons can suffice because each cycle only needs to identify a good first move. The price is computation per step, which is why the efficiency of imagined rollouts matters in practice.

There is a deeper way to understand why this works, connecting to classical control theory. MPC predates learned world models by decades, and in its classical form it uses analytically specified dynamics, typically linear equations fit by system identification. The loop structure, filter, optimize, apply first move, re-filter, is unchanged whether the dynamics come from least squares on a linear model or from a recurrent network trained on pixels. What learned world models contribute is scope: the same loop now applies to systems whose dynamics are nonlinear, high-dimensional, or unknown in analytic form. What classical MPC contributes to us is a mature understanding of the loop's robustness properties, including the fact that closed-loop feedback can stabilize systems far beyond the horizon over which the internal predictions remain accurate. Our puck controller is, quite literally, a modern instance of a fifty-year-old idea, with a neural simulator substituted for the identified linear model.