World models and embodied AI

What Are World Models?

Why do we need this?

Sample efficiency: Real-world interactions are expensive (robot might break, medical trial costs lives). Simulate rollouts in your head (model) instead.
Planning: With a model, you can "think ahead" by simulating future trajectories before acting (like chess players visualizing moves).
Transfer: A good world model captures environment dynamics that transfer across tasks (same physics, different goals).

How is this different from model-free RL?

Model-free (DQN, PPO): Learns $Q(s,a)$ or $\pi(a|s)$ directly from trial-and-error. No explicit understanding of consequences.
Model-based: Learns $p(s_{t+1}|s_t, a_t)$ first, then uses it to plan or train a policy. Understands "if I push this, it will fall."

Derivation: How World Models Enable Planning

Step 1: The Forward Model

Start with the Markov Decision Process (MDP) dynamics: $s_{t+1} \sim P(s_{t+1} | s_t, a_t)$

In model-free RL, we never learn $P$ explicitly. We only learn the value function $V(s)$ or policy $\pi(a|s)$ from samples.

World model approach: Learn an approximation $\hat{P}_\theta(s_{t+1} | s_t, a_t)$ (the forward model).

Why learn this explicitly? Because once we have it, we can:

Rollout hypothetical trajectories without real environment interaction
Backpropagate through the model to optimize actions
Data augmentation by simulating synthetic experience

Step 2: Planning with the Model

Given $\hat{P}_\theta$ , we want to find the best action sequence. Use Model Predictive Control (MPC):

$a^*_{0:H-1} = \arg\max_{a_{0:H-1}} \mathbb{E}_{\hat{P}_\theta} \left[ \sum_{t=0}^{H-1} r(s_t, a_t) \right]$

Derivation from first principles:

Start with the trajectory return: $J(\tau) = \sum_{t=0}^{H-1} r(s_t, a_t)$

We want to maximize expected return under our learned dynamics: $J(a_{0:H-1}) = \mathbb{E}_{s_1 \sim \hat{P}(\cdot|s_0, a_0)} \mathbb{E}_{s_2 \sim \hat{P}(\cdot|s_1, a_1)} \cdots \left[ \sum_{t=0}^{H-1} r(s_t, a_t) \right]$

Why this step? Each future state $s_t$ is random because the model is stochastic. We integrate over all possible futures weighted by their probability.

Expanding the nested expectation: $J = \int \hat{P}(s_1|s_0, a_0) \left[ r(s_0, a_0) + \int \hat{P}(s_2|s_1, a_1) \left[ r(s_1, a_1) + \cdots \right] ds_2 \right] ds_1$

Why recursive? The return is additive over time, and each state depends only on the previous (Markov property).

In practice, we approximate this integral by:

Sampling $K$ trajectories from $\hat{P}_\theta$
Averaging their returns
Optimizing $a_{0:H-1}$ via gradient descent or shooting methods (CEM, CMA-ES)

$J \approx \frac{1}{K} \sum_{k=1}^K \sum_{t=0}^{H-1} r(s_t^{(k)}, a_t) \quad \text{where } s_t^{(k)} \sim \hat{P}_\theta(\cdot | s_{t-1}^{(k)}, a_{t-1})$

Key insight: The model lets us evaluate action sequences in simulation before committing in reality.

Formula Model-Based Return Gradient

