🤖 machine learning

Integrating Causal DAGs in Deep RL: Activating Minimal Markovian States with Multi-Order Exposure

This paper addresses the challenge of constructing provably Markovian states from longitudinal causal graphs in deep reinforcement learning by introducing MOSE (Multi-Order State Exposure), a method that feeds multi-order historical state constructions into the Q-function to demonstrate that controlled redundancy, rather than minimal sufficiency alone, is essential for unlocking the performance benefits of causal state information.

Original authors: Jiamin Xu, Jacqueline Maasch, Kyra Gan

Published 2026-05-11

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jiamin Xu, Jacqueline Maasch, Kyra Gan

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 teaching a robot to play a video game or navigate a maze. To make good decisions, the robot needs to know its current "state." In a perfect world, the robot would only need to look at the exact moment right now to know what to do next. This is called the Markov property.

However, in the real world, things are messy. The robot's sensors (like cameras) give it raw data, but that data often doesn't tell the whole story. For example, if a robot sees a ball, it doesn't know if the ball is rolling toward it or away from it unless it remembers where the ball was a second ago.

Here is the problem: If the robot forgets the past, it makes bad guesses. If it remembers everything (every single pixel from the last 100 seconds), it gets overwhelmed and learns too slowly.

This paper, "Integrating Causal DAGs in Deep RL," proposes a clever solution to this "Goldilocks" problem: finding the just right amount of history to remember.

The Core Idea: The "Minimal" vs. The "Redundant"

The authors tackle this in two steps, using a mix of logic (causality) and a bit of "controlled chaos" (redundancy).

1. The "Minimal" State (The Perfectly Packed Suitcase)

First, the authors use a Causal Graph (a map showing which variables cause which others) to figure out the absolute minimum amount of information needed to make a perfect decision.

The Analogy: Imagine you are packing for a trip. You want to bring the absolute minimum amount of clothes to survive. You calculate exactly what you need: one shirt, one pair of pants, and socks. You leave everything else behind.
The Result: Theoretically, this "minimal suitcase" is perfect. It has no fluff.
The Catch: When the authors tried feeding this "minimal suitcase" to a modern AI (a deep neural network), it failed. The AI got confused. It turns out that AI networks are like students who learn better when they have a little bit of extra context, not just the bare facts. The "minimal" state was too sparse, making it hard for the AI to learn the patterns.

2. The Solution: MOSE (Multi-Order State Exposure)

To fix this, the authors invented MOSE (Multi-Order State Exposure).

The Analogy: Instead of giving the student just the "minimal suitcase," MOSE gives them a series of suitcases of different sizes.
- Suitcase A: Just the current moment.
- Suitcase B: The current moment + the last 1 second.
- Suitcase C: The current moment + the last 2 seconds.
- ...and so on.
How it works: The AI looks at all these different suitcases at the same time. It gets to see the "minimal" version, but it also gets to see versions with extra history.
Why it helps: This acts like a training wheel. The AI can start with the simple, short history and gradually learn to use the longer history when it helps. It's like giving a student a hint, then a bigger hint, then the full answer, all at once, so they can figure out which clues are actually important.

3. The "Best of Both Worlds" (Causal-MOSE)

The authors also tried a hybrid approach called Causal-MOSE. This combines the "perfectly packed minimal suitcase" (derived from the causal map) with the "multiple suitcases" approach.

The Result: This was often the winner. It gave the AI the "perfect core" of information (guaranteed by the math) but allowed it to add extra "redundant" information if it helped the learning process.

What the Experiments Showed

The team tested this on:

Synthetic Games: Made-up worlds where they knew the exact rules (the causal graph).
Real Games: Specifically, an Atari game called GOPHER.

The Findings:

Standard Method (Frame Stacking): This is the current industry standard, where you just stack the last 4 video frames together. It works okay, but it's like carrying a suitcase full of junk you don't need.
Minimal State: Using only the mathematically perfect, minimal history actually made the AI perform worse than the standard method.
MOSE: The new method consistently beat both the standard method and the minimal method.
The Big Lesson: The paper concludes that "Minimal sufficiency is not enough." Just because a state has the minimum amount of information to be theoretically correct doesn't mean it's the best for a neural network to learn from. You need controlled redundancy (a little bit of extra, messy history) to help the AI learn faster and better.

Summary in One Sentence

The paper teaches us that to train a smart AI, you shouldn't just give it the bare minimum facts (which confuses it); instead, you should give it a mix of short and long histories so it can figure out exactly what it needs to remember to win.

Technical Summary: Integrating Causal DAGs in Deep RL with MOSE

Problem Statement

Online reinforcement learning (RL) algorithms fundamentally rely on the Markov property, which assumes that the future is conditionally independent of the past given the present state. However, in real-world applications involving raw observations (e.g., sensor streams, video frames), the observed variables often do not inherently form a Markovian state because system dynamics may depend on an unknown amount of historical context.

While existing causal RL literature utilizes causal Directed Acyclic Graphs (DAGs) to accelerate learning or design robust rewards, these methods typically assume that a valid, pre-specified Markovian state representation is already available. A critical gap remains: given a longitudinal causal graph over observed variables, there is no established procedure to construct a state representation that provably satisfies the Markov property. Furthermore, in the context of Deep RL, simply constructing a theoretically "minimal" Markovian state often fails to improve empirical performance, as neural networks may struggle to generalize from such sparse representations due to overfitting and a lack of redundant features that aid function approximation.

