In previous chapters you studied motion planning, reinforcement learning, and sim-to-real transfer. Embodied AI synthesizes these ideas around the notion of an agent with a physical body whose actions change the world and whose sensors provide partial observations. A world model is an internal model that predicts how the world evolves as the robot acts, enabling reasoning in imagination rather than only through real executions.

We formalize the robot–world interaction as a partially observable Markov decision process (POMDP)

\[ \mathcal{M} = (\mathcal{S}, \mathcal{A}, \mathcal{O}, p, \Omega, r, \gamma) \]

where \( \mathcal{S} \) is the (unobservable) world state space (including robot configuration and environment state), \( \mathcal{A} \) is the action space, \( \mathcal{O} \) is the observation space, \( p(s' \mid s, a) \) is the state transition kernel, \( \Omega(o \mid s) \) is the observation model, \( r(s, a) \) is the reward, and \( \gamma \in (0,1) \) is the discount factor.

A policy is a mapping \( \pi_\theta(a_t \mid h_t) \) from history \( h_t = (o_0, a_0, \dots, o_{t}) \) to actions. In embodied AI, world models aim to compress the history into a latent state \( z_t \) sufficient for decision-making.

flowchart TD
  W["Physical world state s_t"] --> B["Robot body"]
  B --> SENS["Sensors"]
  SENS --> ENC["Encoder: history -> latent z_t"]
  ENC --> WM["World model: predict next z and observation"]
  WM --> POL["Policy: choose action a_t from z_t"]
  POL --> ACT["Actuators"]
  ACT --> W
        

Conceptually, the world model internalizes the dynamics \( p(s' \mid s, a) \) and observation process \( \Omega(o \mid s) \), but often in a lower-dimensional latent state that captures task-relevant information more compactly than the raw robot and environment state.

World models in modern robotics are typically implemented as probabilistic latent dynamical systems. Instead of modeling \( p(s_{t+1}, o_t \mid s_t, a_t) \) directly in the high-dimensional physical state space, we introduce latent variables \( z_t \in \mathbb{R}^d \) and factorize the model as

\[ p_\theta(o_{1:T}, z_{1:T} \mid a_{1:T-1}) = p(z_1) \prod_{t=1}^T p_\theta(o_t \mid z_t)\, p_\theta(z_t \mid z_{t-1}, a_{t-1}) , \]

where \( p_\theta(z_t \mid z_{t-1}, a_{t-1}) \) is the latent transition model and \( p_\theta(o_t \mid z_t) \) is the decoder that generates observations from latent states. When rewards are learned together with the world model, we also specify

\[ p_\theta(r_t \mid z_t, a_t), \]

which allows planning in latent space without directly querying the real robot for rewards at every imagined step.

In practice, we often assume Gaussian distributions parameterized by neural networks, for example

\[ p_\theta(z_t \mid z_{t-1}, a_{t-1}) = \mathcal{N}\bigl(z_t \,;\, \mu_\theta(z_{t-1}, a_{t-1}),\, \Sigma_\theta(z_{t-1}, a_{t-1})\bigr), \]

and similarly for \( p_\theta(o_t \mid z_t) \). The parameters \( \theta \) are learned from trajectories collected by the robot (either via random exploration, a hand-designed policy, or a partially trained controller).

Direct maximum likelihood learning of \( p_\theta(o_{1:T} \mid a_{1:T-1}) \) is generally intractable because it requires marginalizing over all latent sequences \( z_{1:T} \). Instead, we introduce an approximate posterior \( q_\phi(z_{1:T} \mid o_{1:T}, a_{1:T-1}) \) (the inference model or encoder) and optimize a variational lower bound.

Starting from the log marginal likelihood,

\[ \log p_\theta(o_{1:T} \mid a_{1:T-1}) = \log \int p_\theta(o_{1:T}, z_{1:T} \mid a_{1:T-1}) \,\mathrm{d}z_{1:T}, \]

we multiply and divide integrand by \( q_\phi(z_{1:T} \mid o_{1:T}, a_{1:T-1}) \) and apply Jensen's inequality:

\[ \begin{aligned} \log p_\theta(o_{1:T} \mid a_{1:T-1}) &= \log \int q_\phi(z_{1:T} \mid o_{1:T}, a_{1:T-1}) \frac{p_\theta(o_{1:T}, z_{1:T} \mid a_{1:T-1})}{q_\phi(z_{1:T} \mid o_{1:T}, a_{1:T-1})} \,\mathrm{d}z_{1:T} \\ &\ge \int q_\phi(\cdot) \log \frac{p_\theta(o_{1:T}, z_{1:T} \mid a_{1:T-1})}{q_\phi(z_{1:T} \mid o_{1:T}, a_{1:T-1})} \,\mathrm{d}z_{1:T} \\ &\equiv \mathcal{L}(\theta, \phi; o_{1:T}, a_{1:T-1}), \end{aligned} \]

where \( \mathcal{L}(\theta, \phi; \cdot) \) is the evidence lower bound (ELBO). Using the factorization of the generative model and a factorized inference model \( q_\phi(z_{1:T} \mid o_{1:T}, a_{1:T-1}) = \prod_{t=1}^T q_\phi(z_t \mid z_{t-1}, o_t, a_{t-1}) \), the ELBO decomposes into reconstruction and regularization terms:

\[ \mathcal{L} = \mathbb{E}_{q_\phi} \Biggl[ \sum_{t=1}^{T} \log p_\theta(o_t \mid z_t) \Biggr] - \sum_{t=1}^{T} \mathrm{KL}\!\bigl( q_\phi(z_t \mid z_{t-1}, o_t, a_{t-1}) \,\|\, p_\theta(z_t \mid z_{t-1}, a_{t-1}) \bigr). \]

The first term encourages latent states to retain information useful for reconstructing observations; the second term regularizes latent dynamics so that the learned posterior remains close to the prior transition dynamics. In practice, both terms can be estimated by Monte Carlo samples from \( q_\phi \) and optimized via stochastic gradient methods.

When rewards are modeled, we add another reconstruction term \( \log p_\theta(r_t \mid z_t, a_t) \), which encourages latent states to be predictive of task performance, not just sensory observations.

Once a world model is learned, the robot can plan in the latent space instead of repeatedly querying the real environment. Let \( z_0 \sim q_\phi(z_0 \mid h_0) \) be the encoded initial latent state. Given a parametric policy \( \pi_\psi(a_t \mid z_t) \), we define the imagined rollouts

\[ z_{t+1} \sim p_\theta(z_{t+1} \mid z_t, a_t), \quad a_t \sim \pi_\psi(a_t \mid z_t), \]

and consider the model-based objective

\[ J_{\text{model}}(\psi) = \mathbb{E} \Biggl[ \sum_{t=0}^{H-1} \gamma^{t} r_\theta(z_t, a_t) \Biggr], \]

where the expectation is taken with respect to the latent dynamics and the policy. Here \( r_\theta(z_t, a_t) \) is either a learned reward model or a known function of the latent state and action (for example if we encode joint configurations and velocities explicitly into \( z_t \)).

If the world model exactly matches the true environment dynamics and rewards, then the optimal policy in latent space coincides with the optimal policy in the real world (for horizon \( H \)). More formally:

Theorem (Perfect model equivalence). Suppose there exists an injective mapping \( f : \mathcal{S} \rightarrow \mathcal{Z} \) such that \( z_t = f(s_t) \) and for all states and actions

\[ p_\theta(z' \mid z, a) = p\bigl( f(s') \mid f(s), a \bigr) \quad\text{and}\quad r_\theta(z, a) = r(s, a), \]

where \( s = f^{-1}(z) \). Then for any policy \( \pi(a \mid z) \) we have

\[ J_{\text{model}}(\pi) = J_{\text{env}}(\pi \circ f), \]

where \( J_{\text{env}} \) is the expected return in the original POMDP with fully observed state \( s_t \), and \( \pi \circ f \) is the induced policy on original states.

Sketch of proof. By injectivity of \( f \), trajectories in \( \mathcal{S} \) and \( \mathcal{Z} \) correspond one-to-one via \( z_t = f(s_t) \). The assumed equality of transition kernels and rewards implies that the joint distribution over trajectories in latent space under \( p_\theta, r_\theta, \pi \) matches exactly the distribution induced in state space under \( p, r, \pi \circ f \). Therefore the expected discounted sum of rewards agrees for every policy.

In practice, world models are imperfect. Model-predictive control (MPC) mitigates imperfections by replanning at each real time step using short-horizon imagined rollouts and receding horizon optimization, e.g. with random shooting, cross-entropy methods, or gradient-based trajectory optimization in latent space.

World models for robots must capture structure particular to embodiment: rigid-body dynamics, contact, friction, and sensor geometry. Rather than learning arbitrary dynamics, we can inject inductive biases:

• Structured latent states. Split \( z_t = (x_t, e_t) \) into robot configuration \( x_t \) (joint positions and velocities) and environment latent factors \( e_t \). Known kinematic chains are encoded analytically in \( x_t \), while harder-to-model aspects (such as deformable objects) are represented in \( e_t \).
• Equivariant models. If the robot operates in Euclidean space, dynamics and observations should be equivariant with respect to rigid motions of the workspace. This can be expressed as

  \[ g \cdot z_{t+1} = f_\theta(g \cdot z_t, u_t) \quad \text{for transformations } g \text{ in a symmetry group}, \]

  where \( g \cdot z \) denotes the group action on latent states.
• Information-theoretic objectives. We can encourage the latent state to be both predictive and compact by maximizing mutual information between latent states and future observations:

  \[ I(z_t ; o_{t+1:t+K}) = H(o_{t+1:t+K}) - H(o_{t+1:t+K} \mid z_t) , \]

  where \( K \) controls the prediction horizon.

These structures make world models more data-efficient and better aligned with robot embodiment than unconstrained deep networks, which might otherwise waste capacity modeling irrelevant visual texture.

We now construct a tiny world model for a 1D point-mass robot. The physical dynamics are

\[ x_{t+1} = x_t + \Delta t \, v_t, \quad v_{t+1} = v_t + \Delta t \, a_t, \]

with action \( a_t \) (acceleration command). We approximate this with a neural world model in latent space \( z_t = (x_t, v_t) \) and train it to predict \( z_{t+1} \) from \( (z_t, a_t) \). For simplicity, we use mean-squared error as a surrogate for Gaussian negative log-likelihood.

import torch
import torch.nn as nn
import torch.optim as optim

# Simple fully connected world model: (z_t, a_t) -> predicted z_{t+1}
class WorldModel(nn.Module):
    def __init__(self, state_dim=2, action_dim=1, hidden_dim=64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim + action_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, state_dim)
        )

    def forward(self, z_t, a_t):
        inp = torch.cat([z_t, a_t], dim=-1)
        return self.net(inp)

# Generate synthetic data from known physics for supervised training
def generate_trajectories(num_traj=256, horizon=50, dt=0.05):
    z_list = []
    a_list = []
    z_next_list = []

    for _ in range(num_traj):
        x = torch.zeros(1)
        v = torch.zeros(1)
        for t in range(horizon):
            a = torch.randn(1) * 0.5
            x_next = x + dt * v
            v_next = v + dt * a

            z = torch.stack([x, v], dim=-1)          # shape (1, 2)
            z_next = torch.stack([x_next, v_next], dim=-1)
            z_list.append(z)
            a_list.append(a.view(1, 1))
            z_next_list.append(z_next)

            x, v = x_next, v_next

    z = torch.cat(z_list, dim=0)
    a = torch.cat(a_list, dim=0)
    z_next = torch.cat(z_next_list, dim=0)
    return z, a, z_next

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
wm = WorldModel().to(device)
optimizer = optim.Adam(wm.parameters(), lr=1e-3)

z, a, z_next = generate_trajectories()
dataset = torch.utils.data.TensorDataset(z, a, z_next)
loader = torch.utils.data.DataLoader(dataset, batch_size=128, shuffle=True)

for epoch in range(50):
    total_loss = 0.0
    for z_batch, a_batch, z_next_batch in loader:
        z_batch = z_batch.to(device)
        a_batch = a_batch.to(device)
        z_next_batch = z_next_batch.to(device)

        pred = wm(z_batch, a_batch)
        loss = ((pred - z_next_batch) ** 2).mean()

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * z_batch.size(0)

    total_loss /= len(dataset)
    if epoch % 10 == 0:
        print("epoch", epoch, "mse", total_loss)

# Imagination rollout from an initial latent state under a fixed policy
def imagine_rollout(wm, z0, policy, horizon=30):
    z_traj = [z0]
    a_traj = []
    for t in range(horizon):
        with torch.no_grad():
            z_t = z_traj[-1]
            a_t = policy(z_t)
            z_next = wm(z_t, a_t)
        a_traj.append(a_t)
        z_traj.append(z_next)
    return torch.stack(z_traj, dim=0), torch.stack(a_traj, dim=0)

# Simple proportional policy: accelerate toward target position x_star
def proportional_policy(x_star=1.0, k_p=1.0):
    def policy(z_t):
        x = z_t[:, 0:1]
        error = x_star - x
        a = k_p * error
        return a
    return policy

z0 = torch.tensor([[0.0, 0.0]], device=device)
policy = proportional_policy()
z_traj, a_traj = imagine_rollout(wm, z0, policy, horizon=40)

print("imagined final position", z_traj[-1, 0].item())
      

This example omits stochasticity and reward modeling, but it illustrates the core pattern: learn a dynamics model from experience, then perform imagined rollouts for policy search or MPC without interacting with the real system at every step.

In a high-performance robotics stack, world models are often implemented in C++ to integrate with existing real-time control code. Below is a minimal interface using Eigen for vector algebra. The model is kept abstract so that it can be implemented by linear models, Gaussian processes, or learned neural networks.

#pragma once
#include "Eigen/Dense"

class WorldModel {
public:
    using State = Eigen::VectorXd;
    using Action = Eigen::VectorXd;

    virtual ~WorldModel() {}

    // Predict mean next state from current state and action
    virtual State predictNextState(const State& z,
                                   const Action& u) const = 0;

    // Optionally, return covariance of prediction for uncertainty-aware MPC
    virtual Eigen::MatrixXd predictCovariance(const State& z,
                                              const Action& u) const = 0;
};

// A simple linear world model: z_{t+1} = A z_t + B u_t
class LinearWorldModel : public WorldModel {
public:
    LinearWorldModel(int state_dim, int action_dim)
        : A_(State::Zero(state_dim)),
          B_(Eigen::MatrixXd::Zero(state_dim, action_dim)) {}

    void setA(const Eigen::MatrixXd& A) { A_ = A; }
    void setB(const Eigen::MatrixXd& B) { B_ = B; }

    State predictNextState(const State& z,
                           const Action& u) const override {
        return A_ * z + B_ * u;
    }

    Eigen::MatrixXd predictCovariance(const State&,
                                      const Action&) const override {
        // For a deterministic linear model, covariance is zero.
        return Eigen::MatrixXd::Zero(A_.rows(), A_.rows());
    }

private:
    Eigen::MatrixXd A_;
    Eigen::MatrixXd B_;
};
      

A model-predictive controller may use this interface to repeatedly predict the consequences of candidate action sequences and select the best one according to a cost function over predicted latent trajectories.

For robotics systems built on Java (for instance, certain mobile platforms or educational robots), we can define simple world model and policy interfaces that mirror the RL abstraction while remaining independent of learning details.

public interface WorldModel {
    // Predict next latent state z_{t+1} from z_t and action a_t
    double[] predictNext(double[] z, double[] a);
}

public interface Policy {
    // Choose action a_t given z_t
    double[] act(double[] z);
}

public class LinearWorldModel implements WorldModel {
    private final double[][] A;
    private final double[][] B;

    public LinearWorldModel(double[][] A, double[][] B) {
        this.A = A;
        this.B = B;
    }

    @Override
    public double[] predictNext(double[] z, double[] a) {
        int n = A.length;
        double[] next = new double[n];
        for (int i = 0; i != n; ++i) {
            double sum = 0.0;
            for (int j = 0; j != z.length; ++j) {
                sum += A[i][j] * z[j];
            }
            for (int k = 0; k != a.length; ++k) {
                sum += B[i][k] * a[k];
            }
            next[i] = sum;
        }
        return next;
    }
}

// Example proportional policy toward target position
public class ProportionalPolicy implements Policy {
    private final double xStar;
    private final double kP;

    public ProportionalPolicy(double xStar, double kP) {
        this.xStar = xStar;
        this.kP = kP;
    }

    @Override
    public double[] act(double[] z) {
        double x = z[0];
        double error = xStar - x;
        double a = kP * error;
        return new double[]{a};
    }
}
      

A higher-level planner can roll out trajectories by repeatedly calling predictNext and act in a loop, then evaluate task-specific costs over the imagined trajectory.

MATLAB provides built-in support for state-space models and model predictive control, which makes it convenient to prototype world-model-based controllers for robots that can be approximated by linear dynamics around an operating point.

% State x = [position; velocity], input u = acceleration
A = [1, 0.05;
     0, 1];
B = [0.5 * 0.05^2;
     0.05];
C = eye(2);
D = zeros(2, 1);

sys = ss(A, B, C, D, 0.05);   % discrete-time state-space model

% Define MPC object using this world model
predictionHorizon = 20;
controlHorizon = 5;
mpcobj = mpc(sys, 0.05, predictionHorizon, controlHorizon);

% Weights on output tracking and input effort
mpcobj.Weights.OutputVariables = [1, 0.1];
mpcobj.Weights.ManipulatedVariablesRate = 0.01;

% Simulate closed-loop control toward reference position x_star
x0 = [0; 0];
yref = [1; 0];   % target position 1, zero velocity
v = [];          % no measured disturbance
options = mpcsimopt(mpcobj);
options.PlantInitialState = x0;

[y, t, u] = sim(mpcobj, 60, yref, v, options);

% y contains imagined states under the world model, u the control sequence
figure;
subplot(2,1,1);
plot(t, y(:,1));
xlabel('time step');
ylabel('position');
subplot(2,1,2);
plot(t, u);
xlabel('time step');
ylabel('acceleration');
      

In Simulink, the same world model can be represented by a State-Space block connected in feedback with an MPC Controller block, enabling hardware-in-the-loop testing and rapid prototyping of embodied controllers.

Mathematica is useful for symbolic analysis of world model objectives. Consider a one-dimensional linear Gaussian world model

\[ z_{t+1} \sim \mathcal{N}(a z_t + b u_t,\; \sigma^2), \]

and a target (true) model

\[ z_{t+1} \sim \mathcal{N}(a^\star z_t + b^\star u_t,\; (\sigma^\star)^2). \]

We can symbolically compute the KL divergence between these Gaussians and its gradient with respect to \( a \) and \( b \).

Clear[a, b, aStar, bStar, sigma, sigmaStar, z, u];

mu     = a z + b u;
muStar = aStar z + bStar u;

kl = 1/2 Log[(sigma^2)/(sigmaStar^2)]
     + (sigmaStar^2 + (muStar - mu)^2)/(2 sigma^2)
     - 1/2;

(* Gradients of KL divergence with respect to model parameters a and b *)
gradA = D[kl, a] // FullSimplify
gradB = D[kl, b] // FullSimplify

(* Optionally, substitute particular values and simplify *)
gradAeval = gradA /. {z -> 1.0, u -> 0.5,
                      aStar -> 1.0, bStar -> 0.0,
                      sigmaStar -> 0.1, sigma -> 0.1} // N
gradBeval = gradB /. {z -> 1.0, u -> 0.5,
                      aStar -> 1.0, bStar -> 0.0,
                      sigmaStar -> 0.1, sigma -> 0.1} // N
      

Such symbolic derivations help verify the objectives implemented in code and provide theoretical insight into how model mismatch penalizes parameter errors in world models.

Problem 1 (ELBO Decomposition for World Models). Consider the latent world model \( p_\theta(o_{1:T}, z_{1:T} \mid a_{1:T-1}) \) with factorization

\[ p_\theta(o_{1:T}, z_{1:T} \mid a_{1:T-1}) = p(z_1) \prod_{t=1}^T p_\theta(o_t \mid z_t)\, p_\theta(z_t \mid z_{t-1}, a_{t-1}). \]

Assume a variational posterior \( q_\phi(z_{1:T} \mid o_{1:T}, a_{1:T-1}) \). Show that the ELBO can be written as

\[ \mathcal{L} = \sum_{t=1}^{T} \mathbb{E}_{q_\phi} \bigl[ \log p_\theta(o_t \mid z_t) \bigr] - \sum_{t=1}^{T} \mathbb{E}_{q_\phi} \Bigl[ \mathrm{KL}\bigl( q_\phi(z_t \mid \cdot) \,\|\, p_\theta(z_t \mid z_{t-1}, a_{t-1}) \bigr) \Bigr] + \\ \log p(z_1) - \mathbb{E}_{q_\phi} \bigl[ \log q_\phi(z_1 \mid \cdot) \bigr], \]

where \( q_\phi(z_t \mid \cdot) \) abbreviates the appropriate conditioning on history. Interpret each term.

Solution. Starting from the definition,

\[ \mathcal{L} = \mathbb{E}_{q_\phi} \bigl[ \log p_\theta(o_{1:T}, z_{1:T} \mid a_{1:T-1}) - \log q_\phi(z_{1:T} \mid o_{1:T}, a_{1:T-1}) \bigr], \]

and substituting the factorization yields

\[ \begin{aligned} \mathcal{L} &= \mathbb{E}_{q_\phi} \Biggl[ \log p(z_1) + \sum_{t=1}^{T} \log p_\theta(o_t \mid z_t) + \sum_{t=1}^{T} \log p_\theta(z_t \mid z_{t-1}, a_{t-1}) - \log q_\phi(z_{1:T} \mid \cdot) \Biggr]. \end{aligned} \]

Assume \( q_\phi(z_{1:T} \mid \cdot) = q_\phi(z_1 \mid \cdot) \prod_{t=2}^T q_\phi(z_t \mid z_{t-1}, \cdot) \). Grouping terms for each time step gives the stated decomposition: reconstruction terms \( \mathbb{E}_{q_\phi}[\log p_\theta(o_t \mid z_t)] \) and KL regularizers between the variational transitions and the prior dynamics, plus a prior and posterior term for \( z_1 \). Intuitively, the ELBO trades off reconstruction accuracy against consistency with the dynamical prior.

Problem 2 (Bias from Model Error). Let the true environment dynamics be \( p^\star(s' \mid s, a) \), and suppose the learned world model \( p_\theta(s' \mid s, a) \) differs from the truth. Consider the one-step value estimate

\[ \hat{V}(s) = \mathbb{E}_{a \sim \pi(\cdot \mid s)} \mathbb{E}_{s' \sim p_\theta(\cdot \mid s, a)} \bigl[ r(s, a) + \gamma V(s') \bigr]. \]

Show that the Bellman bias at state \( s \) equals

\[ \operatorname{Bias}(s) = \hat{V}(s) - T^\pi V(s) = \gamma \mathbb{E}_{a \sim \pi(\cdot \mid s)} \Bigl[ \mathbb{E}_{s' \sim p_\theta(\cdot \mid s, a)}[V(s')] - \mathbb{E}_{s' \sim p^\star(\cdot \mid s, a)}[V(s')] \Bigr]. \]

