English

Neural Co-state Policies: Structuring Hidden States in Recurrent Reinforcement Learning

Machine Learning 2026-05-12 v2

Abstract

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.

Keywords

Cite

@article{arxiv.2605.05373,
  title  = {Neural Co-state Policies: Structuring Hidden States in Recurrent Reinforcement Learning},
  author = {David Leeftink and Max Hinne and Marcel van Gerven},
  journal= {arXiv preprint arXiv:2605.05373},
  year   = {2026}
}

Comments

17 pages, 5 figures

R2 v1 2026-07-01T12:53:34.367Z