🤖 machine learning

Policy Gradient Methods for Non-Markovian Reinforcement Learning

This paper introduces a reward-centric framework for non-Markovian reinforcement learning that jointly optimizes agent state dynamics and control policies, establishing a novel policy gradient theorem and the ASMPG algorithm with theoretical convergence guarantees and superior empirical performance over predictive baselines.

Original authors: Avik Kar, Siddharth Chandak, Rahul Singh, Soumitra Sinhahajari, Eric Moulines, Shalabh Bhatnagar, Nicholas Bambos

Published 2026-05-12

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Avik Kar, Siddharth Chandak, Rahul Singh, Soumitra Sinhahajari, Eric Moulines, Shalabh Bhatnagar, Nicholas Bambos

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot to navigate a maze, but there's a catch: the robot is blindfolded. It can't see the walls or the exit. All it knows are the sounds it hears (like a creaking floorboard) and the feelings it gets (like bumping into a wall).

In the world of Reinforcement Learning (RL), this is called a Non-Markovian problem. The robot's current situation isn't just about now; it depends entirely on everything that happened before. If the robot bumps into a wall, it doesn't know which wall it is unless it remembers where it started and what turns it took.

Most standard AI methods struggle here because they try to guess the future based only on the "now," or they try to build a perfect map of the past, which gets too heavy and complicated to carry around.

This paper introduces a new way to teach these blindfolded robots, called ASMPG (Agent State-Markov Policy Gradient). Here is how it works, using simple analogies:

1. The Problem: The "Amnesiac" vs. The "Over-Thinker"

The Amnesiac (Standard MDP): Imagine a robot that forgets everything the moment it takes a step. It only knows, "I am here, I am hungry." If the environment is complex (like a conversation or a maze), this robot fails because it doesn't know the context.
The Over-Thinker (History-Based): Imagine a robot that tries to remember every single word of a conversation or every single step of a maze. While this contains all the information, the list of memories grows infinitely long. It becomes impossible to process.

2. The Solution: The "Smart Diary" (Agent State)

The authors propose a middle ground. Instead of forgetting everything or remembering everything, the robot keeps a Smart Diary (called an "Agent State").

How it works: Every time the robot takes an action or sees something new, it updates its diary. It doesn't write down the whole history; it just writes a summary.
- Example: In a chatbot, instead of remembering the entire 100-page conversation, the diary just says: "User is asking about their order status, and they seem impatient."
The Twist: In previous methods, scientists would try to write this diary summary by asking, "Can you predict what the user will say next?" (a predictive objective).
The Innovation: This paper says, "Stop guessing the future. Just write the summary that helps you get the reward (the happy customer)." They teach the robot to write the diary and decide what to do, all at the same time, specifically to maximize the score.

3. The Method: The "Twin-Engine" Approach

The paper introduces a new algorithm called ASMPG. Think of it as a twin-engine plane where both engines are optimized together:

Engine A (The Scribe): Updates the diary (the Agent State) based on new inputs.
Engine B (The Pilot): Reads the diary and decides what action to take.

In older methods, the Scribe was fixed or trained separately to be a "good predictor." In ASMPG, the Scribe and the Pilot are trained jointly. If the Pilot needs a specific detail in the diary to make a good decision, the Scribe learns to include that detail. If the Pilot doesn't need a detail, the Scribe learns to ignore it. They work as a team to win the game.

4. The Proof: Why It Works

The authors did the math to prove that this "joint training" approach is valid.

They derived a new formula (a "Policy Gradient Theorem") that shows exactly how to adjust the Scribe and the Pilot to get better scores.
They proved that if you keep making small adjustments based on this formula, the robot will eventually learn a very good strategy (mathematically guaranteed to converge).

5. The Results: Winning the Game

They tested this new "Smart Diary" approach on five different tricky tasks where the robot couldn't see the whole picture:

CheeseMaze: A robot finding cheese in a maze where different spots look identical.
Hallway Navigation: Walking down a hallway where you can only see the walls right next to you.
Healthcare: Deciding on medical treatments where the patient's reaction depends on their hidden history of past treatments (toxicity and resistance).
Machine Repair: Fixing a machine where you can only see if it's "sick" or "healthy," but the real cause is hidden wear and tear from the past.
CartPole: Balancing a pole on a cart when you can only see the speed, not the position.

The Outcome: In all five cases, the ASMPG robot (the one with the jointly trained Smart Diary) learned faster and got higher scores than robots that tried to learn by predicting the future or using fixed memory systems.

Summary

This paper is about teaching AI agents how to handle situations where "the present" isn't enough to make a decision. Instead of trying to remember everything or guessing the future, the authors teach the AI to maintain a dynamic, evolving summary of its past. Crucially, they teach the AI to build this summary specifically to win the game, rather than just to be a good historian. The result is a smarter, more efficient learner for complex, real-world problems.

Technical Summary: Policy Gradient Methods for Non-Markovian Reinforcement Learning

Problem Statement

Reinforcement Learning (RL) traditionally relies on the Markov Decision Process (MDP) framework, assuming that dynamics and rewards depend solely on the current state and action. However, many real-world environments—such as partially observable systems, dialogue agents, recommender systems, and financial markets—are inherently Non-Markovian Decision Processes (NMDPs). In these settings, observations and rewards depend on the entire interaction history.

