From Causal Graphs to Markovian States: Multi-Order State Exposure for Deep RL

The gist: MOSE consistently outperforms both minimal state construction and single-window policies on common benchmarks and synthetic datasets, establishing that controlled redundancy is necessary to unlock the benefit of causal state information

Minimal Markovian State Construction from Time Series Causal DAGs

The work establishes a direct mapping from a known longitudinal causal graph to a minimal Markovian state representation that satisfies the Markov property and is minimal in the sense that removing any variable would, for some consistent environment, either break the Markov property or reduce the optimal value The authors provide a procedure that returns a provably minimal state space satisfying validity This construction is formalized by Theorem 4.1, which presents a graphical criterion for constructing valid states given a time series causal DAG The procedure, Algorithm F.1, constructs the state sets from the full time graph backwards starting from an empty set at time T + 1.

Multi-Order State Exposure (MOSE)

To address the issue where minimal representation alone fails to improve performance in deep RL, MOSE is proposed as a heuristic method for state space construction that does not require knowledge of the causal DAG Instead of feeding only the minimal (Wth order) state to the Q-function, MOSE feeds all orders 1,..., W simultaneously. The lower-order, non-Markovian inputs act as a form of structured noise that reshapes the loss landscape, creating a gradual, informative gradient from suboptimal to optimal representations. Empirically, MOSE substantially outperforms both the minimal state and standard frame stacking in multiple environments.

Causal-MOSE and Controlled Redundancy

Causal-MOSE is introduced as an extension of MOSE that includes the minimal representation alongside the multi-order input. The authors show that Causal-MOSE performed as well or better than MOSE in experiments, suggesting that controlled redundancy can unlock the benefit of causal state information. The results indicate that combining the minimal Markovian state with MOSE performs as well or better on the experimental settings explored in this work.

Experimental Results and Insights

Experiments tested MOSE and Causal-MOSE against two common ways of constructing state with time dependence: Reward Parent, Window Policy, and DAG-State. In synthetic data generating processes (DGPs), MOSE outperforms all competing benchmarks throughout the full training horizon for both Linear ANM and PNL models. A particularly notable observation is that the benchmark based only on the minimal Markovian state does not improve performance, and in fact remains similar to window policies across all three orders W. This suggests that although some variables are redundant for satisfying the Markov property, they may still contain auxiliary statistical information that helps the network estimate the Q-function more accurately in finite samples.

Conclusion

The work established a formal connection between time series causal DAGs and MDPs, derived a graphical criterion for constructing a minimal, Markovian state representation, and provided a constructive backward algorithm that extracts such a representation directly from the full time graph. MOSE is proposed as a practical heuristic that varies the observation window length to encourage the agent to discover a useful state representation on its own. The findings support a nuanced conclusion: the best performance is achieved neither by using the full raw state nor by restricting the learner to the minimal Markovian state alone, but by combining a principled Markovian core with a mechanism that can incorporate extra informative variables when they are beneficial for learning.

How it works

1. Minimal Markovian State Construction: This involves constructing a set of valid states based on the graphical criterion, which requires conditions such as reward parent inclusion and the exclusion of ancestral re-entry.

2. Multi-Order State Exposure (MOSE): This heuristic feeds all orders 1,..., W simultaneously into the Q-function. The lower-order inputs act as a form of structured noise that reshapes the loss landscape.

3. Causal-MOSE: This combines the minimal Markovian state with MOSE, exposing the network to this set of information, allowing it to start from the minimal representation and add variables that are useful.

Key Findings

