Source-linked AI summary
Imagination-Augmented Agents for Deep Reinforcement Learning
Théophane Weber, Sébastien Racanière, David P. Reichert, Lars Buesing, Arthur Guez, Danilo Jimenez Rezende, Adria Puigdomènech Badia, Oriol Vinyals, Nicolas Heess, Yujia Li, Razvan Pascanu, Peter Battaglia, Demis Hassabis, David Silver, Daan Wierstra
TL;DR
Deep reinforcement learning needs agents that are data-efficient, flexible across tasks, and robust when learned environment models are imperfect. I2As combine model-free decisions with imagined trajectories and learn to interpret model predictions rather than follow fixed planning procedures. They outperform model-free baselines, including on Sokoban, and can use imperfect models, while trading additional computation for fewer environment interactions.
Problem
Model-free reinforcement learning often requires substantial data and generalizes poorly, while model-based methods struggle with imperfect learned environment models.
Method
I2As combine a model-free policy path with imagined trajectories from an approximate environment model, learning an encoder to interpret predictions as additional policy context.
Results
I2As outperform model-free baselines on MiniPacman and Sokoban and successfully use imperfect models, including models without reward predictions.
Takeaways & Limitations
I2A broadens the applicability of model-based reinforcement learning by using imperfect models without relying exclusively on simulated returns.
Takeaways & Limitations
I2As trade environment interactions for computation, with computation growing linearly with the number and depth of rollouts.
Abstract
from arXiv · showhide
We introduce Imagination-Augmented Agents (I2As), a novel architecture for deep reinforcement learning combining model-free and model-based aspects. In contrast to most existing model-based reinforcement learning and planning methods, which prescribe how a model should be used to arrive at a policy, I2As learn to interpret predictions from a learned environment model to construct implicit plans in arbitrary ways, by using the predictions as additional context in deep policy networks. I2As show improved data efficiency, performance, and robustness to model misspecification compared to several baselines.
1 Introduction
Model-free reinforcement learning can require substantial data and generalize poorly, while model-based methods promise better efficiency but struggle when learned models are imperfect. I2As address this gap by learning to interpret approximate model predictions and outperforming model-free baselines with less data.
- Model-free agents map raw observations directly to actions or values but usually require large amounts of training data and generalize poorly to novel tasks.
- Model-based reinforcement learning uses learned world models to simulate future outcomes, potentially improving generalization, data efficiency, and computational scalability.
- Model-based agents have had limited success in complex domains because standard planning methods are vulnerable to errors in learned environment models.
- I2As learn end-to-end to interpret imperfect model simulations, using information beyond predicted rewards and without assuming a particular model structure.
- I2As outperform model-free baselines across several domains, including Sokoban, while achieving better performance with less data despite imperfect models.
2 The I2A architecture
I2A combines imagined trajectories from an action-conditional environment model with a model-free processing path. A rollout encoder learns which information from these predictions should inform the final policy and value estimates.
- An environment model predicts future observations and potentially rewards from the current observation or history together with the current action.
- The imagination core combines the environment model with a rollout policy to generate multiple imagined trajectories of predicted observations and rewards.
- Because learned models may produce erroneous predictions, I2A does not rely solely on predicted rewards or values derived from predicted states.
- A rollout encoder processes each imagined trajectory as a whole, extracting useful information or ignoring the trajectory when necessary.
- The policy module combines model-based rollout information with a model-free representation to produce the policy vector and estimated value.
3 Architectural choices and experimental setup
The experiments use action-conditioned rollouts, recurrent trajectory encoding, pretrained environment models, and model-free and copy-model baselines. These choices isolate the value of imagination while controlling architectural and training factors.
- 3.1 Rollout strategy: Experiments perform one rollout for each possible environment action, with subsequent actions generated by a shared rollout policy.
- 3.1 Rollout strategy: The rollout policy can be random or pretrained, but distilling the imagination-augmented policy into a model-free policy was particularly efficient.
- 3.2 I2A components and environment models: The rollout encoder uses a convolutional encoder with an LSTM that processes trajectory features in reverse order, while the aggregator concatenates the resulting summaries.
- 3.2 I2A components and environment models: The environment model can be pretrained and frozen or jointly trained with the agent; pretraining produced faster I2A runtime and was adopted.
- Experimental setup: The environment-model training data comes from trajectories of a partially trained model-free agent, requiring accounting for pretraining and data-generation steps.
- Experimental setup: The baselines include standard model-free networks, a larger-capacity variant, and a copy-model agent that preserves the I2A architecture without imagination.
4 Sokoban experiments
Sokoban experiments show that I2A improves performance and data efficiency over model-free and copy-model baselines, benefits from longer imagined rollouts, and remains robust to inaccurate models. I2A also uses substantially fewer simulations than MCTS and generalizes to levels with more boxes.
- 4.1 I2A performance vs. baselines on Sokoban: I2A reached 85% of levels solved, versus under 60% for the standard model-free baseline and 70% for the increased-capacity baseline.
- 4.1 I2A performance vs. baselines on Sokoban: 15-step rollouts exceeded 90% of levels solved, while longer rollouts generally produced diminishing returns and were slower.Five-step rollouts were selected as canonical because they were substantially faster; the best agents solved levels in about 50 steps on average.
- 4.1 I2A performance vs. baselines on Sokoban: Even including model pretraining, I2A outperformed baselines after about 3e8 total frames.The model-pretraining budget was conservatively estimated below 1e8 frames.
- 4.2 Learning with imperfect models: I2A performed similarly well with a clearly flawed environment model, whereas the encoder-free agent degraded catastrophically with the poor model.The results indicate that the rollout encoder can use earlier predictions while discounting later errors, although the poor model unexpectedly outperformed the good model in one experiment.
- 4.4 Imagination efficiency and comparison with perfect-model planning methods: MCTS required 25k model simulation steps to reach 87% solved, compared with 1.4k environment-model calls for I2A; I2A with Monte-Carlo search reached 95% using 4k steps under a high-accuracy model.MCTS reached 95% only with 100k simulation steps, while the I2A Monte-Carlo variant required a nearly perfect model.
- 4.4 Imagination efficiency and comparison with perfect-model planning methods: An I2A trained on four-box levels solved more than half of seven-box levels, nearly matching the standard agent trained on four-box levels.
5 Learning one model for many tasks in MiniPacman
MiniPacman tests whether one shared environment model can support multiple tasks with common dynamics but different rewards. I2A outperforms baselines across tasks, especially when rewards are sparse.
- Shared model and tasks: A single MiniPacman environment model predicts frames and events shared across five tasks with different reward vectors.The tasks assign rewards to moving, eating food, eating power pills, eating ghosts, and being eaten.
- Results: I2A outperforms the standard agent on every task and the copy-model baseline on all but one task.Baselines and I2As were trained separately for each task, with baselines also receiving the event variable.
- Results: The performance gap is particularly large on tasks 4 and 5, which have sparse rewards and require anticipating ghost dynamics.The authors posit that I2A uses its environment and reward model to explore more effectively.
- Results: After 300 million environment steps, I2A clearly outperforms the other agents on all sparse-reward tasks.The figure reports performance for different agents across all tasks.
6 Related work
Related work has used learned or explicit models for planning, synthetic data, uncertainty handling, and neural planning architectures. I2A differs by learning to interpret imperfect environment-model predictions within reinforcement learning.
- Learned models and planning: Neural model-based RL often combines learned environment models with classical planning, but these studies do not address mismatch between learned and true models.This limitation is central to I2A’s motivation.
- Other model-based settings: Other related work uses given simulation models for policy pretraining, extracts trajectories for imitation learning, or relies on hand-crafted visual abstractions.These settings differ from learning and interpreting low-level environment models for RL.
- Model imperfection: Prior approaches to model imperfection include scheduled-sampling-like techniques, Bayesian uncertainty models, and hierarchies of models with known fidelity.The cited methods address different aspects of imperfect modeling, with uncertainty-based approaches carrying computational or calibration assumptions.
- Synthetic experience: Dyna-style methods use models to create synthetic training data, but the models are not used by the agent at test time.These approaches target data efficiency rather than test-time model-based decision making.
- Neural planning: Neural planning architectures mimic iterative planning without explicit environment models, whereas I2A models predict low-level observations and exploit additional unsupervised signals.I2A therefore differs in both its model representation and its use of model predictions.
- Adjacent architectures: Related architectures query learned expert models or iteratively construct plans, but target contextual bandits or explicit sequential planning rather than I2A’s approach.The general idea of using an internal model in arbitrary ways was also discussed in earlier work.
7 Discussion
I2A combines model-free decisions with learned interpretation of environment-model rollouts, achieving strong performance across MiniPacman and Sokoban. Its benefits come with extra computation, while its rollout efficiency can exceed MCTS even with a perfect model.
- Contribution: I2A combines model-free and model-based ideas by learning to interpret environment models to augment model-free decisions.The approach successfully uses imperfect models, including models without reward predictions.
- Results: I2A outperforms model-free baselines on MiniPacman and the challenging combinatorial domain of Sokoban.The discussion presents this as the paper’s cross-environment performance result.
- Limitation: I2As trade environment interactions for computation because they ponder before acting.In experiments, I2A was always less than an order of magnitude slower per interaction than model-free baselines.
- Planning efficiency: At comparable Sokoban performance, I2A requires far fewer model function calls than MCTS with a perfect environment model.A learned rollout policy guides I2A rollouts toward relevant parts of the state space.
A Training and rollout policy distillation details
The agents use stochastic policies trained with advantage-based policy gradients and value estimates, while optional entropy and policy-distillation losses support exploration and rollout-policy learning.
- Training: A3C updates policy parameters with a policy gradient weighted by an estimated advantage function.The advantage is computed from a bootstrapped k-step return minus the current value estimate.
- Training: The value function is produced by a neural network whose input is the policy network’s second-to-last layer.Its parameters are updated toward the bootstrapped k-step return.
- Training: An entropy regularizer with λent = 10^-2 is added throughout experiments to encourage exploration.
- Rollout policy distillation: When applicable, policy distillation adds cross-entropy between the behavioral policy and rollout policy, scaled by λdist.The rollout-policy gradients are not backpropagated through the behavioral policy in the stated formulation.
B Agent and model architecture details
I2A combines model-free and imagined-information pathways, using convolutional and recurrent components with environment models tailored to MiniPacman and Sokoban.
- Agent architecture: The standard model-free baseline maps the current observation through a convolutional network and fully connected layer to policy and value heads.
- Agent architecture: The I2A model-free path matches the baseline CNN, while a rollout encoder processes model-generated frames and feeds a shared LSTM across five rollouts.For MiniPacman and Sokoban, the LSTM has 256 and 512 units, respectively.
- Agent architecture: The copy-model agent keeps the I2A architecture but replaces its environment model with an identity function returning the input observation.
- Environment models: The environment models are autoregressive next-frame predictors trained from the last observation and a broadcasted one-hot last-action representation.Separate order-1 models are pretrained for raw-pixel MiniPacman and Sokoban observations.
- Environment models: MiniPacman’s model uses a multiscale CNN and pool-and-inject layer for long-range pixel dependencies, whereas Sokoban uses a residual CNN with a reward-prediction pathway.The Sokoban model is nearly local, while MiniPacman must model nonlocal ghost–Pacman interactions.
Task collection
The evaluation spans five MiniPacman tasks sharing environment dynamics but differing in reward structures and level-termination criteria, plus Sokoban with shaped rewards for exploration.
- MiniPacman: Five MiniPacman tasks share maze and movement dynamics while varying rewards and level termination.The task collection includes Regular, Avoid, Hunt, Ambush, and Rush.
- MiniPacman: Regular clears a level when all food is eaten, while Avoid clears after 128 steps.
- MiniPacman: Hunt and Ambush clear when all ghosts are eaten or after 80 steps, whereas Rush clears when all power pills are eaten.
- MiniPacman: A pink shrinking bar displays the remaining duration of Pacman’s power-pill effect, and episodes end when Pacman is eaten because there are no lives.
- Sokoban: Sokoban uses shaping rewards because random actions solve levels with vanishing probability, creating extreme exploration difficulty.
- Sokoban: Sokoban applies -0.1 per step, +1 for pushing a box onto a target, -1 for pushing one off, and +10 for finishing.Levels terminate on completion or after 120 steps, and repeated training or testing levels are rare.
D.2 Additional experiments
Additional experiments examine reward prediction, mental retries, value-network search, and procedural Sokoban generation, extending I2A evaluation beyond the main comparisons.
- Reward prediction: I2A with reward prediction converged shortly after 1e9 steps, while the version without reward prediction reached nearly 80% solved levels after 3e9 steps.
- Mental retries: With a near-perfect Sokoban model, mental retries increased solved levels from around 87% for base I2A to around 95%.Although up to 16 retries were allowed, the performance increase was obtained within the first 10.
- Value-network search: A value network was trained from 1e8 frame–return pairs generated by a model-free agent that solved close to 60% of episodes.Training and test errors were comparable, and more training points were not expected to substantially improve quality.
- Value-network search: During search, the trained value network evaluates leaf nodes, while root actions are selected by maximum Q value and the tree is reused between steps.
- Level generation: Sokoban levels are procedurally generated through topology creation, position configuration, and reverse playing.Topology uses a random walk, then target and player positions are randomly selected before reverse exploration.
- Level generation: Reverse playing explores configurations from boxes placed on targets and returns the room with the highest RoomScore under search limits.Repeated box/player configurations are skipped, and defaults include depth 300 and one million visited positions.