TL;DR
Adding recurrent architectures (like LSTMs) to off-policy reinforcement learning typically carries massive computational overheads due to repeatedly backpropagating through long sequences of convolutional feature extractions. The authors propose RISE (Recurrent Integration via Simplified Encodings), a dual-stream architecture that feeds precomputed, non-learnable visual embeddings (e.g., from a frozen pre-trained ResNet-18) to the recurrent context stream while running an expensive learnable CNN only on the current frame. This setup boosts Atari-57 human-normalized interquartile mean (IQM) scores by 35.6% while remaining computationally accessible on a single desktop GPU.
What Problem Does This Paper Address?
Off-policy Deep Reinforcement Learning (RL) algorithms like Rainbow and Beyond The Rainbow (BTR) achieve excellent sample efficiency and performance. However, scaling them to handle partial observability and long-term credit assignment often requires recurrent neural networks (RNNs/LSTMs).
In standard recurrent off-policy RL (e.g., R2D2), calculating a single Q-value update over a history window requires passing each historical frame through an expensive, learnable convolutional neural network (CNN) encoder. Prior solutions rely on learning from full sequential trajectories with "burn-in" periods. This wastes compute passes on non-updated steps, forces large batch sizes, and induces temporal correlation across training batches. Consequently, high-performance recurrent off-policy RL has historically required large compute clusters, remaining out of reach for researchers with desktop-scale hardware.
Methods
RISE decouples the representation needed for temporal context from the representation needed for fine-grained current-state evaluation:
- Dual-Stream Architecture:
- Recurrent (Context) Stream: Historical observations ($o_{t-k}, \dots, o_t$) pass through a fixed, non-learnable encoder $\varphi_\theta^\dagger$ (default: penultimate layer of an ImageNet-pretrained ResNet-18). Because this encoder is frozen, its embeddings can be computed once during environment interaction, stored in the replay buffer, and retrieved instantaneously without recomputation. The embedding sequence is fed through an LSTM.
- Learnable (Current) Stream: Only the immediate observation ($o_t$) passes through the standard learnable CNN encoder $\phi$.
- Stream Combination: The LSTM output is projected via a linear layer to match the dimension of the CNN features, processed by a sigmoid activation, and multiplied element-wise (Hadamard product) with the learnable CNN embedding—acting as an attention gating mechanism before reaching the value/advantage heads.
- Transition-Based Sampling: Unlike trajectory-based RNN algorithms (R2D2), RISE samples single transitions $(s_t, a_t, r_t, s_{t+1})$ alongside historical precomputed vectors. It only computes two online/target CNN passes per sample ($2b$ forward passes per batch), avoiding temporally correlated trajectory updates.
Datasets
The authors evaluated RISE across four benchmark suites:
- Atari 2600 (Atari-57 and Atari-5 / Atari-10 subsets): Standard discrete action visual benchmark evaluated at 200M frames with sticky actions.
- Procgen Benchmark: 16 procedurally generated visual environments tested on the "hard" distribution.
- ViZDoom: 7 first-person 3D navigation and shooting tasks with heavy partial observability.
- MiniWorld: 8 3D goal-oriented navigation environments designed around partial observability.
No new benchmarks were created.
Key Insights and Ideas
- Pre-trained representations are sufficient for memory: The recurrent unit does not need backpropagated task-specific gradients through the entire historical visual pipeline to retain useful temporal context. Frozen representations from out-of-domain vision models (ImageNet classification) provided enough signal for tracking long-term dependencies.
- Resistance to representation staleness: Even when using an agent's own earlier, non-updated convolutional embeddings ("stale" embeddings), performance only dropped slightly compared to frozen ResNet features.
- Recurrence choices matter less than having recurrence: LSTMs outperformed or matched modern sequence architectures like Transformers, FFM, and MAMBA (State Space Models) on RL context lengths ($k = 160$), while being significantly faster in wall-clock time.
- Qualitative strategy shifts: Recurrent context allowed agents to solve long-horizon credit assignment failures: for instance, timing oxygen pickups in NameThisGame, learning non-suicidal multi-life play in Phoenix and BattleZone, or avoiding circular head-turning traps in ViZDoom HealthGathering.
Key Results
- Atari-57 Performance: Integrating RISE into Beyond The Rainbow (BTR) improved the human-normalized interquartile mean (IQM) from 5.918 to 8.299 (+35.6%) and median score from 4.199 to 6.134. BTR + RISE outperformed human performance on 50 out of 57 games.
- Walltime & Compute Efficiency: BTR + RISE completed 200M frames on Atari in ~20 hours on a single consumer desktop GPU (RTX 4090) (and ~32 hours on an A100). Compared to standard full-CNN LSTM training on BTR, RISE achieved a 22.9% performance gain while slashing walltime by 84%.
- Comparison to Heavy Baselines: BTR + RISE reached an IQM comparable to leading distributed architectures (MEME: 9.131, Dreamer-v3: 9.268) while requiring a fraction of the compute walltime (MEME used 274 hours on 2$\times$TPU-v4; Dreamer-v3 used 185 hours on an A100).
Limitations and Future Work
- Domain Dependency: On Procgen and specific Atari games that do not rely heavily on memory or credit assignment, RISE yielded negligible improvements over base BTR.
- Fixed Pre-Trained Encoder Bias: If a frozen encoder fails to extract features relevant to a specific domain, the recurrent stream cannot compensate, trading off domain adaptability for training throughput.
- Limited Scope: The experimental validation concentrated primarily on image-based inputs, leaving extensions to other modalities (e.g., text, audio) to future investigation.