**/The authors prove a graphical criterion for constructing a valid state space given a time series causal DAG. The proof of Theorem 4.1 shows that any state space satisfying the criterion satisfies the Markov property and optimality preservation. The minimality of the state space returned by Algorithm F.1 is proven, showing that removal of any variable would violate validity or reduce optimal value. The analysis in Lemma F.3 demonstrates that omitting even a single necessary variable can strictly reduce the optimal value. The empirical results show that MOSE consistently outperforms all competing benchmarks throughout the full training horizon for both Linear ANM and PNL models. The comparison between MOSE and Causal-MOSE in Figure 2 shows that Causal-MOSE consistently outperforms MOSE over the entire 1M-episode training period in Linear ANM. The improvement in Figure 3 indicates that this more targeted state construction can lead to better sample efficiency and stronger performance in practice. The runtime of MOSE and Causal-MOSE is comparable with other policies even when it has more gradient steps per batch. The paper concludes that the best performance is achieved neither by using the full raw state nor by restricting the learner to the minimal Markovian state alone, but by combining a principled Markovian core with a mechanism that can incorporate extra informative variables when they are beneficial for learning. The limitations include assuming a known full time graph and the fact that non-Markovian states sometimes perform comparably. The paper also suggests future work on how to construct valid state given DAG and other further information when causal Markovian condition fails. The authors note that the tanh transformation compresses large-magnitude values toward 1 and amplifies smaller values relative to them, thereby reducing the distinction among different state components in some settings. The benefit of Causal-MOSE over MOSE becomes smaller as W increases and is nearly negligible when W = 10 in PNL models. The authors conclude that the performance gap between different state constructions becomes much smaller when the state values are bounded. The paper demonstrates that by selectively constructing the state representation, MOSE reduces the burden of learning from irrelevant or weakly informative variables for Q-function, leading to better performance across training. The improvement in Figure 3 indicates that this more targeted state construction can lead to better sample efficiency and stronger performance in practice. The paper shows that by directly increasing gradient steps on the stacked frame, the performance drops when using the whole stacked frame as state whereas when using MOSE to selectively choose state, Rainbow may be able to avoid this problem. The runtime of MOSE and Causal-MOSE is comparable with other policies even when it has more gradient steps per batch. The authors report that the performance gap between different state constructions becomes much smaller when the state values are bounded. The paper demonstrates that by selectively constructing the state representation, MOSE reduces the burden of learning from irrelevant or weakly informative variables for Q-function, leading to better performance across training.

1. Wth Order MDP State Construction: Implement a procedure that constructs a provably minimal state representation from raw observations based on a known longitudinal causal DAG, as described in Theorem 4.2, to ensure the constructed state space is both Markovian and minimal in the sense that removing any variable would, for some consistent environment, either break the Markov property or reduce the optimal value.

2. MOSE Heuristic Integration: Integrate Multi-Order State Exposure (MOSE) into deep RL algorithms by feeding all orders 1,..., W simultaneously, which is shown to consistently outperform both the minimal state construction and single-window policies on common benchmarks and synthetic datasets.

3. Causal-MOSE for Unknown Graphs: Develop Causal-MOSE, which includes the minimal representation alongside the multi-order input, allowing deep RL agents to leverage causal information even when the underlying DAG is unknown, as this approach performed as well or better than MOSE in experiments.

Online reinforcement learning (RL) relies on the Markov property for guaranteed performance, but real-world applications often lack well-defined states given raw observed variables. While causal RL has attracted growing interest, existing work typically assumes Markovian states are provided and uses causality only to accelerate learning. This leaves a fundamental gap: given a longitudinal causal graph over observed variables, how does one construct MDP states that provably satisfy the Markov property? We provide a procedure that constructs a provably minimal such state. In deep RL, the minimal representation alone empirically fails to improve performance, indicating that neural networks cannot automatically exploit Markovian minimality. We therefore propose MOSE (Multi-Order State Exposure), which exposes a single shared Q-function to history-based state constructions of orders 1 through W simultaneously. An attention-based selection layer scores the dimensions of this combined input by their estimated relevance to value-function estimation and selects the high-scoring ones. Causal MOSE further constrains the selection layer to retain every variable of the minimal state, guaranteeing sufficiency while leaving the model free to admit useful redundancy. MOSE consistently outperforms both the minimal state construction and single-window policies on common benchmarks and synthetic datasets. Including the minimal representation alongside MOSE can further improve performance. Our results establish a core principle for causal deep RL: minimal sufficiency is not enough. Controlled redundancy gives the model flexibility to learn its state representation end-to-end. Causal state information provides structural guidance that can make this learning more efficient.