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

TL;DR

This paper links recurrent policy hidden states to Pontryagin co-states from optimal control, introducing a co-state loss to improve interpretability and robustness in partially observable reinforcement learning tasks.

Contribution

It establishes a formal connection between recurrent hidden states and Pontryagin co-states, enabling structured and interpretable internal dynamics in reinforcement learning policies.

Findings

Matching or improving performance on DMControl tasks

Robustness against zero-shot sensor masking

Provides a principled approach to policy design

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.