Latent Causal Discovery in Reinforcement Learning and Large Language Models

Overview

Latent causal discovery seeks to uncover hidden causal mechanisms in data where the true causal variables are unobserved. Our research develops theoretical foundations and practical algorithms for identifying latent causal structures, particularly in reinforcement learning (RL) and large language models (LLMs).

Research Focus

Causal Reinforcement Learning

In RL, learning disentangled causal state representations is crucial for robust decision-making, generalization, and transfer learning. Our work challenges traditional assumptions and redefines state disentanglement through causal constraints and interventions.

• Rethinking State Disentanglement in Causal RL
  • Proposing a new framework for interpretable and robust state representations.
• Silver Linings: When Distribution Shifts Enhance Identifiability
  • Investigating favorable shifts that improve latent variable identifiability in RL.

Causal Representation Learning in Large Language Models

Do LLMs learn meaningful causal representations? We analyze whether next-token prediction is sufficient for learning human-interpretable causal concepts and propose methods to enhance causal learning.

• I Predict Therefore I Am
  • Examining whether transformer-based models implicitly encode causal structures.
• Identifiable Latent Polynomial Causal Models
  • Leveraging distribution changes to identify latent causal factors in LLMs.

Applications

🚀 AI for Science – Improving interpretability and robustness in scientific AI applications.
🤖 Autonomous Systems – Enabling RL agents to adapt to dynamic environments.
⚖️ Fair & Robust AI – Reducing bias by ensuring models learn true causal relationships rather than spurious correlations.

Selected Publications

📄 ICLR 2024 – Identifiable Latent Polynomial Causal Models Through the Lens of Change.
📄 JMLR (Submitted) – Identifying Weight-Variant Latent Causal Models.
📄 ICML 2025 (Submitted) – Rethinking State Disentanglement in Causal Reinforcement Learning.
📄 ICML 2025 (Submitted) – Silver Linings: On the Types of Distribution Shifts that Enhance Identifiability in Causal Representation Learning.
📄 ICML 2025 (Submitted) – I Predict Therefore I Am: Is Next Token Prediction Enough to Learn Human-Interpretable Concepts from Data?

Get in Touch

For further details or collaboration opportunities, feel free to reach out! 🚀