Research map / Probabilistic and scientific modeling / Causal discovery
Fine-Grained Causal Dynamics Learning with Quantization for Improving Robustness in Reinforcement Learning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes FCDL, a causal dynamics model that jointly learns a VQ codebook quantizing the state-action space into subgroups, with a local causal graph per subgroup, improving robustness of model-based RL to locally spurious correlations and unseen states.
Problem
Existing causal dynamics models infer only a global causal graph and miss context-dependent (fine-grained) causal relationships, while sample-specific approaches lack interpretable contexts; this hurts robustness in RL to unseen states and locally spurious correlations.
Approach
Quantize the state-action space with a discrete latent codebook (vector quantization) so each code pairs a subgroup (context) with a local causal graph (adjacency matrix), learned end-to-end with the dynamics model under sparsity regularization, with identifiability guarantees for local causal graphs and contexts.
Key results
- On Chemical full-fork test with n=2 noisy nodes, FCDL achieves 15.27 average reward vs 10.95 for NCD and 9.37 for CDL.
- On Chemical full-chain ID-to-OOD prediction accuracy at n=2, FCDL reaches 60.34% vs 40.04% for NCD and 34.90% for CDL.
- On Magnetic test, FCDL achieves 4.81 average reward vs 1.31 for NCD and 1.10 for CDL.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use subgroup quantization to infer context-specific local causal graphs instead of sample-specific dependency inference | FCDL versus NCD · Chemical full-fork and full-chain downstream tasks | average episode reward: 10.95 -> 15.27; 9.11 -> 14.73; 10.32 -> 13.62; 9.60 -> 13.36; 8.86 -> 12.35; 10.32 -> 12.00 | Helped |
| Use subgroup quantization to infer context-specific local causal graphs instead of sample-specific dependency inference | FCDL versus NCD · Magnetic Test | average episode reward: 1.31 -> 4.81 | Helped |
| Use subgroup quantization to infer context-specific local causal graphs instead of sample-specific dependency inference | FCDL versus NCD · Chemical downstream tasks with noisy nodes | prediction accuracy: 41.60 -> 66.44; 37.47 -> 58.49; 42.27 -> 49.09; 40.04 -> 60.34; 37.47 -> 56.64; 41.19 -> 53.29 | Helped |
| Set quantization degree K=2 rather than other tested quantization degrees (ablation) | FCDL · Chemical full-fork downstream tasks | average episode reward | Hurt |
| Use quantization degrees other than K=2 (ablation) | FCDL · Chemical full-fork downstream tasks | average episode reward | Helped |
| Set quantization degree K=1, yielding one global causal graph (ablation) | FCDL · Magnetic non-magnetic samples | structural Hamming distance: mean SHD score of 6 | Hurt |
Surprising findings
- Codebook size K=2 causes instability and codebook collapse, while K>2 trains stably.
Limitations
- Evaluated only on state-based environments with fully observable factored states; extension to high-dimensional image observations is left open.
- Choosing the number of codes K is not studied in depth; over-parametrization is assumed to be cheap.
- Identifiability guarantees rely on assumptions such as E-faithfulness and sufficient quantization degree K>=H.
Experimental setup
Datasets: Chemical Magnetic
Benchmarks: Chemical (full-fork) Chemical (full-chain) Magnetic
Metrics: average episode reward prediction accuracy structural hamming distance
Baselines: MLP Modular GNN NPS CDL GRADER Oracle NCD
Builds on: CDL NCD vector quantization (VQ-VAE) context-set specific independence Gumbel-Softmax
Closest papers by meaning
- Tackling Non-Stationarity in Reinforcement Learning via Causal-Origin Representation (ICML 2024 · similarity 0.87)
- Building Minimal and Reusable Causal State Abstractions for Reinforcement Learning (AAAI 2024 · similarity 0.87)
- Towards Generalizable Reinforcement Learning via Causality-Guided Self-Adaptive Representations (ICLR 2025 · similarity 0.87)
- Hierarchical Reinforcement Learning with Targeted Causal Interventions (ICML 2025 · similarity 0.87)
- Learning Causal Dynamics Models in Object-Oriented Environments (ICML 2024 · similarity 0.87)
- Modeling Unseen Environments with Language-guided Composable Causal Components in Reinforcement Learning (ICLR 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.