Imagination-Augmented Agent (I2A)

The I2A architecture is built upon a standard model-free reinforcement learning agent, such as an A3C (Asynchronous Advantage Actor-Critic). This backbone provides the primary policy and value function that are updated using real environment experience. The 'imagination' components are auxiliary modules that augment, rather than replace, this core learning loop.

• Primary Role: Directly maps observations to actions and value estimates.
• Training Signal: Learns from real reward sequences and advantages.
• Foundation: The world model's imagined rollouts provide enriched context for this backbone's decisions.

At the heart of I2A is a learned world model—a neural network that predicts future environment states and rewards given the current state and a sequence of actions. This model is trained in a supervised fashion on the agent's real experience to minimize prediction error.

• Function: Encodes environment dynamics: f(s_t, a_t) -> (s_{t+1}, r_{t+1}).
• Purpose: Serves as an internal simulator for planning.
• Key Property: It is trained jointly with the policy, allowing the imagination to specialize in trajectories relevant to the task.

This component executes 'imagination' rollouts using the learned world model. Starting from the current state, it samples multiple action sequences (from the agent's policy or a rollout policy) and uses the world model to simulate potential futures. The outcomes of these rollouts are then encoded into a fixed-size vector.

• Process: Generates multiple simulated trajectories.
• Output: A compact embedding summarizing the imagined futures (e.g., expected rewards, final states).
• Integration: This embedding is concatenated with the real observation and fed into the model-free backbone's policy and value networks.

The standard actor (policy) and critic (value) networks of the model-free backbone are modified to accept an augmented input. This input is the concatenation of the raw environment observation and the imagination embedding from the rollout encoder.

• Architectural Change: Network Input = [Observation, Imagination Embedding].
• Result: The policy can condition its decisions not just on the present, but on simulated consequences.
• Benefit: Enables foresight and more strategic action selection, especially for sparse or delayed reward tasks.

During imagination, the agent needs to sample action sequences to simulate. The rollout policy determines which actions to try in the internal model. This can be:

• The Agent's Own Policy: Using the current policy for rollouts is simple and ensures imagination is relevant to learned behavior.
• A Separate, Trained Policy: A dedicated rollout policy can be optimized specifically for effective exploration within the model.
• Random or Heuristic Search: Methods like shotgun sampling can be used for broad coverage.

The choice of rollout policy critically affects the quality and diversity of the imagined futures.

I2A is trained end-to-end with a composite loss function that includes terms for both the model-free agent and the world model:

• Model-Free RL Loss: Standard policy gradient (e.g., A3C loss) computed on real experience.
• World Model Loss: Supervised prediction loss (e.g., MSE) for next-state and reward prediction.
• Gradient Flow: Gradients from the RL loss flow back through the imagination embedding and the world model, encouraging the world model to learn dynamics useful for decision-making.
• Key Challenge: Balancing the training of the potentially unstable world model with the policy learning.

At the heart of I2A is a learned environment model or world model. This is a neural network trained to predict future states and rewards given the current state and a sequence of proposed actions. It enables the agent to perform internal rollouts or 'imagine' potential futures without interacting with the real environment. This model is typically trained on the agent's past experience, learning a compressed latent state representation that facilitates efficient prediction.

This component processes the outcomes of the imagination rollouts. For each imagined trajectory, it aggregates the sequence of predicted latent states and rewards into a fixed-length context vector. This encoding summarizes the potential consequences of different action sequences. The rollout encoder's output is then fed as additional input to the model-free core, enriching its decision-making context with foresight derived from the world model.

The I2A architecture retains a standard model-free reinforcement learning agent (e.g., an A2C or A3C agent) as its foundation. This core agent consists of a policy network and a value network. Its key modification is that its input is augmented. Instead of receiving only the current state observation, it receives the concatenation of the current state and the context vector from the rollout encoder, allowing it to base decisions on both immediate perception and imagined futures.

The final policy π(a|s, g) is a function of the raw state s and the aggregated imagination encodings g. This allows the agent to:

• Evaluate long-term consequences of short-term actions.
• Plan over multiple horizons by varying the depth of imagination rollouts.
• Improve sample efficiency because the world model is trained on all collected experience, allowing the policy to learn from simulated data. The policy learns to interpret the rollout encodings to bias its actions toward imagined successful outcomes.

The environment model is trained unsupervised on the agent's collected experience buffer. Using a reconstruction loss, it learns to encode raw observations (e.g., pixels) into a latent state. It is also trained with a dynamics prediction loss to accurately predict the next latent state and reward given the current latent state and action. This is often framed as a variational inference problem, similar to training a Variational Autoencoder (VAE) for states.

I2A differs from traditional Model-Based RL (MBRL) in a key way: it does not use the world model for explicit trajectory optimization (like Model-Predictive Control (MPC)). Instead, it uses the model to generate informative features for a model-free policy. This hybrid approach is more robust to model bias—errors in the learned world model—because the model-free core can learn to ignore misleading imagination outputs, whereas a pure planner would fail catastrophically.

A world model is a learned or engineered internal representation of an environment that enables an agent to predict future states and outcomes without direct interaction. In an I2A, this model is the core 'imagination' engine used to generate synthetic rollouts.

• Function: Serves as a forward dynamics model, predicting the next state and reward given a current state and action.
• Architecture: Often implemented as a Recurrent State-Space Model (RSSM) or similar latent dynamics model.
• Purpose: Provides a compact, abstract simulation for planning, drastically improving sample efficiency compared to pure model-free learning.

Model-Based Reinforcement Learning (MBRL) is a paradigm where an agent learns an explicit model of the environment's dynamics and uses it for planning or policy optimization. The I2A architecture is a specific, influential instance of MBRL.

• Core Idea: Decouples learning the environment (the model) from learning the policy.
• Advantage: Enables planning via internal simulation, reducing the need for expensive real-world interactions.
• I2A's Role: I2A enhances a model-free policy with model-based rollouts, creating a hybrid that leverages the strengths of both approaches.

A Recurrent State-Space Model (RSSM) is a specific world model architecture that combines a deterministic recurrent neural network with stochastic latent variables. It is the foundational model used in agents like Dreamer and conceptually underpins the I2A's imagination module.

• Components: Encodes past observations into a deterministic hidden state and a stochastic latent state.
• Benefit: Learns a temporally coherent and compact state representation ideal for long-horizon rollouts.
• Output: Predicts future rewards and observations (or their latent equivalents), enabling imagination-based planning.

Model-Predictive Control (MPC) is an online planning method that uses a dynamics model to simulate multiple potential action sequences, selects the best one, and executes only the first action before replanning. I2A's imagination rollouts are a form of learned, neural MPC.

• Process: At each step, the agent runs multiple imagination rollouts from its current state using its world model.
• Outcome: The outcomes of these rollouts are aggregated to form a richer feature vector for the policy and value networks.
• Contrast: Unlike classical MPC which optimizes actions directly, I2A uses rollouts to inform a learned policy network.

Model-Based Value Expansion (MBVE) is a technique closely related to I2A's mechanism. It uses short-horizon rollouts from a learned dynamics model to generate more accurate target values (Q-targets) for training a model-free Q-function.

• Mechanism: The model is used to 'expand' a single real experience into a short simulated trajectory.
• Goal: Produces a better bootstrap target for Temporal Difference (TD) learning, reducing bias.
• Relation to I2A: I2A can be seen as a more integrated architecture where imagination features directly condition the policy, whereas MBVE focuses specifically on value function training.

State representation is the process of encoding raw, high-dimensional observations (like pixels) into a compact, informative vector sufficient for decision-making. The quality of this representation is critical for the I2A's world model and rollouts.

• Challenge: Transforming pixels or sensor data into a Markovian state where the future is predictable given the current state.
• I2A's Approach: Uses the latent state of its world model (e.g., from an RSSM) as its state representation.
• Importance: A good representation enables effective long-horizon imagination and robust planning by abstracting away irrelevant details.