An Empirical Study on the Power of Future Prediction in Partially Observable Environments

TL;DR

The paper investigates whether future observation prediction can yield representations that support reinforcement learning in partially observable environments. By decoupling representation learning from RL and training history encoders via a predictive state loss ($PSR$ loss), the framework $DRL^2$ demonstrates a strong correlation between prediction quality and RL performance across memory-demanding benchmarks, architectures, and tasks. Across grids, memory-intensive sequences, long sequential games, and sparse-reward tasks, decoupled learning often yields faster convergence and more stable learning than end-to-end training, though end-to-end can excel in short-term memory regimes. The work provides practical guidance for leveraging auxiliary future-prediction tasks to improve sample efficiency and offers insights into when decoupled training is advantageous in partially observable domains.

Abstract

Learning good representations of historical contexts is one of the core challenges of reinforcement learning (RL) in partially observable environments. While self-predictive auxiliary tasks have been shown to improve performance in fully observed settings, their role in partial observability remains underexplored. In this empirical study, we examine the effectiveness of self-predictive representation learning via future prediction, i.e., predicting next-step observations as an auxiliary task for learning history representations, especially in environments with long-term dependencies. We test the hypothesis that future prediction alone can produce representations that enable strong RL performance. To evaluate this, we introduce $\texttt{DRL}^2$, an approach that explicitly decouples representation learning from reinforcement learning, and compare this approach to end-to-end training across multiple benchmarks requiring long-term memory. Our findings provide evidence that this hypothesis holds across different network architectures, reinforcing the idea that future prediction performance serves as a reliable indicator of representation quality and contributes to improved RL performance.

Figures (5)

• Figure 1: GridWorld environment with noise in observations.
• Figure 2: Future prediction performance during the burn-in phase (left) and RL performance (right) across different sequential models. In the RepeatPrevious experiment, we compare prediction accuracy using transformers (TF), GRUs, and Amago-GPT as the history summarizing model across varying difficulty levels defined by k. Amago-GPT consistently achieves fast and stable convergence across different k, while TF and GRUs struggle at higher k values. In these cases, E2E fails to learn recall, whereas DRL$^2$ successfully does.
• Figure 3: Temporal credit assignment tasks performance (left), and effects of the update ratio between PSR and RL (right). (Left) DRL$^2$ demonstrate better convergence compared to E2E. (Right) We evaluate the impact of varying update ratios between PSR and RL in the delayed catch environment, and compare its performance with the E2E method. Here, $0.1$x means that for every one update step of RL, PSR is updated only 0.1 times on average, i.e., once every 10 RL updates.
• Figure 4: Performance Comparison of DRL$^2$ and E2E Methods in Noisy MuJoCo Environments. In these MuJoCo environments where E2E-GRU performs well, PSR-TF slightly underperforms relative to E2E-TF.
• Figure 5: Performance of DRL$^2$ on POPGym benchmark. The dashed line represents the best reported on-policy or off-policy method. The x-axis shows timesteps on a million scale.