Learning to Control Coupled-Dynamics Environments with Joint Markov Decision Processes

Ege C. Kaya; Aliasghar Pourghani; Mahsa Ghasemi; Vijay Gupta; Abolfazl Hashemi

cs.LGAug 24, 2026

Learning to Control Coupled-Dynamics Environments with Joint Markov Decision Processes

Authors: Ege C. Kaya, Aliasghar Pourghani, Mahsa Ghasemi, Vijay Gupta, Abolfazl Hashemi

Organizations: Elmore Family School of Electrical and Computer Engineering Purdue University

Abstract

Coupled-dynamics environments expose the one-step outcomes that would follow from several possible counterfactual actions under a common realization of exogenous randomness. The ordinary Markov decision process formalism allows one to reason about the marginal law of each action but discards dependence across these counterfactual outcomes. The Joint Markov decision process (JMDP) formalism preserves that dependence. Prior work established the formalism and solved the fixed-policy joint moment evaluation problem in JMDPs. This paper develops optimal-control methods. We define a nonparametric distributional Bellman optimality operator for JMDPs, and prove that when the induced marginal MDP has a unique optimal policy, its iterates converge in Wasserstein distance to the optimal joint return law. For the first two moments, we establish convergence under a weaker condition that permits several mean-optimal actions as long as their tie resolutions share a second-moment fixed point. We also derive sampled targets for neural approximation.

Figures & tables

Figure 1: Exact tabular distributional control. In the top row, the optimal action is unique at every nonterminal state, and the joint law, mean vector, and covariance blocks converge to their ground-truth values. The bottom panels use an alternating sequence of mean-greedy tie resolutions, which switches between two optimal return laws.

Figure 2: Joint return law at the start state of the coupled reward chain. Dots are atoms, i.e., support points with positive probability. Contours smooth the same discrete mass for readability.

Figure 3: Moment control in a stochastic windy gridworld. The residual decays to zero, the learned covariance blocks match a separate fixed-policy JMDP moment evaluation, and finite-horizon Monte Carlo rollouts give an independent rollout-based comparison. The grayscale matrices show action correlations at representative states.

Figure 4: Route choice with coupled failures. Circles denote single-route rules and the square denotes the solution of equation 42 . “diag MV” sets off-diagonal covariances to zero, and “full MV” uses the full matrix. The right panel shows the resulting traffic allocation. “Superiority” uses $\arg\max_{a}\min_{b\neq a}\Pr(Z_{a}>Z_{b})$ .

Figure 5: Neural approximation in the coupled inverted-pendulum JMDP. All panels evaluate neural estimators. The first two show moment models and the last two show particle models. MAE denotes mean absolute error over off-diagonal entries. The covariance panels share a horizontal scale. Energy distance compares the learned and Monte Carlo joint laws, with lower values indicating better agreement. Dots are five seeds and bars show the standard error.

Semi-Markov Decision Process

See all 10 papers

Joint MDPs and Reinforcement Learning in Coupled-Dynamics Environments
Mar 6, 2026Ege C. Kaya, Mahsa Ghasemi, Abolfazl HashemiSemi-Markov Decision Process Markov Decision Processes
Elmore Family School of Electrical and Computer Engineering, Purdue University, West Lafayette, IN, USA
Lyapunov-Based Sample Complexity Analysis for Weakly-Coupled MDPs
Jun 12, 2026Tianhao Wu, Matthew Zurek, Weina Wang +1Multi-Armed Bandits Robust Markov Decision Processes
Department of Industrial and Systems Engineering, University of Wisconsin-Madison · Department of Computer Sciences, University of Wisconsin-Madison · Computer Science Department, Carnegie Mellon University
Policy Gradient for Continuous-Time Robust Markov Decision Processes
Jun 3, 2026Tanya Veeravalli, David M. Bossens, Atsushi NitandaRobust Markov Decision Processes Markov Decision Processes
Centre for Frontier AI Research, Agency for Science, Technology and Research (A*STAR) · Institute of High Performance Computing, Agency for Science, Technology and Research (A*STAR) · College of Computing and Data Science, Nanyang Technological University

Mar 6, 2026J/K move · Enter open · S save

Joint MDPs and Reinforcement Learning in Coupled-Dynamics Environments

Ege C. Kaya, Mahsa Ghasemi, Abolfazl Hashemi

Elmore Family School of Electrical and Computer Engineering, Purdue University, West Lafayette, IN, USA

Semi-Markov Decision Process Markov Decision Processes Distributional RL

Open paper AskPDF arXiv

Many distributional quantities in reinforcement learning are intrinsically joint across actions, including distributions of gaps and probabilities of superiority. However, the classical Markov decision process (MDP) formalism specifies only marginal laws and leaves the joint law of counterfactual one-step outcomes across multiple possible actions at a state unspecified. We study coupled-dynamics environments with a multi-action generative interface which can sample counterfactual one-step outcomes for multiple actions under shared exogenous randomness. We propose joint MDPs (JMDPs) as a formalism for such environments by augmenting an MDP with a multi-action sample transition model which specifies a coupling of one-step counterfactual outcomes, while preserving standard MDP interaction as marginal observations. We adopt and formalize a one-step coupling regime where dependence across actions is confined to immediate counterfactual outcomes at the queried state. In this regime, we derive Bellman operators for

n

th-order return moments, providing dynamic programming and incremental algorithms with convergence guarantees.

Abstract

Semi-Markov Decision Process

Explore similar work

Joint MDPs and Reinforcement Learning in Coupled-Dynamics Environments

Lyapunov-Based Sample Complexity Analysis for Weakly-Coupled MDPs

Policy Gradient for Continuous-Time Robust Markov Decision Processes