Solution. The Bellman operator for policy \( \pi \) is

\[ T^\pi V(s) = \mathbb{E}_{a \sim \pi(\cdot \mid s)} \mathbb{E}_{s' \sim p^\star(\cdot \mid s, a)} \bigl[ r(s, a) + \gamma V(s') \bigr]. \]

Subtracting from \( \hat{V}(s) \) gives

\[ \begin{aligned} \hat{V}(s) - T^\pi V(s) &= \mathbb{E}_{a \sim \pi} \Bigl[ \mathbb{E}_{s' \sim p_\theta}[r(s, a) + \gamma V(s')] - \mathbb{E}_{s' \sim p^\star}[r(s, a) + \gamma V(s')] \Bigr] \\ &= \gamma \mathbb{E}_{a \sim \pi} \Bigl[ \mathbb{E}_{s' \sim p_\theta}[V(s')] - \mathbb{E}_{s' \sim p^\star}[V(s')] \Bigr], \end{aligned} \]

since the one-step reward \( r(s, a) \) cancels. Thus the bias is proportional to the discrepancy between the world model and true dynamics as seen through the value function \( V \).

Problem 3 (Linear Latent Dynamics and Stability). Suppose the latent dynamics of a world model are linear and deterministic,

\[ z_{t+1} = A z_t + B u_t, \]

