Train Short, Act Long: Multistability and the Horizon Independence of Recurrent Policies

The gist In reinforcement learning, agents acting in partially observable Markov decision processes (POMDPs) must rely on memory, typically encoded in a recurrent neural network (RNN), to integrate information from past observations, and longhorizon POMDPs present challenges like poor generalization and severe sample inefficiency that require a formal characterization of temporal horizon generalization.

Definition of Temporal Horizon Generalization

Temporal horizon generalization (THG) is defined as the property that a policy remains optimal for all horizons T ≥ 0, meaning a same policy solves the task regardless of the length of the idle phase and this property is connected to planning invariance where agents take the same decisions regardless of whether they are given a distant goal or an intermediate waypoint The necessary and sufficient condition for THG is that rho is a compatible and separating read-out in the sense of Definition 2.2.

Multistability as the Key Dynamical Regime

The analysis reveals that multistability is necessary for temporal horizon generalization and, in simple tasks, sufficient specifically when considering constant idle dynamics where the update reduces to a fixed map on H. A monostable map precludes THG because no policy architecture can achieve both simultaneously, as a monostable model forces the read-out rho to be constant on the basin of attraction, meaning rho cannot distinguish any two initial states, including h0(σ) and h0(σ′) for σ ≠ σ′. Conversely, multistability enables THG when there exists a read-out ρ: H → D that is invariant within each basin and this read-out separates the basins.

Contrast Between RNN Architectures

The paper contrasts different RNN variants based on their dynamical properties regarding THG. Modern parallelizable architectures, such as state space models (SSMs) and linear gated RNNs, are monostable by construction and consequently fail to generalize across temporal horizons because their steady-state solutions are explicit and unique for any input x. In contrast, nonlinear RNNs can exhibit multistability, allowing them to reach THG when the encoding places h0(σ) in a specific basin Bk(σ) and a separating read-out exists.

Synergy of Dynamics for Complex Tasks

For more complex environments like the LookupTreeMaze, the research highlights that transient dynamics and multistability are both essential because even multistable models require transient dynamics to reach equilibrium, which are considerably slower when trained on longer mazes. The authors demonstrate this synergy by showing that an hybrid architecture combining half minGRU and half BMRU combines the capabilities of both minGRU and BMRU: THG with high performance on a small number of T-mazes. This suggests that transient dynamics and multistability are both essential for environments combining sparse observations with frequently updated ones.

Conclusion for Scalable Architectures

The research advocates for a shift away from monolithic parallelizable architectures, suggesting that no current parallelizable RNN exhibits both multistability and transient dynamics and that transient dynamics and multistability are two essential and complementary dynamical regimes for horizon generalization. The proposed direction is designing hybrid parallelizable frameworks that effectively leverage the benefits of transient dynamics alongside multistability.

Limitations and Future Directions

A limitation noted is that reaching these policies during training is not guaranteed, as evidenced by agents trained on long horizons failing on shorter ones. A promising avenue for future research is developing generalization-aware training strategies to fully exploit the ability of multistable models to generalize. Furthermore, extending the analysis to environments with multiple, competing objectives and nested temporal horizons will further clarify their practical significance.

Computation of Variability Among Attractors (VAA)

The variability among attractors (VAA) quantifies multistability by measuring the number of different stable states reached by a model from a set of initial states H under a constant input x. The paper notes that the rigorous definition involves computing limits to infinity, so an approximation is used where M, the number of iterations of f, and ϵ, the tolerance to consider two states as identical are introduced. The VAA is defined mathematically as a ratio involving the Kronecker delta function.

Training Procedure Details

The optimal policies and value functions are approximated using proximal policy optimization (PPO) [32] and an advantage actor-critic [44] method implemented with RNNs, trained using a modified version of the LSTM PPO implementation from CleanRL [45]. The training involves 30 models for each recurrent cell (GRU, BRC, nBRC, minGRU, and BMRU) initialized with different seeds and trained for a total of 250 iterations on the T-maze benchmark. For the LookupTreeMaze experiments, 180 models were trained.

Environment Specifications

The T-maze environment is formally defined as a POMDP parameterized by the maze length L ∈ Z+, with a state space S composed of the position of the agent P and the position of the goal G, where S = P × G. The observation space O is deterministic and includes the initial observation of the goal at the first timestep, 0 otherwise. The LookupTreeMaze environment extends this by including multiple T-mazes N:= [Nlow, Nhigh] and varying maze lengths L:= [Llow, Lhigh].

RNN Equations

