VAEs in Model-Based Reinforcement Learning

Variational Autoencoders (VAEs) offer powerful mechanisms for learning compact and structured representations of data. This capability extends significantly into the domain of model-based reinforcement learning (MBRL), where understanding and predicting an environment's behavior is central to an agent's success. In MBRL, an agent aims to learn a model of the environment's dynamics, which it can then use for planning or improving its policy. When dealing with high-dimensional state spaces, such as images from a camera feed in a robotics task, learning this world model directly in the observation space can be incredibly challenging and data-intensive.

VAEs provide an elegant solution by first learning a lower-dimensional latent representation of these high-dimensional observations. The core idea is to train a VAE to encode an observation $s_t$ into a latent vector $z_t = q_\phi(z_t | s_t)$ and to decode this latent vector back into a reconstruction of the observation $\hat{s}_t = p_\theta(s_t | z_t)$. The VAE is trained by maximizing the Evidence Lower Bound (ELBO), which encourages good reconstructions while regularizing the latent space, often to approximate a standard Gaussian distribution $p(z)$.

Once a VAE is trained (or while it is being trained), a separate model can be learned to predict the dynamics entirely within this compressed latent space. This latent dynamics model aims to predict the next latent state $z_{t+1}$ given the current latent state $z_t$ and the action $a_t$ taken by the agent:

$$\hat{z}_{t+1} \sim p_{dyn}(z_{t+1} | z_t, a_t)$$

This dynamics model, $p_{dyn}$, can be a neural network, such as an MLP or an RNN if temporal dependencies beyond a single step are important. Learning dynamics in the latent space $z$ rather than the high-dimensional observation space $s$ offers several advantages:

• Efficiency: The latent space is much lower-dimensional, making the dynamics model simpler and faster to train and use for predictions.
• Focus on Salient Features: The VAE, through its information bottleneck, tends to capture the most salient features of the observation in $z$, potentially filtering out noise or irrelevant details. This can make the dynamics easier to learn.
• Abstract Prediction: The model predicts transitions between abstract states, which can be more robust than predicting pixel-level changes.

A typical architecture for an MBRL agent incorporating a VAE involves several components that interact:

Components of a VAE-based model-based RL agent. Observations from the environment are encoded into a latent space. A dynamics model predicts future latent states based on current latent states and actions. A policy or planner uses these latent predictions to choose actions. The VAE's decoder is used primarily for training the encoder to produce useful latent representations.

In this setup:

1. The agent receives an observation $s_t$ from the environment.
2. The VAE's encoder maps $s_t$ to the latent state $z_t$.
3. The policy $\pi(a_t | z_t)$ (or a planner using the learned models) selects an action $a_t$.
4. The action $a_t$ is executed in the environment, leading to a new observation $s_{t+1}$ and reward $r_t$.
5. The transition $(s_t, a_t, r_t, s_{t+1})$ is stored in a replay buffer.
6. The VAE is trained to reconstruct observations: $s_t \approx p_\theta(s_t | q_\phi(z_t | s_t))$.
7. The latent dynamics model $p_{dyn}(z_{t+1} | z_t, a_t)$ is trained to predict the encoding of the next state, $z_{t+1} = q_\phi(z_{t+1} | s_{t+1})$, given $z_t = q_\phi(z_t | s_t)$ and $a_t$. This is often done by minimizing a loss like Mean Squared Error: \mathcal{L}_{dyn} = || \text{target_}z_{t+1} - p_{dyn}(z_t, a_t) ||^2 where \text{target_}z_{t+1} is the "true" next latent state obtained by encoding $s_{t+1}$ (often with gradients stopped).
8. Optionally, a reward predictor $p_{rew}(r_t | z_t, a_t)$ can also be trained.