If the model is differentiable, we can compute: $\frac{\partial J}{\partial a_t} = \mathbb{E}_{\hat{P}_\theta} \left[ \frac{\partial r(s_t, a_t)}{\partial a_t} + \sum_{t'=t+1}^{H-1} \frac{\partial r(s_{t'}, a_{t'})}{\partial s_{t'}} \frac{\partial s_{t'}}{\partial a_t} \right]$

Why? Actions affect future states via the chain rule $\frac{\partial s_{t'}}{\partial a_t} = \prod_{k=t}^{t'-1} \frac{\partial \hat{P}_\theta(s_{k+1}|s_k, a_k)}{\partial s_k}$ .

This enables gradient-based planning (backprop through time in the model).

Step 3: Dyna Architecture (Learning + Planning)

The Dyna framework unifies model-learning and model-free RL:

Real experience: $(s, a, r, s') \rightarrow$ update policy $\pi$ (model-free) AND model $\hat{P}_\theta$
Simulated experience: Sample $(s, a) \sim$ replay buffer, generate $s' \sim \hat{P}_\theta(s'|s,a)$ → update policy

Why does this work?

Real data is scarce → use model to generate synthetic data
Model errors compound over long rollouts → use short $k$ -step rollouts
Policy improves faster because it gets $N_{\text{real}} + N_{\text{sim}}$ updates per real sample

The Dyna update rule: $\theta_\pi \leftarrow \theta_\pi + \alpha \left[ \nabla_\theta J_{\text{real}} + \lambda \nabla_\theta J_{\text{sim}} \right]$ where $\lambda$ trades off model trust.

Why the $\lambda$ weight? If the model is poor, simulated gradients are biased → reduce their influence.

What Is Embodied AI?

Why embodiment matters:

Grounding: Abstract concepts (e.g., "heavy") are learned through physical interaction, not definitions
Affordances: What you can do depends on your body (a quadruped can't open doors the way a humanoid can)
Exploration: Embodied agents naturally explore by moving, manipulating, trying—this generates diverse data
Partial observability: Real bodies have limited sensors → agent must actively gather information

Historical context: Classical AI (GOFAI) tried to build intelligence via logic and symbols. The embodied cognition movement argued this ignores how biological intelligence is inseparable from the body.

Architecture: World Models (Ha & Schmidhuber, 2018)

The seminal World Models paper proposed a three-component architecture:

V (Vision): Variational autoencoder learns latent representation $z_t$ of observation $o_t$ $z_t = \text{Encoder}(o_t), \quad o_t \approx \text{Decoder}(z_t)$
M (Memory): Mixture Density Network RNN predicts next latent state $P(z_{t+1} | z_t, a_t, h_t) = \sum_{i=1}^K \pi_i \mathcal{N}(\mu_i, \sigma_i^2)$ where $h_t$ is RNN hidden state.
C (Controller): Small linear policy $a_t = W z_t + Wh h_t + b$

Why this architecture?

V compresses high-dim observations (pixels) → low-dim $z_t$ (like "there's a car ahead")
M learns temporal dynamics → predicts future $z_{t+1}$ (like "if I steer left, the car will move right in my view")
C is tiny because it operates on abstract $z_t$ , not raw pixels → easy to optimize

Key innovation: Train M on real trajectories, then train C entirely inside M's dream (simulated rollouts). No real environment needed after collecting initial data.

Derivation of M's loss:

We want $P(z_{t+1} | z_t, a_t)$ . Use maximum likelihood: $\mathcal{L}_M = -\log P_M(z_{t+1} | z_t, a_t, h_t)$

For the Mixture Density Network: $P_M(z_{t+1}) = \sum_{i=1}^K \pi_i(z_t, a_t, h_t) \cdot \mathcal{N}(z_{t+1}; \mu_i, \sigma_i^2)$

Why mixture of Gaussians? Because the future is stochastic and multimodal (e.g., a car might turn left OR right). A single Gaussian would predict the average (going straight), which is wrong.

The negative log-likelihood: $\mathcal{L}_M = -\log \left( \sum_{i=1}^K \pi_i \exp\left( -\frac{(z_{t+1} - \mu_i)^2}{2\sigma_i^2} \right) \right)$

Why this loss? It's the cross-entropy between the true next state distribution and our model's prediction. Minimizing it makes the model match reality.

Worked example Example 1: Car Racing

Task: Agent controls a car in a top-down racing game. Setup:

Observation: 64×64 RGB image
V: VAE compresses to $z_t \in \mathbb{R}^{32}$
M: LSTM with 256 hidden units predicts $z_{t+1}$
C: Linear layer $a_t = W [z_t; h_t] + b$ , actions = {steer, gas, brake}

Training:

Collect 10,000 random rollouts
Train V to reconstruct frames: $\mathcal{L}_V = ||o_t - \hat{o}_t||^2 + \beta \cdot D_{KL}(q(z|o) || p(z))$

Why the KL term? Regularizes latent space to be smooth, continuous (enables interpolation)

Encode all frames to $z_t$ , train M to predict $z_{t+1}$
Train C withMA-ES inside M's predictions (no real environment)

Why CMA-ES? It's gradient-free → works even though M is stochastic/non-differentiable Result: Agent learns to race using only the model after seeing 10k rollouts. In real env, it drives competently because M captured the dynamics.

Why this step? The model learned that "steering left makes track features move right" and "speed affects how fast scenery changes." C exploits this.

Worked example Example 2: Robotic Manipulation

Task: Robot arm must grasp objects from pixels.

Challenge: High-dimensional continuous control, contact dynamics are hard to model.

Approach:

Learn world model $\hat{P}(s_{t+1}, r_t | s_t, a_t)$ where $s_t$ = proprioception + image features
Use model for trajectory optimization: given goal (object in hand), search over action sequences
Planning method: Cross-Entropy Method (CEM)

Sample $N$ candidate action sequences from $\mathcal{N}(\mu, \sigma^2)$
Rollout each in the model, evaluate return
Keep top $M$ elites, refit $\mu, \sigma$ to their distribution
Repeat 5-10 iterations

Why CEM?

Derivative-free → works with stochastic models
Iterative refinement → starts with broad exploration, narows to best region Step-by-step rollout:

Initialize: μ = [0, 0, 0], σ² = [1, 1, 1]  (3D action)
Iteration 1:
 Sample 100 sequences from N(μ, σ²)
 Rollout in model: s₁ ~ P̂(·|s₀, a₀), r₀ = reward(s₀, a₀), ...
 Return J = Σ r_t
 Top 10 elites have mean μ' = [0.3, -0.2, 0.5]
 Update μ ← μ', σ² ← std²(elites)
Iteration 2:
 Sample around new μ (tighter σ)...
After 10 iterations: converge to best action sequence
Execute a₀ in real robot, observe s₁, replan

Why replan every step? Model is imperfect → errors compound. Closed-loop replanning corrects drift.

Common Mistakes

Common mistake Mistake 1: Model Overfitting to Training Data

Wrong idea: "I'll train a world model on1000 trajectories, then plan infinitely many rollouts in it."

Why feels right: More simulated data should improve the policy.

Why it's wrong: The model is an approximation. It has modeling error $\epsilon_M = |\hat{P} - P||$ . Errors compound exponentially with rollout length: $\text{Error at step } H \approx \epsilon_M \cdot \gamma^H \cdot H$ where $\gamma$ is the error propagation rate.

After $H \sim 10$ steps, the simulated trajectory diverges wildly from reality. The policy learns to exploit model bugs (e.g., "clip through walls" in a bugy physics sim).

Fix: Use short rollouts (k=1 to 5), or trust-region methods that bound divergence from real data. The model-based value expansion: $Q(s,a) \approx \sum_{t=0}^{k-1} \gamma^t r_t + \gamma^k V(s_k) \quad \text{(k-step model, then model-free V)}$ Why? Model is accurate short-term, fall back to learned $V$ for long-term.

Common mistake Mistake 2: Ignoring Partial Observability

Wrong idea: "The agent sees an image, so $s_t = \text{image}_t$ . That's the full state."

Why it feels right: The image contains all visible information.

Why it's wrong: An image is a2D projection → loses depth, oclusions, velocities. The true state includes hidden variables (e.g., "is the enemy behind the wall?"). This is a POMDP (Partially Observable MDP).

In a POMDP: $o_t = h(s_t)$ (observation function), $s_t$ is latent. Consequence: A Markov policy $\pi(a|o_t)$ is suboptimal because $_t$ doesn't satisfy $P(s_{t+1}|s_t, a_t)$ (Markov property breaks at observation level).

Fix: Maintain belief state $b_t = P(s_t | o_{0:t}, a_{0:t-1})$ or use recurrent policy $\pi(a | h_t)$ where $h_t = \text{RNN}(h_{t-1}, o_t, a_{t-1})$ . The RNN hidden state acts as memory.

Why RNN helps? $h_t$ aggregates history → implicitly estimates hidden state. For example, if the enemy disappeared behind a wall, $h_t$ "remembers" their last position and velocity to predict where they'll emerge.

Common mistake Mistake 3: Training World Model on Expert Data Only

Wrong idea: "I'll train my model on expert trajectories, then use it to train a policy."

Why it feels right: Expert data is high-quality (good outcomes), should teach the model correct dynamics.

Why it's wrong: The learner policy will make different mistakes than the expert. It will visit states the expert never saw. The model has no data there → makes wild predictions → policy gets misleading gradients.

This is the distributional shift problem. $P_{\text{data}}(s,a) \neq P_{\pi}(s,a)$ .

Fix: Use on-policy data (collect with current policy) or DAger (Dataset Aggregation):

Train policy $\pi_1$ on expert data
Collect rollouts with $\pi_1$ , query expert for correct actions
Aggregate: $\mathcal{D} \leftarrow \mathcal{D} \cup \{(s_{\pi_1}, a_{\text{expert}})\}$
Retrain $\pi_2$ on augmented $\mathcal{D}$
Repeat

Why this works? The learner's mistakes become training data with expert corrections → model learns to recover from errors.

Connections to Other Concepts

Model-FreeRL: World models are the complement—model-based methods need sample efficiency but are complex to train
Variational-Autoencoders: VAEs provide the V (Vision) component for latent state representation
Recurrent-Neural-Networks: RNNs (LSTMs, GRUs) are the standard M (Memory) for temporal dynamics
Reinforcement-Learning-Basics: World models fit into the MDP framework, learn the transition function $P(s'|s,a)$
Imitation-Learning: Behavioral cloning ignores world models; inverse RL can learn reward, then use model to plan
Sim-to-Real-Transfer: Embodied AI must bridge the "reality gap" between simulated world models and real physics
Vision-Transformers: Modern V components might useViT encoders instead of CNs/VAEs for better representation
Diffusion-Models: Emerging approach for generative world models (predict future frames as denoising process)

Recall Explain to a12-Year-Old

Imagine you're learning to ride a bike. At first, you fall a lot because you don't know what happens when you turn the handlebars or pedal faster. But your brain is building a mental model: "If I turn left, the bike goes left. If I pedal hard going uphill, I slow down less."

After enough practice, you can imagine what will happen before you do it. You see a turn coming and think "If I lean and turn now, I'll make it smoothly." You're planning in your head using your model of bike physics.

That's a world model! It's your brain's simulator. World models in AI are the same: the robot learns how its body and the environment work, then "thinks ahead" by simulating actions in its mind before doing them in real life.

Embodied AI means the robot has a body (not just a brain in a computer). It learns by doing—pushing things, moving around, touching objects—just like you learned by actually riding the bike, not by reading a book about it. The body and experience are essential to being smart.

#flashcards/ai-ml

What is a world model in reinforcement learning? :: A learned probabilistic model $p(s_{t+1}, r_t | s_t, a_t)$ that predicts next state and reward given current state and action—an internal simulation of environment dynamics.

Why do world models improve sample efficiency?

They allow agents to simulate rollouts internally instead of requiring real environment interactions, which are often expensive or dangerous in real-world settings.

What are the three components of the World Models architecture?

V (Vision/VAE encoder), M (Memory/RNN dynamics model), C (Controller/policy)—V compresses observations, M predicts future latents, C outputs actions.

What is embodied AI?

AI systems that have a physical or simulated body and learn through sensorimotor interaction with rich environments, not just abstract symbolic reasoning or text processing.

Why do world model errors compound over long rollouts?

Each prediction depends on the previous, so error

\epsilon_M

at step

t

affects all future steps—total error grows as

\epsilon_M \cdot H

where

H

is horizon length.

What is the Dyna architecture?

A framework that combines model-free RL updates from real experience with model-based updates from simulated experience generated by the world model.

What is the distributional shift problem in world models?

When a model is trained on expert data but used by a learner policy that visits different states, the model has no data in the learner's distribution and makes poor predictions.

How does a Mixture Density Network handle stochasticity?

It predicts multiple Gaussian components

\sum_i \pi_i \mathcal{N}(\mu_i, \sigma_i^2)

to represent multimodal futures instead of averaging to a single incorrect prediction.

Why use short rollouts instead of long ones in model-based RL?

Model errors compound exponentially—short

k

-step rollouts keep predictions accurate, then fall back to model-free value estimates for long-term planning.

What is the difference between model-free and model-based RL?

Model-free learns policy/value directly from trial-and-error; model-based explicitly learns environment dynamics

P(s'|s,a)

first, then plans or trains policy using the model.

Concept Map

Hinglish (regional understanding)

Intuition Hinglish mein samjho

Dekho, yaha pe sabse important cheez ye samajhna hai ki world model kya karta hai. Normal reflex jaisa AI hota hai jo bas "agar X dikhe to Y karo" seekh leta hai — bina samjhe ki actually ho kya raha hai. Par world model iske ulta kaam karta hai: ye apne andar ek chhota sa "physics engine" ya mental simulation banata hai, jisse wo predict kar sakta hai ki "agar main Y action karunga, to next state kya hoga aur kitna reward milega." Formally isko hum $p(s_{t+1}, r_t | s_t, a_t)$ likhte hai — matlab current state aur action dekh kar agla state aur reward guess karna. Embodied AI ka matlab hai ki agent ke paas ek body hoti hai (real ya simulated), aur usse seekhne ke liye actually kuch karna padta hai, sirf text padhna kaafi nahi.

Ab ye matter kyu karta hai? Do bade reasons hai. Pehla — sample efficiency. Real world me interaction mehenga hota hai; robot toot sakta hai, ya medical trial me zindagi ka risk hota hai. Agar agent apne andar hi rollouts simulate kar le (apne model me), to wo bina real damage ke practice kar sakta hai. Dusra — planning. Jaise chess player pehle mind me chaal socha karta hai, waise hi model hone se agent future trajectories ko simulate karke best action choose kar sakta hai commit karne se pehle. Isi ko Model Predictive Control (MPC) kehte hai, jaha hum aisa action sequence dhoondte hai jo simulated future me maximum total reward de.

Technically, kyunki model stochastic hota hai, har future state random hota hai, isliye hum saare possible futures pe unki probability ke hisaab se average lete hai (expectation). Practically hum $K$ trajectories sample karte hai, unke returns ka average nikalte hai, aur actions ko optimize karte hai — ya to gradient descent se (agar model differentiable ho) ya shooting methods jaise CEM se. Bottom line ye samajh lo: world model se hum reality me galti karne se pehle "dimag me" hi try karke dekh lete hai kaunsa action best hai. Yahi cheez ise plain model-free RL (jaise DQN, PPO) se zyada powerful aur transferable banati hai, kyunki ek achha world model environment ki physics capture kar leta hai jo alag-alag tasks me kaam aati hai.