# Fine-Grained Causal Dynamics Learning with Quantization for Improving Robustness in Reinforcement Learning

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

## Key facts

- Cited by 11 papers in the map (data through 14 Oct 2026).
- Placed in Deep RL algorithms and representations, one of 15 directions in Deep reinforcement learning methods.
- Builds on 56 of its 81 references that are in the map.
- First posted on arXiv 5 Jun 2024.

## Placed in

- 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 56)

- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Empirical Evaluation of Gated Recurrent Neural Networks on Sequence Modeling](https://mlsanity.com/paper/1412.3555/) (2014)
- [Generating Diverse High-Fidelity Images with VQ-VAE-2](https://mlsanity.com/paper/1906.00446/) (2019)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [robosuite: A Modular Simulation Framework and Benchmark for Robot Learning](https://mlsanity.com/paper/2009.12293/) (2020)
- [Model-Based Reinforcement Learning for Atari](https://mlsanity.com/paper/1903.00374/) (2019)
- [Causal Confusion in Imitation Learning](https://mlsanity.com/paper/1905.11979/) (2019)
- [Recurrent Independent Mechanisms](https://mlsanity.com/paper/1909.10893/) (2019)
- [Differentiable Causal Discovery from Interventional Data](https://mlsanity.com/paper/2007.01754/) (2020)
- [Contrastive Learning of Structured World Models](https://mlsanity.com/paper/1911.12247/) (2019)
- [Explainable Reinforcement Learning Through a Causal Lens](https://mlsanity.com/paper/1905.10958/) (2019)
- [Counterfactual Off-Policy Evaluation with Gumbel-Max Structural Causal Models](https://mlsanity.com/paper/1905.05824/) (2019)
- [Invariant Causal Prediction for Block MDPs](https://mlsanity.com/paper/2003.06016/) (2020)
- [Counterfactual Data Augmentation using Locally Factored Dynamics](https://mlsanity.com/paper/2007.02863/) (2020)
- [Woulda, Coulda, Shoulda: Counterfactually-Guided Policy Search](https://mlsanity.com/paper/1811.06272/) (2018)
- [Causal Discovery in Physical Systems from Videos](https://mlsanity.com/paper/2007.00631/) (2020)
- [Characterizing and Learning Equivalence Classes of Causal DAGs under Interventions](https://mlsanity.com/paper/1802.06310/) (2018)
- [Neural Production Systems: Learning Rule-Governed Visual Dynamics](https://mlsanity.com/paper/2103.01937/) (2021)
- [Causal Influence Detection for Improving Efficiency in Reinforcement Learning](https://mlsanity.com/paper/2106.03443/) (2021)
- [SQ-VAE: Variational Bayes on Discrete Representation with Self-annealed Stochastic Quantization](https://mlsanity.com/paper/2205.07547/) (2022)
- [Causal Dynamics Learning for Task-Independent State Abstraction](https://mlsanity.com/paper/2206.13452/) (2022)
- [Counterfactual Credit Assignment in Model-Free Reinforcement Learning](https://mlsanity.com/paper/2011.09464/) (2020)
- [Sample-Efficient Reinforcement Learning via Counterfactual-Based Data Augmentation](https://mlsanity.com/paper/2012.09092/) (2020)
- [Hierarchical Quantized Autoencoders](https://mlsanity.com/paper/2002.08111/) (2020)
- [Causal Curiosity: RL Agents Discovering Self-supervised Experiments for Causal Representation Learning](https://mlsanity.com/paper/2010.03110/) (2020)
- [Systematic Evaluation of Causal Discovery in Visual Model Based Reinforcement Learning](https://mlsanity.com/paper/2107.00848/) (2021)
- [An Investigation into Pre-Training Object-Centric Representations for Reinforcement Learning](https://mlsanity.com/paper/2302.04419/) (2023)
- [Vector Quantized Models for Planning](https://mlsanity.com/paper/2106.04615/) (2021)

## Built on by (most cited first, 11 of 11)

- [Learning Interactive World Model for Object-Centric Reinforcement Learning](https://mlsanity.com/paper/2511.02225/) (2025)
- [SPARTAN: A Sparse Transformer World Model Attending to What Matters](https://mlsanity.com/paper/2411.06890/) (2024)
- [Curious Causality-Seeking Agents Learn Meta Causal World](https://mlsanity.com/paper/2506.23068/) (2025)
- [Disentangling Dynamical Systems: Causal Representation Learning Meets Local Sparse Attention](https://mlsanity.com/paper/2603.14483/) (2026)
- [Learning Task-Sufficient World Models by Synergizing Agentic Exploration and Structured Modeling](https://mlsanity.com/paper/2607.04409/) (2026)
- [Towards Causal Representation Learning with Observable Sources as Auxiliaries](https://mlsanity.com/paper/2509.19058/) (2025)
- [Dynamic Sparsity: Challenging Common Sparsity Assumptions for Learning World Models in Robotic Reinforcement Learning Benchmarks](https://mlsanity.com/paper/2511.08086/) (2025)
- [On Transportability for Structural Causal Bandits](https://mlsanity.com/paper/2511.17953/) (2025)
- [Consistency evaluation of benchmarks used for causal discovery](https://mlsanity.com/paper/2606.01789/) (2026)
- [Learning Object Manipulation from Scratch via Contrastive Interaction](https://mlsanity.com/paper/2606.11525/) (2026)
- [From Image Interpretation to Clinical Reasoning: Upstream Physician-Context-Aware Multimodal Learning with Causal Reinforcement Learning](https://mlsanity.com/paper/2609.38924/) (2026)

## Abstract (arXiv)

Causal dynamics learning has recently emerged as a promising approach to enhancing robustness in reinforcement learning (RL). Typically, the goal is to build a dynamics model that makes predictions based on the causal relationships among the entities. Despite the fact that causal connections often manifest only under certain contexts, existing approaches overlook such fine-grained relationships and lack a detailed understanding of the dynamics. In this work, we propose a novel dynamics model that infers fine-grained causal structures and employs them for prediction, leading to improved robustness in RL. The key idea is to jointly learn the dynamics model with a discrete latent variable that quantizes the state-action space into subgroups. This leads to recognizing meaningful context that displays sparse dependencies, where causal structures are learned for each subgroup throughout the training. Experimental results demonstrate the robustness of our method to unseen states and locally spurious correlations in downstream tasks where fine-grained causal reasoning is crucial. We further illustrate the effectiveness of our subgroup-based approach with quantization in discovering fine-grained causal relationships compared to prior methods.