Source-linked AI summary
Bootstrap Latent-Predictive Representations for Multitask Reinforcement Learning
Daniel Guo, Bernardo Avila Pires, Bilal Piot, Jean-bastien Grill, Florent Altché, Rémi Munos, Mohammad Gheshlaghi Azar
TL;DR
Deep RL needs representations that capture unknown environments, especially across multitask and partially observable settings where predicting complete future observations can be difficult. PBL predicts future latent observation embeddings and bootstraps them through agent-state predictions, enabling latent-space multimodal learning. The paper reports improved performance across DMLab-30 and Atari-57 multitask settings, while identifying transfer and state-information analysis as future work.
Problem
Deep RL needs useful representations for complex multitask and partially observable environments, but accurate multistep prediction of future observations can be difficult.
Method
PBL learns history representations by predicting future latent observation embeddings, while training those embeddings to predict corresponding agent states.
Results
PBL outperformed state-of-the-art representation learning in DMLab-30 and improved performance across multitask DMLab-30 and Atari-57 settings.
Takeaways & Limitations
Meaningful latent representations of future observations benefit multitask agents over longer prediction horizons, and latent-space prediction supports multiple observation modalities.
Takeaways & Limitations
The paper leaves transfer to unseen tasks and the amount of POMDP-state information preserved by PBL representations for future study.
Abstract
from arXiv · showhide
Learning a good representation is an essential component for deep reinforcement learning (RL). Representation learning is especially important in multitask and partially observable settings where building a representation of the unknown environment is crucial to solve the tasks. Here we introduce Prediction of Bootstrap Latents (PBL), a simple and flexible self-supervised representation learning algorithm for multitask deep RL. PBL builds on multistep predictive representations of future observations, and focuses on capturing structured information about environment dynamics. Specifically, PBL trains its representation by predicting latent embeddings of future observations. These latent embeddings are themselves trained to be predictive of the aforementioned representations. These predictions form a bootstrapping effect, allowing the agent to learn more about the key aspects of the environment dynamics. In addition, by defining prediction tasks completely in latent space, PBL provides the flexibility of using multimodal observations involving pixel images, language instructions, rewards and more. We show in our experiments that PBL delivers across-the-board improved performance over state of the art deep RL agents in the DMLab-30 and Atari-57 multitask setting.
1. Introduction
Deep RL representation learning is valuable but difficult in complex, multitask, partially observable environments because accurate multistep prediction can be hard. PBL addresses this by learning history representations through bootstrapped predictions of future latent observations and improves performance in partially observable multitask evaluation.
- Representation learning and auxiliary tasks have improved deep RL performance, including in complex and multitask settings.
- Predicting future observations is a common representation-learning strategy, but accurate multistep prediction may be difficult in partially observable problems.
- PBL learns history representations by predicting future observation embeddings whose training is bootstrapped through predictions of agent states.
- Latent-space bootstrapping can capture contextual information about environment dynamics, incorporate multimodal observations, and use information from the far future.
- In DMLab-30, PBL outperformed pixel control, DRAW, and CPC when used as an auxiliary task for the PopArt-IMPALA agent.
2. Background
The background frames multitask partially observable RL as policy learning under hidden state, where recurrent networks compress histories into agent states. Predictive representations provide a related self-supervised route for learning those states.
- Partially Observable Environments: A POMDP models hidden states, actions, observations, transitions, rewards, and discounting, with multitask dynamics or rewards potentially indexed by unobserved state.
- Partially Observable Environments: A policy maps observation-action histories to action distributions, and RL seeks to maximize expected discounted rewards.
- Recurrent History Representations: Recurrent networks compress full histories into agent states that can support decision making, while partial-history networks are used only for future predictions.
- Predictive Representations: Predictive representations learn agent states by predicting statistics of future observations as self-supervised or auxiliary objectives.
- Recurrent History Representations: Figure 1 depicts recurrent history compression, with omitted observation and action-processing networks and dashed links aligning histories in time.
3. Predictions of Bootstrapped Latents (PBL)
PBL uses two coupled auxiliary prediction tasks: compressed histories predict future observation latents, while observation latents predict immediate agent states. Their one-way training creates a bootstrapping cycle intended to enrich representations with structured information about environment dynamics.
- PBL objectives: PBL combines forward action-conditional prediction from compressed partial histories to future latent observations with reverse prediction from latent observations to agent states.These are the method’s two auxiliary prediction tasks.
- Forward prediction: The forward task predicts future observation embeddings Z_t+k = f(O_t+k) from a compressed partial history, using multiple future offsets through a chosen horizon.The prediction horizon ranges from one to a maximum future index.
- Reverse prediction: PBL trains the latent observation embedding network by predicting the learned agent state B_t from Z_t, making latent observations predictive of immediate agent states.This reverse-prediction direction trains the observation latents rather than treating them as fixed random projections.
- Bootstrapping: Forward prediction trains agent states to predict future latent observations, while reverse prediction trains latent observations to predict immediate agent states, forming a bootstrapping cycle.The two objectives optimize disjoint parameter sets and do not pass gradients into prediction targets.
- Bootstrapping: The bootstrapping effect may make observation latents encode structural information about the paths leading to observations and propagate useful information from far into the future.The paper illustrates this with an unlocked-door observation whose latent can encode information about previously finding a key.
4. Experiments
Experiments evaluate PBL across DMLab-30, Atari-57, prediction horizons, architectures, and representation probes. PBL generally outperforms comparison methods, with benefits linked to multistep prediction and structured memory of environment dynamics.
- DMLab 30: PBL outperforms all compared methods throughout DMLab-30 training by a significant margin.Figure 3 reports mean capped human normalized scores averaged across 16 independent runs with 95% confidence intervals.
- DMLab 30: PBL improves performance on most DMLab-30 tasks and reaches parity with pixel control on the remainder.The per-task comparison reports broad gains rather than improvement limited to a small subset of tasks.
- Prediction horizon: Performance improves monotonically with longer forward-prediction horizons, although the gains diminish as the horizon increases.One-step prediction matches pixel control, while multistep prediction provides additional performance gains.
- Prediction horizon: Disabling reverse prediction makes longer horizons ineffective, producing performance slightly below pixel control during training but the same final performance.Without reverse prediction, future latents are random projections, indicating that meaningful latent targets matter in addition to predicting farther ahead.
- Architecture choice: The larger network architecture gives the base RL method a significant boost in overall performance.This architecture was therefore adopted for the reported experiments.
- Atari-57: PBL improves overall Atari-57 performance; it statistically outperforms other methods on seven tasks, and no method statistically outperforms PBL on the others.The benchmark reports median human normalized scores across 57 levels and 95% confidence intervals.
- Decoding the PBL representation: In a 3D-room probe, PBL learns to preserve the red cube’s position after the agent looks away.The probe evaluates whether the representation encodes object position without sending gradients back into the representation.
5. Related Work
Prior work learns predictive representations of future observations for reinforcement learning, including predictive-state, generative, pixel-based, and contrastive approaches. These methods differ in what they predict and how they compress or model future information.
- Predictive state representations compress history into compact statistics that can replace the underlying state for decision-making.They use action-conditional predictions of future events or observations.
- Future-observation prediction is a common strategy for learning representations that support deep reinforcement learning.The surveyed approaches condition predictions on actions and use compressed partial histories to model future observations.
- PreCo and Simcore DRAW use generative likelihood or autoencoding objectives to shape representations of partial and full histories.Simcore DRAW is described as a VAE-based technique that autoencodes observations while compressing histories.
- Pixel control learns predictive representations through Q-learning over changes in future pixels.The method treats control of pixel changes as a representation-learning task.
- CPC learns representations in latent space with a contrastive objective rather than by predicting latent embeddings.This contrasts with PBL’s latent-prediction approach while retaining a lightweight alternative to generative models.
6. Conclusion
PBL addresses representation learning in partially observable, multitask deep RL by learning meaningful future latent observations through bootstrapped predictions. It outperforms a state-of-the-art representation-learning technique in DMLab-30 and suggests extensions involving multimodal inputs, transfer, generative latent models, and representation analysis.
- PBL learns meaningful future latent observations through bootstrapped forward and reverse predictions for partially observable, multitask deep RL.
- PBL outperforms the state-of-the-art representation learning technique in DMLab-30.
- Learning meaningful latent embeddings lets PBL improve as the prediction horizon increases, unlike random projections of future observations.
- Latent-space prediction supports combining multiple observation modalities, including proprioception, natural language, touch, smell, sound, and temperature.
- A generative latent model could extend PBL to stochastic dynamics and generate rollouts for Monte Carlo planning, although it may not help deterministic DMLab30.
- Future work includes evaluating transfer to unseen tasks and studying how much information about the POMDP state predictive representations capture and preserve.
A. Implementation Details
The implementation-details section documents architecture and parameter choices, including hyperparameter sweeps, software libraries, and plotting procedures.
- The section describes implementation details, architecture and parameter choices, hyperparameter-sweep protocols, software libraries, and plotting procedures.
- Hyperparameters selected from sweeps are accompanied by the tested protocol and considered values.
- The implementation uses documented software libraries and a specified plotting protocol.
A.1. Agent Architecture
The appendix specifies the shared agent architecture and auxiliary-learning implementations used to compare PBL with CPC, DRAW, pixel control, and related methods.
- Agent Architecture: The main agent networks process observations, recurrent histories, values, and policies using ResNets, LSTMs, and MLPs, with larger networks than the reference configuration.
- Agent Architecture: Language instructions are embedded word-by-word, processed by an LSTM, and combined with visual observations, previous actions, and rewards.
- Agent Architecture: PBL, CPC, and Simcore DRAW compress partial histories using recurrent states initialized from the full-history agent network.
- Agent Architecture: Training minibatches contain T × B sequences and subsampled compressed histories that generate H-step partial histories, with H = 20 outside scalability experiments.
- Agent Architecture: Representation-loss weights are selected among powers of 10 from −2 to 2, while prediction horizons are generally fixed at 20.
- Agent Architecture: PBL combines action-conditional forward prediction from partial histories to future latent observations with reverse prediction from latent observations to agent states.
- Agent Architecture: PBL's observation encoder outputs a 592-dimensional latent embedding combining first-person vision, language, previous action, and reward information.
- Agent Architecture: CPC discriminates jointly distributed histories and observations from independently sampled pairs, while pixel control uses mean loss over 4 × 4 cells.
A.10. Plotting
The plotting procedure aggregates returns across time bins, tasks, runs, and confidence estimates to produce reported performance summaries.
- Plotting: Per-episode returns are binned into 108-frame intervals and averaged before conversion to capped mean human-normalized scores.
- Plotting: Mean estimates in Figure 4 and Table 9 average performance over the final 500M training frames.
- Plotting: Confidence intervals in Figure 4 are computed by combining the confidence intervals for PBL and pixel control.
B.1. DMLab-30
The DMLab-30 results include task-level performance breakdowns and uncapped mean human normalized scores across methods, PBL prediction horizons, and architectures.
- B.1. DMLab-30: The section reports final performance across all DMLab-30 tasks for every compared method.These results are presented in Table 9.
- B.1. DMLab-30: Mean uncapped human normalized score is compared across methods and shown over training frames.The method comparison uses frames from 0 to 1e10.
- B.1. DMLab-30: PBL is evaluated across different prediction horizons using mean uncapped human normalized score over training frames.The horizon comparison spans 0 to 1e10 frames.
- B.1. DMLab-30: Architecture comparisons are also reported using mean uncapped human normalized score.Figure 13 presents this architecture comparison.
B.2. Atari-57
The Atari-57 results report human normalized scores across tasks during the final 5% of training, with statistically significant improvements shown in bold.
- B.2. Atari-57: Atari-57 task scores are reported for the last 500M of 10B training frames.The table covers the final 5% of training.