with a linear feedback policy \( u_t = K z_t \). Show that the closed loop dynamics can be written as \( z_{t+1} = (A + B K) z_t \) and derive a sufficient condition for mean-square stability in terms of eigenvalues of \( A + B K \).

Solution. Substituting the policy into the dynamics gives

\[ z_{t+1} = A z_t + B K z_t = (A + B K) z_t. \]

This is a linear time-invariant system. A sufficient condition for asymptotic stability (and hence mean-square stability under bounded disturbances) is that all eigenvalues of \( A + B K \) have magnitude strictly less than one. Then \( (A + B K)^t \rightarrow 0 \) as \( t \rightarrow \infty \), implying \( z_t \rightarrow 0 \) for any initial condition.

Problem 4 (Information Bottleneck for Latent States). Consider a single-step observation \( o_t \) and latent state \( z_t \). The information bottleneck objective is

\[ \mathcal{J} = I(z_t ; o_{t+1}) - \beta I(z_t ; o_t), \]

with \( \beta > 0 \). Explain qualitatively why this encourages latent states to discard irrelevant information about \( o_t \) while retaining predictive information about \( o_{t+1} \), and how this is beneficial for embodied world models.

Solution. The term \( I(z_t ; o_{t+1}) \) encourages latent states that are maximally informative about the future observation; increasing it requires that \( z_t \) capture aspects of \( o_t \) that influence \( o_{t+1} \) through the dynamics. The penalty \( \beta I(z_t ; o_t) \) discourages \( z_t \) from preserving unnecessary details of the current observation. For embodied robots, this means that latent states focus on task-relevant quantities (such as object poses, contact states, and robot configuration) while ignoring nuisance factors like lighting or background texture, leading to more robust and data-efficient world models.

Problem 5 (Training Loop as an Algorithmic Pipeline). Sketch an algorithmic pipeline for training a world model from robot experience and then using it for planning with MPC.

Solution (pipeline):

flowchart TD
  D["Collect trajectories from robot"] --> FIT["Fit world model parameters"]
  FIT --> EVAL["Evaluate prediction error on validation data"]
  EVAL -->|model good enough| PLAN["Plan actions in latent space with MPC"]
  EVAL -->|model poor| D
  PLAN --> EXEC["Execute first action on robot"]
  EXEC --> D
        

The robot continuously alternates between data collection, world model training, and model-based planning. When prediction error degrades, additional experience is gathered to update the model before deploying new plans.

In this lesson we framed embodied AI as sequential decision-making for a physical agent and introduced world models as probabilistic latent dynamical systems that approximate environment dynamics and rewards. We derived the variational ELBO for learning such models, discussed their use in model-based planning and control, and analyzed how imperfect models introduce Bellman bias. We also emphasized the role of embodiment-specific inductive biases and demonstrated minimal implementations of world models in Python, C++, Java, MATLAB/Simulink, and Mathematica. These concepts form the basis of many current research directions in long-horizon, data-efficient, and generalist robotic autonomy.