Beyond Sliding Windows: Learning to Manage Memory in Non-Markovian Environments

TL;DR

This work tackles memory efficiency for reinforcement learning in real-world, non-Markovian environments by introducing Adaptive Stacking, a meta-algorithm that learns which past observations to retain within a fixed memory budget. It provides convergence guarantees for AS under unbiased value estimates and demonstrates substantial compute and memory reductions for MLP, LSTM, and Transformer agents, compared to standard Frame Stacking. Through extensive experiments on memory-heavy tasks (e.g., TMaze variants) and generalization tests across diverse domains (XorMaze, Rubik's Cube, FetchReach), AS matches oracle baselines while using far less memory and improving generalization. The approach reframes memory management as a reward-guided decision problem, integrating memory eviction into the RL loop and offering practical benefits for scalable, memory-constrained agents in big worlds.

Abstract

Recent success in developing increasingly general purpose agents based on sequence models has led to increased focus on the problem of deploying computationally limited agents within the vastly more complex real-world. A key challenge experienced in these more realistic domains is highly non-Markovian dependencies with respect to the agent's observations, which are less common in small controlled domains. The predominant approach for dealing with this in the literature is to stack together a window of the most recent observations (Frame Stacking), but this window size must grow with the degree of non-Markovian dependencies, which results in prohibitive computational and memory requirements for both action inference and learning. In this paper, we are motivated by the insight that in many environments that are highly non-Markovian with respect to time, the environment only causally depends on a relatively small number of observations over that time-scale. A natural direction would then be to consider meta-algorithms that maintain relatively small adaptive stacks of memories such that it is possible to express highly non-Markovian dependencies with respect to time while considering fewer observations at each step and thus experience substantial savings in both compute and memory requirements. Hence, we propose a meta-algorithm (Adaptive Stacking) for achieving exactly that with convergence guarantees and quantify the reduced computation and memory constraints for MLP, LSTM, and Transformer-based agents. Our experiments utilize popular memory tasks, which give us control over the degree of non-Markovian dependencies. This allows us to demonstrate that an appropriate meta-algorithm can learn the removal of memories not predictive of future rewards without excessive removal of important experiences. Code: https://github.com/geraudnt/adaptive-stacking

Figures (30)

• Figure 2: Illustration of Frame Stacking (FS) and Adaptive Stacking (AS) with $k=4$ in the passive-TMaze with $k^*=L+2=5$. In order to free up space to push the new observation $x_t$, the FS agent always pops the last observation ($i_t=0$) while the AS agent chooses which observation to pop ($i_t \in \{1, \dots, k\}$). In this figure, corresponds to an empty memory slot.
• Figure 3: Experimental domains ($10^6$ training steps for all). (a-b) are as described in Section \ref{['sec:AdaptiveStacking']}. (c) Episodic ($T=100$) with similar observations as the Active-TMaze, but the agent must go to the XOR of the two goal cues. (d) Episodic ($T=100$) where the agent only sees one cube face at a time. (e) Episodic ($T=1445$) with a similar goal as the Active-TMaze. (f) Continual with hidden joint velocities. The goal position also changes every $50$ steps and is only observable for the first two steps after each change (Assumption \ref{['ass:partial-order']} does not hold). See Appendix \ref{['apdx:domains']} for detailed descriptions.
• Figure 4: Continual TMazes ($T=10^6$) with Q-learning ($N_{rs}=20$). AS matches the oracle FS($k^*$) in returns and memory usage, while outperforming FS($\kappa$) especially for long-term dependencies.
• Figure 5: Episodic Passive-TMaze ($k^*=T=L+2$) with PPO using an MLP ($N_{rs}=10$). AS learns to retain the goal cue regardless of maze and memory lengths, unlike FS with the same $k$.
• Figure 6: Domain-agnostic results, showing performance during training ($N_{rs}=10$ except when stated otherwise) and testing ($N_{rs}=100$). The shaded regions show 95% confidence intervals.

Theorems & Definitions (11)

• Remark 1
• Definition 1
• Proposition 1
• Theorem 1
• Theorem 2: Partial-order Preserving
• Proposition 1
• Theorem 1
• Corollary 2: Convergence of Episodic Monte Carlo Policy Gradient
• Corollary 3: Convergence of Average-Reward Monte Carlo Policy Gradient
• Theorem 4: Partial-order Preserving