[Paper Review] Bootstrap Latent-Predictive Representations for Multitask Reinforcement Learning

Daniel Guo, Bernardo Ávila Pires+5 more|arXiv (Cornell University)|Apr 30, 2020

Reinforcement Learning in Robotics38 references42 citations

TL;DR

The paper introduces Predictions of Bootstrapped Latents (PBL), a self-supervised representation learning method for multitask deep RL that predicts future latent embeddings and uses a bootstrapping cycle with latent-to-state and state-to-latent predictions, improving performance on DMLab-30 and Atari-57.

ABSTRACT

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.

Motivation & Objective

• Motivate improved representation learning for multitask, partially observable RL settings.
• Develop a self-supervised auxiliary task that focuses on predicting latent embeddings of future observations.
• Introduce a bootstrapping mechanism between latent observations and agent states to enrich representations.
• Enable multimodal observation integration by operating entirely in latent space.
• Empirically validate PBL against state-of-the-art baselines on DMLab-30 and Atari-57.

Proposed method

• Define Z_t as a latent embedding of observations using a learned encoder f(O_t).
• Forward prediction: predict Z_{t+k} from compressed partial histories B_{t,k} using a predictor g; minimize ||g(B_{t,k})-Z_{t+k}||^2 for horizons k=1..K.
• Reverse prediction: predict compressed history B_t from latent Z_t using a predictor g'; minimize ||g'(f(O_t))-B_t||^2.
• Train forward and reverse predictors jointly to form a bootstrapping cycle that avoids trivial solutions.
• Use two RNNs: h_f for full histories and h_p for partial histories to compute B_t and B_{t,k}.
• Adopt the PopArt-IMPALA RL base, with a larger architecture for improved performance, and subsample timesteps for efficiency.

Experimental results

Research questions

• RQ1 Does PBL improve multitask RL performance over existing representation learning methods in DMLab-30 and Atari-57?
• RQ2 How does the prediction horizon affect PBL performance and what is the role of reverse prediction in learning meaningful latent representations?
• RQ3 Is PBL stable against collapse to trivial representations, and how does architecture choice influence results?
• RQ4 Can PBL encodings capture cross-task shared structure and generalize to unseen tasks?

Key findings

• PBL outperforms pixel control, CPC, and DRAW as auxiliary representation tasks in the DMLab-30 multitask setting.
• Performance improves as the forward horizon increases, with diminishing returns, and multistep prediction is beneficial over single-step.
• Removing reverse prediction (latent targets random) still benefits from longer horizons, indicating forward predictions alone are insufficient without meaningful latent targets; reverse prediction helps learn useful latent structure.
• PBL remains stable and does not collapse to trivial solutions; using random projections in the latent path still yields competitive results, suggesting robust training dynamics.
• In Atari-57, PBL improves median human-normalized score across tasks, indicating generality beyond DMLab-30; across tasks, PBL outperforms baselines on several tasks.
• The decoding probe shows PBL's representation encodes object position information better and retains it longer than random-projection baselines.