In model-based reinforcement learning, the agent actively constructs its own understanding of the environment's mechanics. Instead of solely relying on trial-and-error to learn a policy or value function directly, the agent first builds a model of the environment. This model typically consists of two main components:

Transition Dynamics Model: This predicts the next state, $s'$ , given the current state, $s$ , and the action taken, $a$ . Formally, it approximates the true environment transition probability distribution $P(s' | s, a)$ .
Reward Function Model: This predicts the immediate reward, $r$ , received after taking action $a$ in state $s$ and transitioning to state $s'$ . It approximates the true reward function, which might be represented as $R(s, a)$ or $R(s, a, s')$ .

Learning these models transforms the RL problem, at least partially, into a supervised learning problem. The agent gathers experience tuples $(s_t, a_t, r_{t+1}, s_{t+1})$ through interaction with the actual environment. This collected data serves as the training set for learning the transition and reward models.

Learning the Transition Model

The complexity of learning the transition model, $\hat{P}_{\theta}(s' | s, a)$ (where $\theta$ represents the model parameters, often neural network weights), depends heavily on whether the environment is deterministic or stochastic, and the nature of the state space.

Deterministic Transitions

In environments where taking action $a$ in state $s$ always leads to the exact same next state $s'$ , the task simplifies significantly. The goal is to learn a function $\hat{s}' = f_{\theta}(s, a)$ that directly predicts the next state.

Approach: This is a standard regression problem.
Function Approximator: Neural networks (like Multi-Layer Perceptrons, MLPs) are commonly used. If the state $s$ includes spatial information (like images), Convolutional Neural Networks (CNNs) might be employed. For sequences or history dependence, Recurrent Neural Networks (RNNs) or Transformers could be suitable.
Training Data: Pairs of $((s_t, a_t), s_{t+1})$ from collected experience.
Loss Function: A common choice is the Mean Squared Error (MSE) between the predicted next state $\hat{s}'_{t+1} = f_{\theta}(s_t, a_t)$ and the actual observed next state $s_{t+1}$ : $L(\theta) = \frac{1}{N} \sum_{i=1}^{N} || f_{\theta}(s_i, a_i) - s_{i+1} ||^2$ Where $N$ is the number of data points in a batch or the dataset.

Stochastic Transitions

When the environment is stochastic, the same $(s, a)$ pair can lead to different next states $s'$ according to the probability distribution $P(s' | s, a)$ . Learning the model now means approximating this entire distribution, not just a single outcome.

Approach: This requires learning a conditional probability distribution $\hat{P}_{\theta}(s' | s, a)$ .
Function Approximator: Again, neural networks are the standard tool, but their output needs to represent a distribution. Common techniques include:
- Categorical Distribution: If the state space is discrete (or discretized), the network can output a probability for each possible next state $s'$ . The loss function is typically the Cross-Entropy between the predicted distribution and a one-hot encoding of the observed $s_{t+1}$ .
- Gaussian Mixture Models (GMMs): For continuous state spaces, the network can output the parameters (means $\mu_k$ , standard deviations $\sigma_k$ , and mixture weights $\pi_k$ ) of a GMM conditioned on $(s, a)$ . The model predicts the next state distribution as $\hat{P}_{\theta}(s' | s, a) = \sum_k \pi_k(s, a) \mathcal{N}(s' | \mu_k(s, a), \sigma_k(s, a))$ . The loss function is the Negative Log-Likelihood (NLL) of the observed $s_{t+1}$ under this predicted distribution: $L(\theta) = - \sum_{i=1}^{N} \log \hat{P}_{\theta}(s_{i+1} | s_i, a_i)$
- Flow-based Models or Other Generative Models: For very high-dimensional state spaces like images, more advanced generative models (Normalizing Flows, VAEs, GANs) can be trained to model $P(s'|s,a)$ , often requiring specialized architectures and training procedures.
Training Data: Tuples of $(s_t, a_t, s_{t+1})$ from experience.

The diagram below illustrates the supervised learning setup for training both models:

Data tuples $(s_t, a_t, r_{t+1}, s_{t+1})$ collected from the environment are used to train separate transition and reward networks via supervised learning. The networks predict the next state (or its distribution) and reward based on the current state and action. These predictions are compared to the actual outcomes to compute loss signals, which are then used to update the network parameters ( $\theta$ and $\phi$ ).

Learning the Reward Model

Learning the reward model, $\hat{R}_{\phi}(s, a, s')$ or sometimes simplified to $\hat{R}_{\phi}(s, a)$ , is typically more straightforward than learning the transition dynamics. Rewards are often scalar values, making this a standard regression task.

Approach: Predict the expected reward $\hat{r} = g_{\phi}(s, a)$ or $\hat{r} = g_{\phi}(s, a, s')$ .
Function Approximator: A separate neural network (MLP often suffices) is common.
Training Data: Pairs of $((s_t, a_t), r_{t+1})$ or $((s_t, a_t, s_{t+1}), r_{t+1})$ from collected experience.
Loss Function: MSE is frequently used: $L(\phi) = \frac{1}{N} \sum_{i=1}^{N} ( g_{\phi}(s_i, a_i) - r_{i+1} )^2$ (or using $g_{\phi}(s_i, a_i, s_{i+1})$ if the reward depends on the next state).

Considerations in Model Learning

Data Collection: The quality and coverage of the data $(s, a, r, s')$ significantly impact the learned model's accuracy. Data might be collected using a random policy, the current learned policy, or a mixture. The agent's exploration strategy directly influences the data available for model training.
Model Bias: The learned model $\hat{P}_{\theta}, \hat{R}_{\phi}$ will almost always be an imperfect approximation of the true environment dynamics. Errors in the model can compound during planning, potentially leading to suboptimal or even detrimental plans if the agent relies too heavily on a flawed model. This is known as model bias.
Model Uncertainty: It's often useful for the agent to know how confident it is in its model's predictions. Areas of the state-action space that haven't been visited frequently will likely have less accurate model predictions. Techniques like using Bayesian neural networks or ensembles of models can help quantify this uncertainty, which can then be incorporated into the planning process (e.g., preferring to plan through well-understood state-action regions).
Computational Cost: Training complex dynamics models, especially for high-dimensional state spaces or stochastic environments, can be computationally intensive.

Once these models are trained, they unlock the ability to perform planning. The agent can use the learned $\hat{P}_{\theta}$ and $\hat{R}_{\phi}$ to simulate potential trajectories, evaluate sequences of actions, and update its policy or value function without necessarily requiring further interaction with the real environment, as we will see in subsequent sections discussing Dyna-Q and MCTS integration.