Towards Causal Model-Based Policy Optimization

• URL: http://arxiv.org/abs/2503.09719v1
• Date: Wed, 12 Mar 2025 18:09:02 GMT
• Title: Towards Causal Model-Based Policy Optimization
• Abstract summary: We introduce Causal Model-Based Policy Optimization (C-MBPO)<n>C-MBPO is a novel framework that integrates causal learning into the Model-Based Reinforcement Learning pipeline.<n>We show that C-MBPO can be shown to be robust to a class of distributional shifts that affect spurious, non-causal relationships in the dynamics.
• Score: 0.24578723416255752
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Real-world decision-making problems are often marked by complex, uncertain dynamics that can shift or break under changing conditions. Traditional Model-Based Reinforcement Learning (MBRL) approaches learn predictive models of environment dynamics from queried trajectories and then use these models to simulate rollouts for policy optimization. However, such methods do not account for the underlying causal mechanisms that govern the environment, and thus inadvertently capture spurious correlations, making them sensitive to distributional shifts and limiting their ability to generalize. The same naturally holds for model-free approaches. In this work, we introduce Causal Model-Based Policy Optimization (C-MBPO), a novel framework that integrates causal learning into the MBRL pipeline to achieve more robust, explainable, and generalizable policy learning algorithms. Our approach centers on first inferring a Causal Markov Decision Process (C-MDP) by learning a local Structural Causal Model (SCM) of both the state and reward transition dynamics from trajectories gathered online. C-MDPs differ from classic MDPs in that we can decompose causal dependencies in the environment dynamics via specifying an associated Causal Bayesian Network. C-MDPs allow for targeted interventions and counterfactual reasoning, enabling the agent to distinguish between mere statistical correlations and causal relationships. The learned SCM is then used to simulate counterfactual on-policy transitions and rewards under hypothetical actions (or ``interventions"), thereby guiding policy optimization more effectively. The resulting policy learned by C-MBPO can be shown to be robust to a class of distributional shifts that affect spurious, non-causal relationships in the dynamics. We demonstrate this through some simple experiments involving near and far OOD dynamics drifts.

Related papers

• IADD-TR: Intervention-Aware Dynamics Decoupling with Targeted Regularization for Model-Based Reinforcement Learning [57.16513298507272]
  We propose IADD-TR, a unified framework combining Intervention-Aware Dynamics Decoupling (IADD) and Targeted Regularization (TR)<n>IADD factorizes transitions into an action-intervention stage and an action-free natural evolution stage, using a zero-action anchor to resolve the non-uniqueness of this two-stage factorization for robust generalization.<n>For policy learning, we derive TR from the efficient influence function of a replay-state policy-gradient functional.
  arXiv  Detail & Related papers  (2026-08-11T08:20:02Z)
• From Noise to Control: Parameterized Diffusion Policies [54.941343992110426]
  We propose a framework for learning diffusion policies conditioned on continuous parameters embedded in a learned behavior manifold.<n>Our approach enables smooth distances between known strategies and efficient adaptation to novel constraints without updating policy weights.
  arXiv  Detail & Related papers  (2026-05-29T20:21:50Z)
• On the Plasticity and Stability for Post-Training Large Language Models [54.757672540381236]
  We identify a root cause as the conflict between plasticity and stability gradients.<n>We propose Probabilistic Conflict Resolution (PCR), a framework that models gradients as random variables.<n>PCR significantly smooths the training trajectory and achieves superior performance in various reasoning tasks.
  arXiv  Detail & Related papers  (2026-02-06T07:31:26Z)
• Rethinking the Trust Region in LLM Reinforcement Learning [72.25890308541334]
  Proximal Policy Optimization (PPO) serves as the de facto standard algorithm for Large Language Models (LLMs)<n>We propose Divergence Proximal Policy Optimization (DPPO), which substitutes clipping with a more principled constraint.<n>DPPO achieves superior training and efficiency compared to existing methods, offering a more robust foundation for RL-based fine-tuning.
  arXiv  Detail & Related papers  (2026-02-04T18:59:04Z)
• Model-Based Reinforcement Learning Under Confounding [3.5690236380446163]
  We investigate model-based reinforcement learning in contextual Markov decision processes (C-MDPs) in which the context is unobserved and induces confounding in the offline dataset.<n>We adapt a proximal off-policy evaluation approach that identifies the confounded reward expectation using only observable state-action-reward trajectories under mild invertibility conditions on proxy variables.<n>The proposed formulation enables principled model learning and planning in confounded environments where contextual information is unobserved, unavailable, or impractical to collect.
  arXiv  Detail & Related papers  (2025-12-08T13:02:00Z)
• Multi-Path Collaborative Reasoning via Reinforcement Learning [54.8518809800168]
  Chain-of-Thought (CoT) reasoning has significantly advanced the problem-solving capabilities of Large Language Models (LLMs)<n>Recent methods attempt to address this by generating soft abstract tokens to enable reasoning in a continuous semantic space.<n>We propose Multi-Path Perception Policy Optimization (M3PO), a novel reinforcement learning framework that explicitly injects collective insights into the reasoning process.
  arXiv  Detail & Related papers  (2025-12-01T10:05:46Z)
• Causal Model-Based Reinforcement Learning for Sample-Efficient IoT Channel Access [39.76683291751265]
  Multi-agent reinforcement learning (MARL) for wireless use case such as medium access control (MAC) is hindered by their sample inefficiency.<n>In this paper, a novel causal model-based MARL framework is developed by leveraging tools from causal learn- ing.<n>The proposed approach can reduce environment interactions by 58%, and yield faster convergence compared to model-free baselines.
  arXiv  Detail & Related papers  (2025-11-13T13:26:33Z)
• Efficient Solution and Learning of Robust Factored MDPs [57.2416302384766]
  Learning r-MDPs from interactions with an unknown environment enables the synthesis of robust policies with provable guarantees on performance.<n>We propose novel methods for solving and learning r-MDPs based on factored state representations.
  arXiv  Detail & Related papers  (2025-08-01T15:23:15Z)
• CTRLS: Chain-of-Thought Reasoning via Latent State-Transition [57.51370433303236]
  Chain-of-thought (CoT) reasoning enables large language models to break down complex problems into interpretable intermediate steps.<n>We introduce groundingS, a framework that formulates CoT reasoning as a Markov decision process (MDP) with latent state transitions.<n>We show improvements in reasoning accuracy, diversity, and exploration efficiency across benchmark reasoning tasks.
  arXiv  Detail & Related papers  (2025-07-10T21:32:18Z)
• Better Decisions through the Right Causal World Model [17.623937562865617]
  Causal Object-centric Model Extraction Tool (COMET) is a novel algorithm designed to learn the exact interpretable causal world models (CWMs) Our results, validated in Atari environments such as Pong and Freeway, demonstrate the accuracy and robustness of COMET.
  arXiv  Detail & Related papers  (2025-04-09T20:29:13Z)
• Understanding Endogenous Data Drift in Adaptive Models with Recourse-Seeking Users [6.782864450313782]
  We study user strategic behaviors and their interactions with decision-making systems under resource constraints and competitive dynamics.<n>We propose two methods--Fair-top-k and Dynamic Continual Learning--which significantly reduce recourse cost and improve model robustness.<n>Our findings draw connections to economic theories, highlighting how algorithmic decision-making can unintentionally reinforce a higher standard and generate endogenous barriers to entry.
  arXiv  Detail & Related papers  (2025-03-12T12:17:34Z)
• SAMBO-RL: Shifts-aware Model-based Offline Reinforcement Learning [9.88109749688605]
  Model-based offline reinforcement learning trains policies using pre-collected datasets and learned environment models.<n>This paper offers a comprehensive analysis that disentangles the problem into two fundamental components: model bias and policy shift.<n>We introduce Shifts-aware Model-based Offline Reinforcement Learning (SAMBO-RL), a practical framework that efficiently trains classifiers to approximate SAR for policy optimization.
  arXiv  Detail & Related papers  (2024-08-23T04:25:09Z)
• Decision Mamba: A Multi-Grained State Space Model with Self-Evolution Regularization for Offline RL [57.202733701029594]
  We propose Decision Mamba, a novel multi-grained state space model (SSM) with a self-evolving policy learning strategy. To address these challenges, we propose Decision Mamba, a novel multi-grained state space model (SSM) with a self-evolving policy learning strategy. To mitigate the overfitting issue on noisy trajectories, a self-evolving policy is proposed by using progressive regularization.
  arXiv  Detail & Related papers  (2024-06-08T10:12:00Z)
• Learning a Diffusion Model Policy from Rewards via Q-Score Matching [93.0191910132874]
  We present a theoretical framework linking the structure of diffusion model policies to a learned Q-function.<n>We propose a new policy update method from this theory, which we denote Q-score matching.
  arXiv  Detail & Related papers  (2023-12-18T23:31:01Z)
• Model-based Causal Bayesian Optimization [74.78486244786083]
  We introduce the first algorithm for Causal Bayesian Optimization with Multiplicative Weights (CBO-MW) We derive regret bounds for CBO-MW that naturally depend on graph-related quantities. Our experiments include a realistic demonstration of how CBO-MW can be used to learn users' demand patterns in a shared mobility system.
  arXiv  Detail & Related papers  (2023-07-31T13:02:36Z)
• Causal Disentangled Variational Auto-Encoder for Preference Understanding in Recommendation [50.93536377097659]
  This paper introduces the Causal Disentangled Variational Auto-Encoder (CaD-VAE), a novel approach for learning causal disentangled representations from interaction data in recommender systems. The approach utilizes structural causal models to generate causal representations that describe the causal relationship between latent factors.
  arXiv  Detail & Related papers  (2023-04-17T00:10:56Z)
• When to Update Your Model: Constrained Model-based Reinforcement Learning [50.74369835934703]
  We propose a novel and general theoretical scheme for a non-decreasing performance guarantee of model-based RL (MBRL) Our follow-up derived bounds reveal the relationship between model shifts and performance improvement. A further example demonstrates that learning models from a dynamically-varying number of explorations benefit the eventual returns.
  arXiv  Detail & Related papers  (2022-10-15T17:57:43Z)
• Model-Based Offline Reinforcement Learning with Pessimism-Modulated Dynamics Belief [3.0036519884678894]
  Model-based offline reinforcement learning (RL) aims to find highly rewarding policy, by leveraging a previously collected static dataset and a dynamics model. In this work, we maintain a belief distribution over dynamics, and evaluate/optimize policy through biased sampling from the belief. We show that the biased sampling naturally induces an updated dynamics belief with policy-dependent reweighting factor, termed Pessimism-Modulated Dynamics Belief.
  arXiv  Detail & Related papers  (2022-10-13T03:14:36Z)
• Causal Dynamics Learning for Task-Independent State Abstraction [61.707048209272884]
  We introduce Causal Dynamics Learning for Task-Independent State Abstraction (CDL) CDL learns a theoretically proved causal dynamics model that removes unnecessary dependencies between state variables and the action. A state abstraction can then be derived from the learned dynamics.
  arXiv  Detail & Related papers  (2022-06-27T17:02:53Z)
• Revisiting Design Choices in Model-Based Offline Reinforcement Learning [39.01805509055988]
  Offline reinforcement learning enables agents to leverage large pre-collected datasets of environment transitions to learn control policies. This paper compares and designs novel protocols to investigate their interaction with other hyper parameters, such as the number of models, or imaginary rollout horizon.
  arXiv  Detail & Related papers  (2021-10-08T13:51:34Z)
• Learning Robust Models Using The Principle of Independent Causal Mechanisms [26.79262903241044]
  We propose a new gradient-based learning framework whose objective function is derived from the ICM principle. We show theoretically and experimentally that neural networks trained in this framework focus on relations remaining invariant across environments.
  arXiv  Detail & Related papers  (2020-10-14T15:38:01Z)
• Invariant Causal Prediction for Block MDPs [106.63346115341862]
  Generalization across environments is critical to the successful application of reinforcement learning algorithms to real-world challenges. We propose a method of invariant prediction to learn model-irrelevance state abstractions (MISA) that generalize to novel observations in the multi-environment setting.
  arXiv  Detail & Related papers  (2020-03-12T21:03:01Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.