Directly conditioning policies on the full history is computationally intractable as the history space grows exponentially. While existing approaches attempt to compress history, they often face limitations:

Fixed Dynamics: Some methods assume the agent's internal state dynamics are fixed and only optimize the control policy (e.g., Dong et al., 2022).
Predictive Objectives: Other methods learn state representations by optimizing auxiliary predictive objectives (e.g., predicting future observations), which may require high-dimensional representations and does not directly align with the ultimate goal of maximizing cumulative reward (e.g., Subramanian et al., 2022).

This paper addresses the challenge of jointly learning a compact state representation and a control policy in NMDPs without relying on fixed dynamics or intermediate predictive objectives.

Methodology

Agent State-Markov (ASM) Policies

The authors introduce a class of Agent State-Markov (ASM) policies. In this framework, the agent maintains an internal state $S_t$ that is recursively updated based on the previous state, the previous action, and the current observation:
$S_t \sim \nu_t(\cdot \mid S_{t-1}, A_{t-1}, O_t)$
The control policy $\phi$ then maps the current agent state to an action:
$A_t \sim \phi_t(\cdot \mid S_t)$
The pair $(\nu, \phi)$ constitutes the ASM policy. Crucially, the authors propose optimizing the parameters of both the state dynamics $\nu$ and the control policy $\phi$ jointly to maximize the expected cumulative reward, rather than fixing $\nu$ or learning it via prediction.

Theoretical Contributions

The paper establishes a Policy Gradient Theorem for ASM policies, extending classical results to both episodic and infinite-horizon discounted NMDPs.

Gradient Expression: For a parameterized ASM policy $\pi_\theta$ , the gradient of the expected return $J(\theta)$ is derived as:
$\nabla_\theta J(\theta) = \mathbb{E}_{\pi_\theta} \left[ \sum_{t} R_{t:\infty} \nabla_\theta \log \pi_\theta(S_t, A_t \mid S_{t-1}, A_{t-1}, O_t) \right]$
where $R_{t:\infty}$ represents the future return (reward-to-go).
Factorization: Due to the structure of ASM policies, the log-probability term factorizes into the state dynamics and the control policy:
$\nabla_\theta \log \pi_\theta = \nabla_\theta \log \nu_\theta + \nabla_\theta \log \phi_\theta$
This allows for the separate but simultaneous update of the state representation and the policy.
Convergence Guarantees: Under standard smoothness and boundedness assumptions, the authors prove:
- Finite-time convergence: The time-averaged squared gradient norm converges at a rate of $\tilde{O}(H^{3.5}/\sqrt{K})$ for episodic tasks (where $H$ is the horizon and $K$ is the number of iterations).
- Almost sure convergence: With decreasing step sizes, the gradient norm converges to zero almost surely.
- Similar guarantees are established for the infinite-horizon discounted setting, where the effective horizon is determined by the discount factor $\gamma$ .

Algorithm: ASMPG

Based on the gradient theorem, the authors propose the Agent State-Markov Policy Gradient (ASMPG) algorithm. It is an online, REINFORCE-style algorithm that:

Samples trajectories using the current ASM policy.
Computes the gradient estimate using the recursive structure of the agent state.
Updates the parameters of both the state dynamics network and the policy network using stochastic gradient ascent.

Experimental Results

The authors evaluated ASMPG on five non-Markovian environments:

CheeseMaze: A partially observable navigation task with aliased observations.
HallwayNavigation: A deterministic maze task where local wall configurations are ambiguous.
HealthcareTreatment: A treatment planning task with hidden toxicity and resistance variables.
MachineRepair: A maintenance task with hidden wear levels affecting machine degradation.
VelocityOnlyCartPole: A control task where only velocity (not position) is observed.

Comparison: ASMPG was compared against AIS-KL and AIS-MMD (Subramanian et al., 2022), which learn state representations via predictive objectives (Approximate Information States).

Findings:

Across all five environments, ASMPG achieved higher average rewards per step or per episode compared to the predictive baselines.
The performance gap was particularly notable in tasks requiring long-term memory, such as CheeseMaze, MachineRepair, and VelocityOnlyCartPole.
The results suggest that directly optimizing the state dynamics for the reward objective is more effective than optimizing for predictive sufficiency.

Significance and Claims

The paper claims that its primary contribution is a reward-centric formulation for non-Markovian RL. By treating the agent state dynamics as a learnable component optimized jointly with the control policy, the method bypasses the need for auxiliary predictive objectives or fixed state representations.

Key claims include:

Theoretical Validity: The derivation of a policy gradient theorem for general NMDPs provides a rigorous foundation for joint optimization, extending classical MDP results.
Practical Efficacy: Empirical results demonstrate that reward-centric joint optimization outperforms methods that rely on predictive state representations, even when the representation sizes are comparable.
Optimality Conditions: The paper notes that if an "ideal" agent state dynamics exists (one that is a sufficient statistic for rewards and future observations), restricting the search to ASM policies incurs no loss of optimality.

The authors acknowledge that the proposed method relies on REINFORCE-type estimators, which may suffer from high variance. They identify the development of actor-critic variants for ASM policies as a natural direction for future work to reduce variance while preserving the reward-centric objective.