The learned latent dynamics model, $p_{dyn}$, and reward predictor, $p_{rew}$, can then be used for planning. For instance, an agent can perform "imaginary rollouts" in the latent space by repeatedly applying the dynamics model for a sequence of actions: $z_t \rightarrow \hat{z}_{t+1} \rightarrow \hat{z}_{t+2} \dots$. These imagined trajectories can be used with techniques like Model Predictive Control (MPC) or Monte Carlo Tree Search (MCTS) to select optimal actions. The VAE's decoder can also be used to visualize these imagined latent trajectories by converting them back to observation space, providing a way to "see" what the model is predicting.

Prominent architectures like "World Models" (Ha & Schmidhuber, 2018) and the "Dreamer" family of agents (Hafner et al., 2019, 2020, 2023) are prime examples of this approach.

• World Models explicitly separate the VAE (V component), the latent dynamics model (M component, often an RNN), and a controller (C component). The controller operates using only the latent representations from the VAE and predictions from the M.
• Dreamer and its successors learn behaviors entirely from imagined trajectories generated by a learned world model operating in a compact latent space. The world model, including the representation learner, dynamics model, and reward predictor, is trained jointly. Policy learning then occurs using these imagined trajectories.

Training Considerations and Objectives

The overall training often involves optimizing several objectives concurrently or iteratively:

1. VAE Reconstruction Loss: Ensures $z_t$ captures information to reconstruct $s_t$. This is part of the ELBO: $$\mathcal{L}_{recon} = \mathbb{E}_{q_\phi(z|s)}[\log p_\theta(s|z)]$$
2. VAE KL Regularization: Regularizes the latent space, typically $D_{KL}(q_\phi(z|s) || p(z))$, where $p(z)$ is a standard Normal prior.
3. Latent Dynamics Loss: Minimizes the error in predicting the next latent state, as shown earlier. $$\mathcal{L}_{dyn} = \mathbb{E}_{z_t, a_t, z_{t+1}} [ \text{distance}(z_{t+1}, p_{dyn}(z_t, a_t)) ]$$
4. Reward Prediction Loss (if applicable): $$\mathcal{L}_{rew} = \mathbb{E}_{z_t, a_t, r_t} [ \text{distance}(r_t, p_{rew}(z_t, a_t)) ]$$
5. Policy Loss: The policy is trained to maximize expected rewards, often using actor-critic methods on trajectories generated either from the real environment or from the learned world model.

A careful balance is needed. If the VAE focuses too much on pixel-perfect reconstruction, $z_t$ might retain high-frequency details not relevant for dynamics prediction. Conversely, if $z_t$ is too compressed or too regularized, it might lose information critical for long-term prediction or control. Some approaches modify the VAE objective or architecture to better suit the needs of control, for example, by ensuring that states distinguishable by their future outcomes are also distinguishable in the latent space.

Benefits and Challenges

The integration of VAEs into MBRL brings several benefits:

• Sample Efficiency: By learning a model, the agent can generate additional "imaginary" experiences, reducing reliance on real-world interactions which can be costly or slow.
• Handling High-Dimensional Data: VAEs effectively compress complex sensory inputs like images into manageable latent vectors.
• Planning Capabilities: The learned latent dynamics model enables sophisticated planning algorithms to be applied.

However, there are also challenges:

• Model Accuracy: The learned world model is an approximation. Errors in the latent dynamics model can compound over long prediction horizons, leading to "model-mismatch" where the imagined trajectories diverge significantly from reality.
• Representation Trade-offs: The latent space must be good for both reconstruction (to train the VAE) and for predicting dynamics and rewards. These objectives are not always perfectly aligned. Achieving disentanglement in $z_t$ might be helpful but is not always a primary objective if the resulting representation serves the dynamics learning well.
• Computational Cost: Training a VAE, a dynamics model, and a policy simultaneously can be computationally demanding.

Despite these challenges, VAEs have become a foundational component in many state-of-the-art model-based reinforcement learning agents, particularly those designed to operate from rich sensory inputs. Their ability to learn structured latent representations and facilitate dynamics modeling in a compressed space continues to drive progress in creating more intelligent and data-efficient agents.