Neural Co-state Policies: Structuring Hidden States in Recurrent Reinforcement Learning
For researchers in reinforcement learning and control, this provides a principled framework to design interpretable and robust recurrent policies under partial observability.
This paper establishes a formal link between hidden states in recurrent reinforcement learning policies and co-states from the Pontryagin minimum principle, introducing a co-state loss to structure internal dynamics. The approach matches or improves performance on partially observable DMControl tasks and shows robustness to zero-shot out-of-distribution sensor masking.
A key capability of intelligent agents is operating under partial observability: reasoning and acting effectively despite missing or incomplete state observations. While recurrent (memory-based) policies learned via reinforcement learning address this by encoding history into latent state representations, their internal dynamics remain uninterpretable black boxes. This paper establishes a formal link between these hidden states and the Pontryagin minimum principle (PMP) from optimal control. We demonstrate that for standard recurrent architectures, latent representations map directly to PMP co-states, which allows the readout layer to be interpreted as performing Hamiltonian minimization. Because standard reward maximization does not naturally discover this alignment, we introduce a PMP-derived co-state loss to explicitly structure the internal dynamics. Empirically, this approach matches or improves performance on partially observable DMControl tasks, and is robust against zero-shot out-of-distribution sensor masking. By framing recurrent networks as dynamic processes governed by the minimum principle, we provide a principled approach to designing robust continuous control policies.