Transition Model

At its core, a transition model is a learned function, typically parameterized by a neural network, that approximates the environment's true dynamics. It takes the current state (s) and an action (a) as input and outputs a prediction of the next state (s'). Formally, it learns $p(s' | s, a)$.

• Input: (State, Action) pair.
• Output: Predicted next state or a distribution over possible next states.
• Purpose: This function forms the agent's internal 'simulator,' allowing it to 'imagine' the future.

Transition models can be architected to capture different levels of environmental uncertainty.

• Deterministic Model: Outputs a single, predicted next state. Simpler and faster but cannot represent stochastic environments or its own uncertainty.
• Probabilistic Model: Outputs a distribution over possible next states (e.g., mean and variance of a Gaussian). This is more expressive and enables uncertainty-aware planning.
• Ensemble Models: A common robust approach uses an ensemble of multiple deterministic or probabilistic networks. Disagreement among the ensemble members provides a practical estimate of epistemic uncertainty (model uncertainty).

The primary technical challenge in MBRL is that the learned model is always an imperfect approximation. Model error—the discrepancy between predicted and true dynamics—is inevitable and problematic.

• Compounding Error: In multi-step imagined rollouts, small errors accumulate, causing predictions to diverge rapidly from reality. This makes long-horizon planning unreliable.
• Mitigation Strategies: Algorithms address this through:
  • Short planning horizons (e.g., in Model Predictive Control).
  • Regular re-planning from the true current state.
  • Using model rollouts primarily for policy optimization (like in MBPO) rather than direct action selection.
  • Uncertainty quantification to avoid exploiting flawed predictions.

For high-dimensional observations like images, learning dynamics directly in pixel space is extremely difficult.

• Pixel Space Model: Predicts future raw observations (pixels). Often high-variance and computationally expensive.
• Latent Dynamics Model: Learns to encode the high-dimensional observation into a compact latent state representation. The transition model then predicts in this latent space. This is far more efficient and improves generalization.
• Example: The Recurrent State-Space Model (RSSM) in the Dreamer algorithm uses a stochastic latent variable and a deterministic RNN to model temporal dependencies, enabling effective learning from pixels.

A transition model is useless without a planning algorithm to leverage it. The model provides the 'physics' for internal simulation, while the planner searches for optimal actions.

• Model Predictive Control (MPC): An online planner that uses the model to simulate multiple action sequences over a finite horizon, selects the best one, executes the first action, and then replans.
• Backpropagation Through Time (BPTT): Used in algorithms like Dreamer. The policy is trained via gradient descent on imagined rollouts, treating the model as a differentiable simulation.
• Monte Carlo Tree Search (MCTS): Used in algorithms like MuZero. The model is queried to simulate trajectories that guide a tree search for the best action.

A sophisticated paradigm shift where the model's purpose is not to reconstruct the true environment state, but to be accurate for planning. The MuZero algorithm learns a value-equivalent model.

• Predicts Planning-Relevant Quantities: Instead of $p(s' | s, a)$, it learns to predict future reward, policy (action probabilities), and value function directly.
• Abstract States: The model's internal 'state' is a latent representation that is sufficient for accurate planning but may not correspond to the true environmental state.
• Key Advantage: The model is optimized directly for the downstream task of finding good policies, which can be more sample-efficient than learning accurate pixel-to-pixel dynamics.

A world model is a comprehensive internal representation that encompasses both a transition model (for dynamics) and a reward model. It allows an agent to simulate and evaluate entire future trajectories—states, actions, and rewards—internally, enabling planning and imagination without direct environment interaction. This is the foundational concept behind algorithms like Dreamer.

Model Predictive Control (MPC) is an online planning algorithm that directly utilizes a transition model. At each step, it:

• Uses the model to simulate multiple action sequences over a finite planning horizon.
• Selects the sequence that maximizes expected reward.
• Executes only the first action from this optimal sequence.
• Repeats the process from the new state. This receding horizon control approach is robust to model inaccuracies and is widely used in robotics and process control.

Model error is the discrepancy between a learned transition model's predictions and the true environment dynamics. This error is a primary challenge in MBRL. Compounding error occurs when this inaccuracy accumulates over the course of a multi-step imagined rollout, leading the model to simulate states that are increasingly unrealistic. Managing these errors through uncertainty quantification and robust planning is critical for deploying MBRL agents successfully.

A latent dynamics model learns to predict future states in a compressed, abstract representation space (the latent space), rather than in the raw, high-dimensional observation space (e.g., pixels). This approach, used in algorithms like Dreamer with its Recurrent State-Space Model (RSSM), offers significant advantages:

• Improved generalization by learning essential features.
• Computational efficiency for planning.
• Better handling of partial observability through recurrent states.

Uncertainty quantification involves estimating the epistemic (model) and aleatoric (environmental) uncertainty in a transition model's predictions. Accurate uncertainty estimates are essential for:

• Robust planning: Avoiding actions in states where the model is highly uncertain (e.g., pessimistic exploration).
• Guided exploration: Actively seeking out states with high uncertainty to improve the model. Common technical approaches include Bayesian Neural Networks (BNNs) and probabilistic ensembles of models.

Sample efficiency measures the number of real environment interactions an agent requires to learn a high-performing policy. It is the primary claimed advantage of Model-Based Reinforcement Learning (MBRL). By learning a transition model, the agent can generate vast amounts of imagined rollouts (synthetic experience) to train its policy internally, drastically reducing the need for costly, slow, or dangerous real-world trials compared to model-free methods.