# Learning Causal States Under Partial Observability and Perturbation

arXiv: https://arxiv.org/abs/2512.00357 · HTML: https://mlsanity.com/paper/2512.00357/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Diffusion and flow policies for robotics, one of 14 directions in Diffusion models.
- Builds on 33 of its 46 references that are in the map.
- First posted on arXiv 29 Nov 2025.

## Placed in

- Diffusion models / [Diffusion and flow policies for robotics](https://mlsanity.com/diffusion/diffusion-and-flow-policies-for-robotics/)
- Deep reinforcement learning methods / [Deep RL algorithms and representations](https://mlsanity.com/deeprl/deep-rl-algorithms-and-representations/)

## Builds on (most cited first, 30 of 33)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Improved Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2102.09672/) (2021)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [Planning with Diffusion for Flexible Behavior Synthesis](https://mlsanity.com/paper/2205.09991/) (2022)
- [Recurrent World Models Facilitate Policy Evolution](https://mlsanity.com/paper/1809.01999/) (2018)
- [Is Conditional Generative Modeling all you need for Decision-Making?](https://mlsanity.com/paper/2211.15657/) (2022)
- [Time-Contrastive Networks: Self-Supervised Learning from Video](https://mlsanity.com/paper/1704.06888/) (2017)
- [Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions](https://mlsanity.com/paper/2209.11215/) (2022)
- [Learning Invariant Representations for Reinforcement Learning without Reconstruction](https://mlsanity.com/paper/2006.10742/) (2020)
- [A Theoretical Analysis of Deep Q-Learning](https://mlsanity.com/paper/1901.00137/) (2019)
- [Scalable methods for computing state similarity in deterministic Markov Decision Processes](https://mlsanity.com/paper/1911.09291/) (2019)
- [State Representation Learning for Control: An Overview](https://mlsanity.com/paper/1802.04181/) (2018)
- [Deep Variational Reinforcement Learning for POMDPs](https://mlsanity.com/paper/1806.02426/) (2018)
- [Recurrent Model-Free RL Can Be a Strong Baseline for Many POMDPs](https://mlsanity.com/paper/2110.05038/) (2021)
- [Policy Representation via Diffusion Probability Model for Reinforcement Learning](https://mlsanity.com/paper/2305.13122/) (2023)
- [PlanT: Explainable Planning Transformers via Object-Level Representations](https://mlsanity.com/paper/2210.14222/) (2022)
- [Diffusion Models Without Attention](https://mlsanity.com/paper/2311.18257/) (2023)
- [Nonparametric Regression on Low-Dimensional Manifolds using Deep ReLU Networks : Function Approximation and Statistical Recovery](https://mlsanity.com/paper/1908.01842/) (2019)
- [Unveil Conditional Diffusion Models with Classifier-free Guidance: A Sharp Statistical Theory](https://mlsanity.com/paper/2403.11968/) (2024)
- [Accelerating Diffusion Models with Parallel Sampling: Inference at Sub-Linear Time Complexity](https://mlsanity.com/paper/2405.15986/) (2024)
- [Bisimulation Makes Analogies in Goal-Conditioned Reinforcement Learning](https://mlsanity.com/paper/2204.13060/) (2022)
- [Action-Sufficient State Representation Learning for Control with Structural Constraints](https://mlsanity.com/paper/2110.05721/) (2021)
- [Learning Causal State Representations of Partially Observable Environments](https://mlsanity.com/paper/1906.10437/) (2019)
- [Causal Diffusion Autoencoders: Toward Counterfactual Generation via Diffusion Probabilistic Models](https://mlsanity.com/paper/2404.17735/) (2024)
- [Lower Bounds for Learning in Revealing POMDPs](https://mlsanity.com/paper/2302.01333/) (2023)
- [Invariant Causal Imitation Learning for Generalizable Policies](https://mlsanity.com/paper/2311.01489/) (2023)
- [On the rate of convergence of a deep recurrent neural network estimate in a regression problem with dependent data](https://mlsanity.com/paper/2011.00328/) (2020)
- [Ensemble Bayesian Decision Making with Redundant Deep Perceptual Control Policies](https://mlsanity.com/paper/1811.12555/) (2018)
- [CAR-DESPOT: Causally-Informed Online POMDP Planning for Robots in Confounded Environments](https://mlsanity.com/paper/2304.06848/) (2023)
- [ODE-based Recurrent Model-free Reinforcement Learning for POMDPs](https://mlsanity.com/paper/2309.14078/) (2023)

## Abstract (arXiv)

A critical challenge for reinforcement learning (RL) is making decisions based on incomplete and noisy observations, especially in perturbed and partially observable Markov decision processes (P2OMDPs). Existing methods fail to mitigate perturbations while addressing partial observability. We propose textitCausal State Representation under Asynchronous Diffusion Model (CaDiff), a framework that enhances any RL algorithm by uncovering the underlying causal structure of P2OMDPs. This is achieved by incorporating a novel asynchronous diffusion model (ADM) and a new bisimulation metric. ADM enables forward and reverse processes with different numbers of steps, thus interpreting the perturbation of P2OMDP as part of the noise suppressed through diffusion. The bisimulation metric quantifies the similarity between partially observable environments and their causal counterparts. Moreover, we establish the theoretical guarantee of CaDiff by deriving an upper bound for the value function approximation errors between perturbed observations and denoised causal states, reflecting a principled trade-off between approximation errors of reward and transition-model. Experiments on Roboschool tasks show that CaDiff enhances returns by at least 14.18% compared to baselines. CaDiff is the first framework that approximates causal states using diffusion models with both theoretical rigor and practicality.