Methodology

The paper proposes a two-pronged approach: a theoretical procedure for constructing minimal states from known causal graphs, and a heuristic method (MOSE) for deep RL when the graph is unknown or to enhance learning stability.

1. Minimal Markovian State Construction from Causal DAGs

The authors establish a formal connection between time-series causal DAGs and Markov Decision Processes (MDPs).

Theoretical Framework: Using Structural Causal Models (SCMs) extended to time series, the paper defines a "Full Time Graph" where nodes represent variables at specific time steps and edges represent causal dependencies with time lags.
Graphical Criterion: Theorem 4.1 provides four necessary and sufficient conditions for a sequence of variable sets $\{S_t\}$ ${S_{t}}$ to constitute a valid state space:
1. Reward Parent Inclusion: All parents of the reward variable $R_t$ must be included in $S_t$ .
2. Extended Next-State Parent Inclusion: All parents (at times $\le t$ ) of the "extended next-state set" (the closure of $S_{t+1}$ under same-time parents) must be included in $S_t$ .
3. No Ancestral Re-entry: Once a variable is removed from the state construction, its ancestors cannot re-enter the state at a later time.
4. Persistence: If a variable appears in the state at two different times, it must appear in all intermediate states.
Algorithm F.1: A backward construction algorithm is proposed to generate a provably minimal state sequence satisfying the above criteria. Starting from the end of the horizon, it iteratively selects the necessary parents of the reward and the next state, ensuring no redundant variables are included while preserving the Markov property and optimal value.

2. Multi-Order State Exposure (MOSE)

Recognizing that minimal representations can lead to poor generalization in deep RL (due to sharp minima and lack of statistical slack), the authors propose MOSE.

Mechanism: Instead of feeding a single minimal state or a fixed window of history to the Q-function, MOSE feeds multiple orders of historical state constructions simultaneously. Specifically, for a window length $W$ , the algorithm iteratively updates the value function using inputs ranging from the current observation ( $l=0$ ) up to the full window ( $l=W$ ).
Rationale: The lower-order, non-Markovian inputs act as structured noise that reshapes the loss landscape. This provides a gradual, informative gradient from suboptimal to optimal representations, acting as a regularizer similar to data augmentation or multi-view learning.
Causal-MOSE: An extension where the minimal Markovian state (derived from the causal DAG) is included alongside the multi-order inputs. This combines the theoretical guarantee of the minimal core with the empirical benefits of redundancy.

Key Contributions

Minimal Markovian State Construction: The paper provides a graphical criterion and a constructive backward algorithm (Algorithm F.1) to derive a provably minimal state space from a known longitudinal causal DAG. This bridges the gap between causal discovery and MDP state definition.
Multi-Order State Exposure (MOSE): A heuristic method for state construction that does not require prior knowledge of the causal DAG. It demonstrates that exposing deep RL agents to multiple temporal windows of history improves performance over both naive frame stacking and minimal causal representations.
Causal-MOSE: A hybrid approach showing that combining the minimal Markovian state with multi-order exposure yields performance equal to or better than MOSE alone, validating the utility of causal structure when available.

Experimental Results

The authors evaluated their methods on three synthetic Data Generating Processes (DGPs) and the Atari game GOPHER.

Synthetic DGPs: Using Linear Additive Noise Models (ANM) and Post-Nonlinear (PNL) models with varying MDP orders ( $W=2, 5, 10$ $W = 2, 5, 10$ ), MOSE consistently outperformed:
- Window Policy: Standard frame stacking (concatenating fixed $W$ frames).
- Reward Parents: Using only parents of the reward as the state.
- DAG-State: The minimal state constructed by Algorithm F.1.
- Key Finding: The minimal state (DAG-State) alone did not improve performance over window policies and sometimes degraded it, confirming that theoretical minimality does not equate to empirical optimality in deep RL.
Causal-MOSE Performance: In Linear ANM environments, Causal-MOSE (minimal state + MOSE) consistently outperformed MOSE alone, suggesting that the minimal state provides a reliable core for the network to build upon.
Atari (GOPHER): MOSE outperformed the standard Rainbow algorithm (which uses 4-frame stacking) over 1 million training steps. This indicates that in high-dimensional pixel environments, not all stacked frames contribute equally to control, and adaptive state selection improves sample efficiency.
Ablations: Experiments confirmed that the gains were not due to increased gradient steps per batch (optimization artifacts) but rather the richer state representations provided by MOSE.

Significance and Claims

The paper establishes a core principle for causal deep RL: minimal sufficiency is not enough, and controlled redundancy is necessary to unlock the benefit of causal state information.

The authors argue that while causal graphs provide a principled way to identify the minimal variables required for the Markov property, deep neural networks benefit from the "statistical slack" provided by redundant, non-Markovian inputs. These redundant inputs help the network avoid overfitting to spurious correlations and facilitate better generalization.

The work contributes a theoretical foundation for constructing valid MDP states from causal graphs and a practical heuristic (MOSE) that makes these causal insights usable in modern deep RL. The authors note that their approach assumes access to a ground-truth longitudinal causal graph; learning such graphs from high-dimensional, non-stationary data remains an open challenge for future work.