World Models Model-Based Reinforcement Learning
# World Models & Model-Based Reinforcement Learning
## 1. Introduction & Motivation
Model-free reinforcement learning has produced remarkable results, from Atari-playing agents to superhuman Go players, but it purchases those results with an enormous number of environment interactions. A model-free agent must experience an event many times before it can reliably act on it, because it has no internal mechanism for predicting the consequences of actions it has not yet tried. Model-based reinforcement learning addresses this inefficiency by having the agent learn an explicit, predictive model of how the environment behaves — a "world model" — and then using that model to plan, imagine rollouts, or generate synthetic training data.
The appeal of world models is straightforward: if an agent can accurately predict "if I do A, the environment will transition to state S' and I will receive reward R," it can evaluate many hypothetical action sequences internally, without paying the cost of executing them in the real environment. This is precisely how humans and animals appear to operate — mental simulation lets us avoid touching a hot stove twice. A learned dynamics model, once accurate enough, can be queried millions of times at negligible cost compared to a single real-world rollout, which is what makes model-based methods attractive in domains such as robotics, autonomous driving, and any setting where real interactions are slow, expensive, or dangerous.
The central tension in model-based RL is model accuracy versus model utility. A model that is even slightly wrong compounds its errors over the length of a planning horizon, and a policy or value function trained purely inside an inaccurate model can end up confidently pursuing actions that fail catastrophically in the real world. This is the "model exploitation" problem, and much of the research in this field — from Dyna-style architectures to uncertainty-aware planning to latent-space world models such as PlaNet, Dreamer, and MuZero — is about managing that tension, either by keeping the model honest, by bounding how far the agent trusts it, or by learning the model in a representation space where errors are less catastrophic.
## 2. Core Concepts & Theory
A world model is, at minimum, a learned approximation of the environment's transition function $\hat{T}(s_{t+1} \mid s_t, a_t)$ and, often, its reward function $\hat{R}(s_t, a_t)$. Given such a model, an agent can plan by simulating trajectories: starting from the current state, it rolls the model forward under candidate action sequences and scores the resulting imagined trajectories, without ever touching the real environment during the search.
There are two broad ways model-based RL exploits a learned model. The first is background planning (also called Dyna-style RL), where the model is used to generate additional synthetic transitions that are mixed into the replay buffer used to train a model-free policy or value function — the real environment provides "ground truth" data to keep the model calibrated, while the model amplifies the effective sample count. The second is decision-time planning, where the model is queried at every timestep to search over action sequences (e.g., via random shooting, the cross-entropy method, or Monte Carlo Tree Search) and the best sequence found is executed, typically in a receding-horizon (Model Predictive Control, MPC) fashion where only the first action is taken and the search restarts from the newly observed state.
A crucial design decision is where the model operates: in raw observation space (pixels, joint angles) or in a learned latent space. Latent world models, popularized by the World Models paper (Ha & Schmidhuber, 2018) and refined by PlaNet and Dreamer, first compress high-dimensional observations into a compact latent state using an encoder (often a variational autoencoder), then learn transition dynamics entirely within that latent space. This has two benefits: the dynamics model no longer has to predict every pixel, which is both wasteful and distracts capacity toward irrelevant detail, and planning in a low-dimensional latent space is dramatically cheaper computationally.
## 3. Mathematical Formulation
The world-model objective is typically decomposed into a representation-learning loss and a dynamics-prediction loss. For an encoder $q_\phi(z_t \mid o_t)$ mapping observations to latent states, a transition model $p_ heta(z_{t+1} \mid z_t, a_t)$, and a decoder/reward model $p_ heta(o_t, r_t \mid z_t)$, a common variational objective (as used in Dreamer-style agents) is
$
\mathcal{L}( heta, \phi) = \mathbb{E}_{q_\phi}\Big[ \sum_t \underbrace{\log p_ heta(o_t \mid z_t)}_{ ext{reconstruction}} + \underbrace{\log p_ heta(r_t \mid z_t)}_{ ext{reward prediction}} - \underbrace{\beta \, D_{KL}\big(q_\phi(z_t \mid o_{\le t}, a_{<t}) \,\|\, p_ heta(z_t \mid z_{t-1}, a_{t-1})\big)}_{ ext{dynamics consistency}} \Big]
$
The KL term forces the transition model's prior prediction of the next latent state to agree with what the encoder infers after actually seeing the next observation, which is what trains the dynamics model to be predictive.
For simpler, non-latent dynamics models, the standard supervised objective is one-step prediction error:
$
\mathcal{L}_{ ext{dyn}}( heta) = \mathbb{E}_{(s_t, a_t, s_{t+1}) \sim \mathcal{D}} \left[ \left\| \hat{f}_ heta(s_t, a_t) - s_{t+1}
ight\|_2^2
ight]
$
Given a model, planning by random shooting evaluates $N$ candidate action sequences $\{a_{t:t+H}^{(i)}\}_{i=1}^N$ of horizon $H$ under the model and picks the one maximizing predicted cumulative reward:
$
a_{t:t+H}^\star = \arg\max_i \sum_{ au=t}^{t+H} \hat{R}\big(\hat{s}_ au^{(i)}, a_ au^{(i)}\big), \qquad \hat{s}_{ au+1}^{(i)} = \hat{f}_ heta\big(\hat{s}_ au^{(i)}, a_ au^{(i)}\big)
$
The Dyna architecture's background-planning update simply augments the real-experience Bellman update with model-generated ones:
$
Q(s, a) \leftarrow Q(s, a) + \alpha \Big[ r + \gamma \max_{a'} Q(s', a') - Q(s, a) \Big], \quad (s, a, r, s') \sim \mathcal{D}_{ ext{real}} \ \cup \ \mathcal{D}_{ ext{model}}
$
where $\mathcal{D}_{ ext{model}}$ consists of synthetic transitions generated by querying $\hat{f}_ heta$ from previously visited states under sampled actions.
## 4. Advanced Theory & Extensions
Compounding error is the central theoretical obstacle in model-based RL. If a one-step model has prediction error $\epsilon$ per step, naive multi-step rollouts accumulate error that, under mild Lipschitz assumptions on the true dynamics, grows at worst exponentially with the rollout horizon $H$, i.e., $O(\epsilon \cdot L^H)$ for Lipschitz constant $L$. This theoretical fact motivates several practical mitigations: short rollout horizons (as in MBPO, Model-Based Policy Optimization, which caps synthetic rollouts to a handful of steps), ensembles of dynamics models whose disagreement estimates epistemic uncertainty and can be used to penalize the reward in uncertain regions (as in PETS and MOPO), and value-equivalent models that only need to be accurate enough to preserve the ranking of actions under the value function rather than accurate in every observable detail (as in MuZero, which learns a model purely to support value and policy prediction, with no reconstruction loss at all).
MuZero represents an important theoretical shift: rather than learning a model to match environment dynamics in observation space, it learns a model whose only job is to correctly predict future rewards, values, and policies — this is the "value equivalence" principle, formalized by Grimm et al., which shows that a model need not resemble the true environment at all as long as planning within it yields the same value estimates as planning in the true environment. This sidesteps the compounding pixel-prediction error problem entirely, since the model is never asked to predict pixels.
Another major axis of extension is uncertainty-aware planning. Probabilistic ensembles (PETS) represent both aleatoric uncertainty (irreducible environment stochasticity, captured by predicting a distribution rather than a point estimate per model) and epistemic uncertainty (uncertainty due to limited data, captured by disagreement across ensemble members). Trajectory optimization that penalizes high-disagreement regions of state-action space (as in MOPO and MOReL, primarily developed for the offline RL setting) directly operationalizes the intuition that the agent should stay within the region where its model is trustworthy.
## 5. Computational Considerations
Learning a world model shifts computational cost from environment interaction to model training and model-based rollouts, which is favorable whenever real interaction is the bottleneck (robotics, healthcare, any physical system) but adds meaningful GPU/CPU cost in simulation-heavy domains where interaction is already cheap. Latent-space world models reduce the per-rollout cost of imagined trajectories dramatically: a dynamics model operating on a 32- or 200-dimensional latent vector is orders of magnitude cheaper to query than one that must generate full-resolution image predictions, which matters enormously when thousands of imagined trajectories are evaluated per real environment step, as in Dreamer's actor-critic training entirely inside imagined rollouts.
Ensemble-based models multiply computational cost by the ensemble size (typically 5-7 models in PETS-style methods), but this cost is usually justified by the improved robustness to model exploitation, and ensemble members can be evaluated in parallel on a GPU with negligible wall-clock overhead relative to a single large model. Decision-time planning methods such as MCTS (used in MuZero) trade off search breadth and depth against wall-clock latency, which matters in real-time control settings; background-planning methods like Dyna-Q instead amortize this cost by baking imagined experience into a value function ahead of time, so acting at decision time is as cheap as a single forward pass.
## 6. Practical Implementation Strategies
A practical world-model implementation typically separates concerns into three learned components: an encoder (if operating from pixels), a dynamics/transition model, and a reward model, with training conducted via truncated backpropagation through time over short sub-sequences sampled from a replay buffer of real trajectories. Care must be taken with the replay buffer: because the dynamics model is trained only on data the policy has actually visited, a policy that drifts too far from the data distribution the model was trained on will encounter compounding model error — this is why iterative "collect real data, retrain model, replan" loops, rather than training the model once and planning forever, are standard practice.
For decision-time planning, random shooting is the simplest baseline: sample $N$ random action sequences, evaluate them under the model, execute the best one's first action, replan. The Cross-Entropy Method (CEM) improves on this by iteratively refitting a sampling distribution (typically diagonal Gaussian) to the top-performing action sequences over several iterations, which concentrates search effort in promising regions of action space far more efficiently than pure random shooting. For discrete action spaces, MCTS with a learned model (as in MuZero) is preferred, since it can exploit tree reuse and value bootstrapping.
A key implementation detail for background planning (Dyna) is limiting the rollout length used to generate synthetic transitions — MBPO's ablations showed that even 1-step model rollouts, mixed liberally into an off-policy replay buffer, deliver most of the sample-efficiency benefit, while longer rollouts quickly become harmful once model error compounds past a few steps.
## 7. Benchmark Datasets & Evaluation
Model-based RL is commonly evaluated on continuous-control benchmarks from the MuJoCo/DeepMind Control Suite (HalfCheetah, Walker2d, Hopper, Cheetah-Run) using sample efficiency (return as a function of real environment steps, not wall-clock training time) as the primary metric, since sample efficiency is exactly what model-based methods are designed to improve. Atari-100k is a standard benchmark specifically designed to stress-test sample efficiency, capping agents to 100,000 environment frames (roughly two hours of real-time play) — a regime where model-based methods such as EfficientZero and DreamerV3 substantially outperform model-free baselines like Rainbow DQN, which need tens of millions of frames to reach comparable performance under standard training.
Board games (Go, Chess, Shogi) and procedurally generated environments (Procgen, Crafter) are used to evaluate planning quality and generalization respectively — Crafter in particular was designed to test whether an agent's world model supports long-horizon, compositional reasoning (crafting tools, managing hunger, avoiding threats) rather than just short-horizon reactive control. Evaluation should always separate "asymptotic performance" (how good is the final policy, given unlimited data) from "sample efficiency" (how good is the policy after a fixed, small budget of real interactions), since model-based methods sometimes trade a small amount of asymptotic performance for a large improvement in the latter.
## 8. Key Challenges & Limitations
The dominant failure mode is model exploitation: because the policy is optimized (via planning or synthetic data) against the learned model rather than the true environment, it can discover and aggressively exploit regions where the model is systematically wrong but predicts high reward, resulting in policies that look excellent inside imagination and fail in reality. This is especially severe in the offline RL setting, where there is no opportunity to collect corrective real data, motivating explicit uncertainty penalization (MOPO, MOReL, COMBO).
Partial observability and stochasticity are also challenging: a deterministic dynamics model trained via mean-squared error will, when the true environment is stochastic or multimodal (e.g., an opponent that could move left or right), learn to predict the *average* of the possible outcomes, which may not correspond to any physically valid outcome at all — this motivates distributional or latent-variable dynamics models that can represent multimodal transition distributions.
Long-horizon compounding error remains fundamentally unsolved in the sense that no dynamics model is perfectly accurate, so any application requiring very long planning horizons (hundreds or thousands of steps) will eventually see planning quality degrade, which is why most successful model-based systems either use short planning/rollout horizons, replan frequently (MPC), or, like MuZero, avoid needing the model to be observation-accurate at all.
## 9. Hyperparameter Tuning
The rollout horizon $H$ used for synthetic data generation or planning is the single most consequential hyperparameter: too short and the model's advantages are underused, too long and compounding error dominates — MBPO-style methods typically sweep $H \in \{1, \dots, 15\}$ and often find a small horizon (1-5 steps) works best even though it seems conservative. Ensemble size for probabilistic dynamics models trades computational cost against uncertainty-estimation quality, with diminishing returns typically observed beyond 5-7 members. The KL weight $\beta$ in variational latent world models (analogous to a $\beta$-VAE) controls the trade-off between reconstruction fidelity and how strongly the transition prior is regularized to match the encoder's posterior; too high a $\beta$ can collapse the latent representation and destroy predictive information, too low a $\beta$ can make the prior a poor predictor, forcing the agent to rely on the (unavailable at planning time) encoder.
The ratio of real to synthetic (model-generated) transitions used to train the policy or value function also matters, as does the frequency of model retraining relative to policy updates — retraining the model too infrequently leaves it stale relative to the evolving policy's visited state distribution, while retraining too frequently can be computationally wasteful and destabilize the downstream policy optimization if the model's predictions shift significantly between updates.
## 10. Real-World Applications & Case Studies
Robotic manipulation and locomotion are the most natural domains for model-based RL, since real robot interaction is slow, requires human supervision, and risks hardware damage; MPC-based approaches using learned dynamics models have been used for tasks like in-hand manipulation and legged locomotion (e.g., work at UC Berkeley and elsewhere on model-based control for quadrupeds), where a model trained on a modest number of real trials enables extensive internal replanning. In autonomous driving research, world models are used both for closed-loop planning and for generating synthetic training scenarios that would be too dangerous or rare to collect in the real world.
MuZero's application to board games (Go, Chess, Shogi) and Atari demonstrated that a learned, value-equivalent model — with no access to the actual game rules — could match or exceed AlphaZero (which is given the true simulator), establishing that model-based planning can be competitive even when the model itself is entirely learned. DreamerV3 was notably applied, with a single set of hyperparameters, across domains ranging from Atari and DMC to Minecraft's "collect a diamond" challenge (a very long-horizon, sparse-reward task), demonstrating that latent world models can support meaningfully long-horizon planning when engineered carefully.
## 11. Integration with Other Methods
Model-based RL integrates naturally with offline RL, where a learned model can serve as a proxy simulator to generate additional synthetic rollouts from a fixed dataset, provided uncertainty penalization prevents the policy from exploiting model errors in out-of-distribution regions (MOPO, MOReL, COMBO — the last of which explicitly combines model-based data generation with conservative Q-learning objectives). It also integrates with hierarchical RL, where a world model's latent state can serve as the observation space for a higher-level policy operating over longer effective timesteps.
Model-based methods pair naturally with exploration bonuses: because the agent has an explicit predictive model, prediction error itself is a natural curiosity signal (as in intrinsic curiosity module approaches), directing exploration toward regions where the model is currently inaccurate and thus most informative to visit. Finally, model-based planning is a natural complement to imitation learning and offline pretraining: a world model can be pretrained on large amounts of passively collected or demonstration data, then fine-tuned or used purely for planning once the agent begins interacting with the environment, reducing the online sample complexity even further.
## 12. Future Research Directions
An active area of research is scaling world models the way language models have been scaled — training large, general-purpose predictive models of physical or simulated dynamics on massive, diverse interaction datasets (video, robot trajectories, simulation logs) and then adapting them to specific downstream control tasks, in the spirit of foundation models. Video-prediction-based world models (e.g., Genie, and various "world model as a generative video model" approaches) push this further by attempting to learn action-conditioned dynamics directly from internet-scale video, without an explicit low-dimensional state representation at all.
Another direction is combining the value-equivalence principle of MuZero with the representation richness of latent generative world models, aiming to get the best of both: models that are cheap and robust to plan in, yet flexible enough to support diverse downstream objectives beyond a single fixed value function. Formal guarantees on when model-based planning is safe — precise, checkable bounds on compounding error and out-of-distribution model reliability — remain an open theoretical problem, particularly relevant as these methods move toward safety-critical, real-world deployment such as robotics and autonomous vehicles, where uncontrolled model exploitation carries physical consequences.
## 13. Summary & Key Takeaways
Model-based reinforcement learning trades additional model-learning machinery for dramatically improved sample efficiency, by letting an agent plan or generate synthetic experience via an internal, learned approximation of environment dynamics rather than relying purely on real interaction. The central theoretical challenge — compounding model error over long rollouts — has motivated a range of practical solutions: short synthetic rollout horizons, ensemble-based uncertainty estimation, value-equivalent models that avoid needing to be observationally accurate, and latent-space dynamics models that are both cheaper to query and less prone to wasting capacity on irrelevant detail. From Dyna-Q's simple background planning to Dreamer's imagination-only actor-critic training to MuZero's value-equivalent tree search, the field's history is largely a story of finding better ways to manage the trust an agent places in its own internal model of the world, and this trade-off remains the defining design axis for any new model-based method.
Keywords: world models, model-based reinforcement learning, Dyna-Q, dynamics model, latent world model, Dreamer, PlaNet, MuZero, model predictive control, random shooting, cross-entropy method, compounding error, model exploitation, value equivalence, ensemble dynamics model, MBPO, PETS, offline model-based RL, MOPO, sample efficiency, background planning, decision-time planning
---
## Appendix: Practical Labs
### Lab 1: Learned Dynamics Model Prediction Error Decreases With Training Data
import numpy as np
def true_dynamics(state, action):
"""Ground-truth 2D dynamics: next_state = state + action + small coupling
term, simulating a simplified point-mass system with mild nonlinearity."""
x, y = state
dx, dy = action
return np.array([x + dx + 0.05 * x * dy, y + dy + 0.05 * y * dx])
def generate_dataset(n_samples, rng):
states = rng.uniform(-2, 2, size=(n_samples, 2))
actions = rng.uniform(-0.5, 0.5, size=(n_samples, 2))
next_states = np.array([true_dynamics(s, a) for s, a in zip(states, actions)])
return states, actions, next_states
def fit_linear_dynamics_model(states, actions, next_states):
"""Fits next_state ~= W @ [state, action, 1] via least squares."""
X = np.hstack([states, actions, np.ones((len(states), 1))])
W, _, _, _ = np.linalg.lstsq(X, next_states, rcond=None)
return W
def predict(W, state, action):
x = np.concatenate([state, action, [1.0]])
return x @ W
def evaluate_model(W, states, actions, next_states):
preds = np.array([predict(W, s, a) for s, a in zip(states, actions)])
return np.mean(np.sum((preds - next_states) ** 2, axis=1))
def test_learned_dynamics_model_reduces_prediction_error_with_more_data():
rng = np.random.RandomState(0)
test_states, test_actions, test_next_states = generate_dataset(500, rng)
train_sizes = [5, 20, 100, 1000]
errors = []
for n in train_sizes:
train_states, train_actions, train_next_states = generate_dataset(n, rng)
W = fit_linear_dynamics_model(train_states, train_actions, train_next_states)
err = evaluate_model(W, test_states, test_actions, test_next_states)
errors.append(err)
print(f"{'train size':>10} | {'test MSE':>10}")
for n, e in zip(train_sizes, errors):
print(f"{n:10d} | {e:10.6f}")
# More training data should yield a model at least as accurate, and the
# smallest dataset should be clearly worse than the largest.
assert errors[-1] < errors[0], "More training data should substantially reduce prediction error"
assert errors[-1] < 0.01, "With ample data, the near-linear dynamics should be fit very accurately"
print("Learned dynamics model data-scaling test passed.")
if __name__ == "__main__":
test_learned_dynamics_model_reduces_prediction_error_with_more_data()### Lab 2: Dyna-Q Improves Sample Efficiency Over Plain Q-Learning
import numpy as np
def make_chain_mdp(n_states=10):
"""A simple chain MDP: states 0..n_states-1, actions {0: left, 1: right},
reward of +1 only for reaching the rightmost state, 0 otherwise, episode
resets there. This favors an agent that can propagate reward information
quickly, which model-based planning should do more efficiently."""
return n_states
def step(state, action, n_states):
if action == 1:
next_state = min(state + 1, n_states - 1)
else:
next_state = max(state - 1, 0)
reward = 1.0 if next_state == n_states - 1 else 0.0
done = next_state == n_states - 1
return next_state, reward, done
def q_learning(n_states, n_real_steps, planning_steps, rng, alpha=0.3, gamma=0.95, eps=0.3):
"""Tabular Q-learning with optional Dyna-style background planning using a
learned (here: tabular, exactly-remembered) model of observed transitions."""
# Small random initialization (rather than all-zeros) avoids a degenerate
# tie-breaking bias toward action index 0 that would otherwise trap the
# agent near the start of the chain regardless of planning.
Q = rng.uniform(-0.01, 0.01, size=(n_states, 2))
model = {} # (state, action) -> (next_state, reward)
state = 0
cumulative_reward = 0.0
rewards_per_step = np.zeros(n_real_steps)
for t in range(n_real_steps):
if rng.uniform(0, 1) < eps:
action = rng.randint(0, 2)
else:
action = int(np.argmax(Q[state]))
next_state, reward, done = step(state, action, n_states)
Q[state, action] += alpha * (reward + gamma * np.max(Q[next_state]) - Q[state, action])
model[(state, action)] = (next_state, reward)
# Background planning: replay random previously observed transitions
if planning_steps > 0 and len(model) > 0:
keys = list(model.keys())
for _ in range(planning_steps):
ps, pa = keys[rng.randint(0, len(keys))]
pns, pr = model[(ps, pa)]
Q[ps, pa] += alpha * (pr + gamma * np.max(Q[pns]) - Q[ps, pa])
cumulative_reward += reward
rewards_per_step[t] = cumulative_reward
state = 0 if done else next_state
return rewards_per_step
def test_dyna_q_improves_sample_efficiency_over_q_learning():
n_states = 8
n_real_steps = 200
n_seeds = 20
final_cumreward_plain = []
final_cumreward_dyna = []
for seed in range(n_seeds):
rng = np.random.RandomState(seed)
plain = q_learning(n_states, n_real_steps, planning_steps=0, rng=rng)
rng2 = np.random.RandomState(seed)
dyna = q_learning(n_states, n_real_steps, planning_steps=10, rng=rng2)
final_cumreward_plain.append(plain[-1])
final_cumreward_dyna.append(dyna[-1])
mean_plain = np.mean(final_cumreward_plain)
mean_dyna = np.mean(final_cumreward_dyna)
print(f"Plain Q-learning cumulative reward after {n_real_steps} real steps: {mean_plain:.2f}")
print(f"Dyna-Q cumulative reward after {n_real_steps} real steps: {mean_dyna:.2f}")
# Dyna-Q's background planning should propagate reward information faster
# through the chain, yielding strictly higher cumulative reward for the
# same number of REAL environment steps.
assert mean_dyna > mean_plain, "Dyna-Q should be more sample-efficient than plain Q-learning"
print("Dyna-Q sample efficiency test passed.")
if __name__ == "__main__":
test_dyna_q_improves_sample_efficiency_over_q_learning()### Lab 3: Compounding Rollout Error Grows With Planning Horizon
import numpy as np
def true_step(state):
"""A mildly expansive rotation used as ground-truth single-step dynamics:
it rotates the state and grows its norm slightly each step, so that small
per-step model errors compound rather than being squashed by saturation
(as a bounded/contractive map like tanh would do)."""
theta = 0.3
rotation = np.array([[np.cos(theta), -np.sin(theta)],
[np.sin(theta), np.cos(theta)]])
return 1.05 * (rotation @ state)
def learned_step(state, noise_std, rng):
"""A model of the true dynamics that adds small zero-mean Gaussian noise
to each prediction, simulating a per-step approximation error."""
return true_step(state) + rng.normal(0, noise_std, size=2)
def rollout_error(initial_state, horizon, noise_std, rng, n_trials=300):
"""Rolls out both the true dynamics and the noisy model dynamics for
`horizon` steps from the same initial state, and measures how far the
model's trajectory diverges from the true trajectory at the final step."""
errors = np.zeros(n_trials)
for i in range(n_trials):
true_state = initial_state.copy()
model_state = initial_state.copy()
for _ in range(horizon):
true_state = true_step(true_state)
model_state = learned_step(model_state, noise_std, rng)
errors[i] = np.linalg.norm(true_state - model_state)
return errors.mean()
def test_rollout_compounding_error_grows_with_horizon():
rng = np.random.RandomState(3)
initial_state = np.array([0.5, -0.3])
noise_std = 0.01
horizons = [1, 3, 6, 10]
errors = [rollout_error(initial_state, h, noise_std, rng) for h in horizons]
print(f"{'horizon':>8} | {'mean divergence':>16}")
for h, e in zip(horizons, errors):
print(f"{h:8d} | {e:16.5f}")
# Divergence from the true trajectory should grow monotonically with
# rollout horizon, since per-step errors compound.
assert errors[0] < errors[1] < errors[2] < errors[3], \
"Compounding rollout error should increase monotonically with horizon"
# The longest horizon should diverge substantially more than a single step
assert errors[-1] > 3 * errors[0], "Long-horizon rollout error should be much larger than single-step error"
print("Compounding rollout error test passed.")
if __name__ == "__main__":
test_rollout_compounding_error_grows_with_horizon()### Lab 4: Random-Shooting MPC With a Learned Model Beats a Random Policy
import numpy as np
def env_dynamics(state, action):
"""Point-mass control task: state is 2D position, action is a 2D
velocity command (clipped), target is the origin, and reward is negative
squared distance to target after moving."""
action = np.clip(action, -1.0, 1.0)
return state + 0.2 * action
def reward_fn(state):
return -np.sum(state ** 2)
def fit_dynamics_model(n_samples, rng):
"""Fits a linear model of env_dynamics from randomly sampled transitions,
used as the 'learned world model' for planning."""
states = rng.uniform(-3, 3, size=(n_samples, 2))
actions = rng.uniform(-1, 1, size=(n_samples, 2))
next_states = np.array([env_dynamics(s, a) for s, a in zip(states, actions)])
X = np.hstack([states, actions, np.ones((n_samples, 1))])
W, _, _, _ = np.linalg.lstsq(X, next_states, rcond=None)
return W
def model_predict(W, state, action):
x = np.concatenate([state, action, [1.0]])
return x @ W
def random_shooting_mpc(state, W, horizon, n_candidates, rng):
"""Samples n_candidates random action sequences of length `horizon`,
evaluates predicted cumulative reward under the learned model, and
returns the first action of the best sequence (receding-horizon MPC)."""
best_return = -np.inf
best_first_action = None
for _ in range(n_candidates):
s = state.copy()
actions = rng.uniform(-1, 1, size=(horizon, 2))
total_reward = 0.0
for a in actions:
s = model_predict(W, s, a)
total_reward += reward_fn(s)
if total_reward > best_return:
best_return = total_reward
best_first_action = actions[0]
return best_first_action
def run_episode(policy_fn, start_state, n_steps, rng, W=None):
state = start_state.copy()
total_reward = 0.0
for _ in range(n_steps):
action = policy_fn(state, W, rng) if W is not None else policy_fn(state, rng)
total_reward += reward_fn(state)
state = env_dynamics(state, action)
return total_reward
def random_policy(state, rng):
return rng.uniform(-1, 1, size=2)
def mpc_policy(state, W, rng):
return random_shooting_mpc(state, W, horizon=5, n_candidates=64, rng=rng)
def test_mpc_planning_outperforms_random_policy():
rng = np.random.RandomState(7)
W = fit_dynamics_model(2000, rng)
start_state = np.array([2.0, -1.5])
n_seeds = 10
n_steps = 15
random_returns = []
mpc_returns = []
for seed in range(n_seeds):
rng_eval = np.random.RandomState(100 + seed)
random_returns.append(run_episode(random_policy, start_state, n_steps, rng_eval))
rng_eval2 = np.random.RandomState(100 + seed)
mpc_returns.append(run_episode(mpc_policy, start_state, n_steps, rng_eval2, W=W))
mean_random = np.mean(random_returns)
mean_mpc = np.mean(mpc_returns)
print(f"Random policy mean return: {mean_random:.3f}")
print(f"MPC (learned model) mean return: {mean_mpc:.3f}")
# Planning with the learned world model should drive the state toward the
# origin far more effectively than random actions, yielding a much higher
# (less negative) cumulative reward.
assert mean_mpc > mean_random, "Model-based MPC planning should outperform a random policy"
assert mean_mpc > 0.5 * mean_random, "MPC should substantially close the gap to zero reward"
print("Random-shooting MPC test passed.")
if __name__ == "__main__":
test_mpc_planning_outperforms_random_policy()