Source-linked AI summary
The Predictron: End-To-End Learning and Planning
David Silver, Hado van Hasselt, Matteo Hessel, Tom Schaul, Arthur Guez, Tim Harley, Gabriel Dulac-Arnold, David Reichert, Neil Rabinowitz, Andre Barreto, Thomas Degris
TL;DR
The paper addresses the challenge of learning models that are effective for planning rather than merely reconstructing environments. It introduces the predictron, which rolls an abstract internal model forward and trains accumulated predictions end-to-end. In random mazes and simulated pool, it produced more accurate value predictions than conventional architectures and improved selected pool decisions.
Problem
Model-based reinforcement learning traditionally trains environment models separately from planning, creating a model-planner mismatch.
Method
The predictron is a differentiable architecture that unrolls an abstract Markov reward process over multiple internal planning steps and accumulates rewards and values.
Results
The predictron significantly outperformed conventional model-free deep architectures in procedurally generated mazes and simulated pool, with 27 versus 10 pocketing events in one pool decision procedure.
Takeaways & Limitations
Accurate value predictions can arise from a fully abstract internal model whose rolled-forward plans need not correspond to human-understandable environment states.
Takeaways & Limitations
The paper focuses on value prediction in uncontrolled environments and leaves transfer to control and internal MDPs with abstract actions for future work.
Abstract
from arXiv · showhide
One of the key challenges of artificial intelligence is to learn models that are effective in the context of planning. In this document we introduce the predictron architecture. The predictron consists of a fully abstract model, represented by a Markov reward process, that can be rolled forward multiple "imagined" planning steps. Each forward pass of the predictron accumulates internal rewards and values over multiple planning depths. The predictron is trained end-to-end so as to make these accumulated values accurately approximate the true value function. We applied the predictron to procedurally generated random mazes and a simulator for the game of pool. The predictron yielded significantly more accurate predictions than conventional deep neural network architectures.
1. Introduction
Model-based reinforcement learning traditionally separates environment-model learning from planning, potentially leaving the model mismatched to the agent’s objective. The predictron integrates abstract modeling, planning, and value prediction through end-to-end training, improving predictions in mazes and pool.
- Model-based reinforcement learning learns an environment model and then rolls it forward to estimate cumulative reward and construct a value function.
- Prior models were trained largely independently of their use in planning, while pixel-perfect reconstruction models had not surpassed state-of-the-art model-free methods in challenging raw-input domains.
- The predictron integrates learning and planning end-to-end using a completely abstract internal model optimized for accurate value prediction.
- The architecture predicts multiple value functions for varied pseudo-rewards and discount factors, including events such as staying alive or reaching a room.
- The predictron was evaluated on procedurally generated random mazes and simulated pool directly from pixels, outperforming conventional model-free deep networks and showing greater robustness to depth.
2. Background
The paper frames value prediction in uncontrolled Markov reward processes, where returns are cumulative discounted rewards and value functions are their expected returns. Model-based methods learn an approximate process and evaluate its corresponding value function.
- An MRP maps a state and IID noise to a next state, reward, and discount factor, allowing stochastic transitions.
- An MRP return is the cumulative discounted reward along a trajectory, with rewards, discounts, and returns also generalized to vectors.
- The value function is the expected return from a state and is consistent with an environment when it satisfies the Bellman equation.
- Model-based reinforcement learning learns an approximate MRP that outputs subsequent states, rewards, and discounts, then defines values through the model’s Bellman equation.
3. Predictron architecture
The predictron encodes inputs into an abstract state, repeatedly applies an internal MRP, and accumulates rewards and values across planning depths. Its λ-preturn combines these depth-specific estimates with state-dependent gates in one forward pass.
- Predictron components: The architecture encodes raw input into an internal abstract state before applying its recurrent planning model.
- Predictron components: The internal model maps each abstract state to a subsequent state, reward, and discount, while a value function estimates remaining internal return.
- k-step predictron: A k-step predictron rolls the internal model forward k steps, accumulating internal rewards before adding a discounted value estimate at depth k.
- λ-predictron: The λ-predictron aggregates k-step preturns using weights λ0 through λK, analogous to the forward-view TD(λ) return.
- λ-predictron: All k-step preturns are computed in one forward pass, while the λ-preturn is formed by a forward model sweep followed by a backward accumulator sweep.
- λ-predictron: State-dependent λk gates can truncate computation or use deeper model layers, enabling adaptive planning depth for each internal state.
4. Predictron learning updates
Predictron learning updates train individual and aggregate preturns toward target outcomes, while consistency updates align estimates across imagined planning depths without requiring environmental labels.
- The joint parameters θ are updated so each k-step preturn gk approaches a target outcome g, such as a Monte-Carlo return.The update minimizes mean-squared error between gk and g.
- Individual preturn losses can be averaged uniformly or weighted by each preturn’s usage in the λ-predictron output.Usage weighting emphasizes preturns according to wk, their contribution to the aggregate output.
- The λ-predictron learns parameters η that set weights wk so the aggregate λ-preturn gλ approaches the target outcome g.The joint parameters instead make individual gk values approach the target.
- Planning rollouts generate abstract trajectories with rewards, discounts, and values, enabling refinement of value estimates through imagined trajectories.This extends the planning opportunity associated with model-based reinforcement learning architectures such as Dyna.
- Consistency updates: Consistency updates adjust each preturn gk toward the aggregate estimate gλ while treating gλ as fixed during the update.The parameters θ change to make gk more similar to gλ, not vice versa.
- Consistency updates: Consistency updates require neither labels nor environmental samples, so they can train estimates for hypothetical or otherwise unlabeled states.This enables self-consistency training and is especially relevant to semi-supervised learning.
5. Experiments
Experiments evaluated the predictron on random mazes and simulated pool, testing sequential planning, architectural components, adaptive depth, semi-supervised learning, and decision quality. Across these settings, predictron variants produced accurate predictions and outperformed conventional deep architectures.
- Experimental domains: The experiments used random mazes and a simulated pool domain to evaluate trajectory, connectivity, and multiple future-event predictions.The maze tasks used generated layouts and random starting positions; pool inputs consisted of RGB frames and targeted simultaneous event prediction.
- Learning sequential plans: Near-perfect maze predictions emerged from weighted preturns that decomposed trajectories into sequential planning steps.The model was not given the trajectories’ sequential construction, yet simpler paths used fewer steps while more complex paths used more.
- Exploring the predictron architecture: Adding the MRP structure, λ aggregation, and usage weighting progressively improved aggregated root mean squared prediction errors, especially in pool.The variants differed in whether they learned internal rewards and discounts, used λ weights, and weighted losses by preturn usage.
- Comparing architectures: All (r, γ, λ)-predictrons outperformed corresponding feedforward and recurrent baselines across random mazes and pool, at every tested depth.This held with and without skip connections and with shared or separate core weights.
- Adaptive depth: Prediction depth varied by prediction type and input, correlated with real-world discount for several pool events, and outperformed a scalar shared λ control.The scalar λ control reduced performance in all scenarios, indicating that heterogeneous depth provided useful flexibility.
- Using predictions to make decisions: Selecting pool initial conditions with predictron predictions pocketed 27 coloured balls in 50 episodes, compared with 10 for an equally deep convolutional network.The selection considered 128 angles and two velocities, giving 256 candidate initial conditions.
6. Related work
The predictron relates to prior architectures for adaptive computation, value learning, predictive representations, and abstract modeling, while differing in how it learns and combines predictions.
- λ-predictron weights are learned conditionally on the input, unlike hand-tuned branching weights, and its intermediate predictions are combined into an aggregate output.
- Adaptive computation time aggregates outputs across pondering steps, whereas the λ-predictron can use different pondering depths for each prediction.
- VINs plan over the full input state space with convolutions, while the predictron plans through imagined trajectories in an abstract state space.
- Predictive state representations, general value functions, and nexting study many future predictions, but prior work had not applied them to learning abstract models.
- Unlike separately trained or unsupervised abstract-model approaches, the predictron trains its abstract model end-to-end for accurate value prediction.
7. Conclusion
The conclusion presents the predictron as a differentiable architecture that plans with a fully abstract internal model and combines predictions across planning depths. The paper demonstrates more accurate value predictions than conventional architectures, while leaving control settings for future work.
- The predictron rolls forward an internal model to estimate external values without requiring the model to correspond to the real environment in human-understandable terms.
- Its architecture accumulates rewards, produces values at multiple depths, combines them into an ensemble, and can enforce self-consistency among predictions.
- The experiments produced more accurate value predictions in reinforcement learning environments than conventional network architectures.
- The study focused on value prediction in uncontrolled environments; extending the ideas to control and internal MDPs remains future work.
A. Architecture
The architecture encodes raw inputs into an abstract state, repeatedly applies a convolutional MRP and λ core, and computes values from the resulting internal states.
- A two-layer convolutional state representation feeds a recurrent core that combines the MRP model and λ-network into a repeatable module.
- A fully connected value network computes v_k from each internal state, while the core can use shared or distinct weights across model steps.
- The core computes next abstract states with convolutions and separately predicts internal rewards, discounts, and lambdas from a shared hidden layer.
- The standard experiments allowed up to 16 model steps, producing 52-layer networks with batch normalization and ReLU nonlinearities between layers.
- The illustrative maze used 6 model steps and 16-channel convolutional feature maps, with compact subnetworks before fully connected layers.
B. Training
Training used supervised Monte-Carlo updates in all experiments except the semi-supervised consistency experiment, with Adam optimizing normalized-return losses.
- All experiments used supervised Monte-Carlo updates except the semi-supervised experiment, which used the consistency update.
C. Comparing architectures of different depths
The experiments test whether predictron advantages persist across network depths, architectural variants, and parameter counts. Predictrons outperform conventional deep networks across these comparisons, with performance depending more on network structure than raw parameter count.
- Depth and skip connections: Predictrons outperform corresponding feedforward and recurrent baselines at all tested depths, with and without skip connections.The comparison covers 2, 4, 8, and 16 model steps, corresponding to 10, 16, 28, and 52 total layers.
- Depth and skip connections: Skip connections make conventional architectures more robust to depth, especially on the pool task.The dashed conventional-network curves nearly overlap across depths, particularly for pool.
- Architectural variants: The alternative unpadded encoder reduces maze and pool abstract states to 16 × 16, making their abstract states equal in size.For mazes, two 3×3 stride-1 convolutions produce 16 × 16 planes; for pool, three 5 × 5 stride-1 convolutions do so.
- Parameter counts: Both shared-core and non-shared-core predictrons outperform every deep network on mazes and pool across the tested parameter settings.This result addresses whether earlier gains were merely an artifact of giving predictrons more parameters.
- Parameter counts: 1.27M vs 4.85M parameters: the 32-node predictron outperforms the 512-node deep network despite the latter having approximately four times as many parameters.The comparison holds in each tested setting, including mazes and pool and both core-sharing conditions.
E.2.1. FIRST TASK
The first task uses randomly generated 20 × 20 mazes with sparse local wall observations and trajectories produced by a fixed policy. Maze generation controls connectivity while randomizing wall placement.
- Task trajectories: The policy observes wall configurations at four adjacent locations and maps each observation to an action for 60 steps from a uniformly random start.The target indicates whether the trajectory traversed each maze location.
- Maze generation: 20 × 20 mazes contain 70% empty locations and 30% walls after uniformly shuffling the wall configuration.The top-left and bottom-right corners are forced to be empty.
- Maze generation: A stochastic line search selects wall counts so the two forced-empty corners are connected in approximately 50% of generated mazes.The walls are then uniformly shuffled.