The equations for the different models used are provided, including the GRU update rule where ht = zt ⊙ ht−1 + (1 − zt) ⊙ nt and the BMRU update rule where ht = zt ⊙ S(nt) ⊙ α + (1 − zt) ⊙ ht−1. The minGRU update equation is given by ht = zt ⊙ ht−1 + (1 − zt) ⊙ nt and the BMRU update equation involves the term βt = Wβxt + bβ.

Experimental Results on Generalization

Experiments show that monostable models fail to generalize on either shorter or longer mazes, whereas multistable BMRU generalizes to all maze sizes in all conditions. Furthermore, the hybrid architecture combining minGRU and BMRU largely outperforms both architectures for a high number of T-mazes, some models maintaining an almost optimal policy to up to 10 000 T-mazes. This indicates that at least a subset of models achieve THG. The results show that monostable minGRU fails to generalize on either shorter or longer mazes, whereas multistable BMRU generalizes to all maze sizes in all conditions. The paper concludes that transient dynamics and multistability are both essential for solving long-horizon, potentially complex tasks.

Hyperparameter Summary

The training hyperparameters are summarized in Table 1, which details the settings for each benchmark, such as PPO epochs being 20 for Benchmark 1 and VF epochs being 10. The learning rate is set to Adam with a learning rate of 0.005 for the policy and 0.001 for the value function. The discount factor γ is set to 0.998 across both experiments.

Observation Distribution in LookupTreeMaze

In the LookupTreeMaze, if we assume that "1 is present n times in the lookup table with 0 < n < τ (and consequently −1 is present τ − n times in the table) then the probability for index i is p(i gt = 1) = 1/n and p(i gt = −1) = 1/τ−n". The observation distribution is deterministic, with specific probabilities assigned based on whether the agent has reached the far right of the maze.

Reward Function in LookupTreeMaze

The reward function for the LookupTreeMaze is defined such that rewards are given based on reaching terminal states, for example, 4 if st ∈ F/, st+1 ∈ F, pt+1 = (L − 1, g0), pt+1 = (L − 1, r), r ∈ G, r ≠ g0. The reward is zero otherwise for non-terminal transitions.

1. Bold header: Multistability-based Policy Generalization for Long-Horizon POMDPs. This improved system can remain optimal for all horizons T ≥ 0 by exploiting multistability, which is a necessary and sufficient condition for temporal horizon generalization (THG).

2. Bold header: Hybrid Parallelizable Architecture Design. The system should move away from purely monostable architectures by combining the benefits of transient dynamics alongside multistability, as this synergy is essential for achieving THG in complex environments like the LookupTreeMaze.

3. Bold header: Robustness to Non-Constant Dynamics via Basin Invariance. The agent can maintain optimal performance even when observations fluctuate during idle phases, as perturbations never push the hidden state across the boundary between two basins of attraction, ensuring that a read-out ρ remains informative.

4. Bold header: Efficient Exploration and Training via Short-Horizon Pre-training. The system should be trained on shorter horizon tasks to learn an optimal policy in only a couple of training steps, with almost no variance, which is more efficient than training directly on long horizons where exploration is much more difficult.

5. Bold header: Task-Specific Read-out Function Design. To achieve generalization, the system must employ a read-out function ρ that is both invariant within each basin and separating for σ, ensuring that different initial observations lead to distinct decision representations.

In partially observable reinforcement learning (RL), a decision often depends on a cue seen many steps earlier, and training at such horizons is costly. Ideally, a policy trained on short horizons would remain optimal at arbitrarily long ones, making the training horizon a matter of convenience. We formalize this property as delay-horizon independence (DHI): optimality at every horizon, at a memory, computation and precision cost that is bounded in the horizon. We study agents that carry the cue in a recurrent memory of fixed size, whose cost does not grow with the delay. This paper characterizes DHI through the dynamics of that memory, in theory and on benchmarks. Although quantified over all horizons, DHI reduces to a property of a single map: a recurrent policy is DHI if and only if it admits a continuous readout of the memory, invariant under the idle dynamics and separating the cues. Monostable memory then precludes DHI at any parameter values, whereas multistable memory enables it. Monostability covers all contracting affine recurrences, hence state space models and gated linear recurrent neural networks (RNNs). Under arbitrarily small perturbations of the idle dynamics, multistability is the only mechanism supporting DHI at bounded memory. On the T-maze, the effective horizon of monostable models resists hyperparameter search: selected horizons fall short by up to 745 times when retrained. Multistable models trained on short horizons stay optimal up to 10 6, on the T-maze and on MiniGrid. Current parallelizable RNNs thus trade horizon independence for scalability: parallelizable multistable architectures are the direction this points to.