Joint MDPs and Reinforcement Learning in Coupled-Dynamics Environments

arXiv:2603.06946 · cs.LG, math.OC · Submitted 2026-08-16 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Next we'll be talking about the paper "Joint MDPs and Reinforcement Learning in Coupled-Dynamics Environments".

Jane: The paper was written by Ege C. Kaya, Mahsa Ghasemi and Abolfazl Hashemi from Purdue University.

Tom: Stay tuned as we take you through the paper and discuss its implications.

Title: Tom: Welcome back to the show, everyone. Today we’re digging into a fresh arXiv preprint called “Joint MDPs and Reinforcement Learning in Coupled-Dynamics Environments.” Jane, I’ve got to say, the title alone had me leaning forward.

Jane: Same here, Tom. And for our listeners who aren’t deep in the reinforcement learning weeds, let me break down what that title actually means. An MDP, a Markov decision process, is the standard mathematical framework for how an agent learns to make decisions by trial and error. It’s the backbone of most modern reinforcement learning.

Tom: Right, and the paper argues that this backbone has a blind spot. The classic MDP only tells you what happens if you take one action at a time. It gives you the odds of a reward or a new state for each action separately, but it stays silent on what would happen if you tried several actions at the exact same moment, under the same random conditions.

Jane: Exactly. Think of it like weather forecasting. The MDP tells you the chance of rain if you stay home and the chance of rain if you go out, but it doesn’t tell you whether those two scenarios are linked. Maybe if it’s cloudy, both get rain, or maybe they’re opposites.

Tom: And that link, that coupling, is precisely what this paper formalizes. They call it a Joint MDP, or JMDP. The authors, Kaya, Ghasemi, and Hashemi from Purdue, are basically saying that the environment itself can have hidden structure that connects outcomes across actions.

Jane: So instead of just asking “what’s the reward for action A,” you can ask “what’s the reward for action A and action B together, given the same random event.” That’s a much richer question, and it opens the door to things like comparing actions fairly, which we’ll get into later.

Tom: And the kicker is that this isn’t just theoretical. In simulation environments, like when you’re testing a robot or a trading algorithm, you can actually sample these counterfactual outcomes. You can run multiple actions under the same random seed.

Jane: That’s the practical hook. The paper gives us a way to model that shared randomness and, crucially, to compute things like the probability that one action beats another. That’s a question you simply cannot answer with the old framework.

Tom: So we’ve got a new formalism, a new way to think about the environment, and a whole class of questions that suddenly become answerable. I’m curious to see how they actually build the math for this, because that’s where it gets tricky.

Jane: Absolutely. And that’s exactly what we’re going to dig into next. They’ve got to define how these joint outcomes behave over time, not just in one step. Stick around.

Paper discussion segment 2: Jane: So, Tom, we’ve established that a Joint MDP lets us see the coupling between actions. But the paper doesn’t stop at just defining the model. It actually builds the machinery to compute things with it.

Tom: Right, and that machinery is the meat of the paper. They focus on policy evaluation, which means they fix a strategy, a policy, and ask: what are the statistical properties of the long-term return? Usually you want the average return, but here they want more.

Jane: They want the full moment structure. So beyond the mean, they want the variance, and crucially, the mixed moments between different actions. That’s the covariance structure. That’s what tells you if two actions tend to do well together or if they’re negatively correlated.

Tom: And they do this under what they call a “one-step coupling regime.” That’s a really important assumption. It means the shared randomness only affects the immediate next step. After that, the future unfolds independently for each branch.

Jane: Let me put that in plain language. Imagine you’re at a fork in the road. The JMDP says the weather today is the same for both paths. But tomorrow, the weather on each path is its own separate thing. That keeps the math tractable, because you don’t have to track an exponentially growing tree of coupled futures.

Tom: Exactly. And with that assumption, they derive a Bellman operator. For the uninitiated, a Bellman operator is like a rule that says: the value of being here equals the reward now plus the discounted value of where you end up next. They build a version of that for these joint moments.

Jane: And the beautiful part is they prove it’s a contraction. That’s a mathematical guarantee that if you keep applying this rule over and over, you’ll converge to the true answer. It’s not a heuristic; it’s a certified algorithm.

Tom: They even give you a stopping criterion. You can measure the Bellman residual, which is basically how much your current guess changes after one application of the rule. If that residual is small, you know you’re close to the truth. That’s a huge practical advantage.

Jane: So we have a dynamic programming algorithm that’s guaranteed to work. But dynamic programming requires you to store a value for every state and every action. That’s fine for small problems, but it explodes for real-world ones.

Tom: And that’s the bridge to the next part of the paper. They don’t just stop at the tabular case. They also give us an incremental, sample-based version, which is the kind of thing you can actually run when the state space is huge.

Jane: So they’ve got the theory and the practical algorithm. I’m really curious about the experiments though. Did they actually show this working on something concrete?

Tom: They did, and we’ll get into those results in a moment. But first, let’s appreciate the leap here. They’ve taken a philosophical gap in the MDP formalism and turned it into a computable quantity.

Paper discussion segment 3: Tom: Alright, so we’ve got the theory down. Now let’s talk about what the paper actually did to prove it works. Jane, what did they run?

Jane: They ran two main tabular experiments. One is a Windy Gridworld, which is a classic navigation task where the wind pushes you around. The other is a Coupled-Reward Chain, which is a simpler setup where two actions give perfectly anti-correlated rewards.

Tom: And in both cases, they tracked the Bellman residual, that error certificate we talked about. The plots show it decaying linearly on a log scale, which is exactly what the contraction proof predicts. So the theory matches the practice.

Jane: But the more interesting part is what they visualize. They compute the correlation matrix between actions at each state. In the gridworld, you can literally see that the wind creates a structured, state-dependent correlation between moving up and moving down. That structure is completely invisible to a standard MDP.

Tom: That’s the money shot. It proves that the coupling isn’t just a theoretical curiosity. It’s a real, measurable property of the environment that affects the joint return distribution.

Jane: And they tie it back to the gap random variable. That’s the difference in return between two actions. They show that with the mixed moments from their algorithm, you can compute the variance of that gap. And from there, you can use something like Chebyshev’s inequality to bound the probability that one action beats another.

Tom: They validate those gap estimates against Monte Carlo simulation, and the agreement is strong. So the algorithm isn’t just converging to some abstract fixed point; it’s converging to the right numbers.

Jane: And then they scale it up. They take the incremental version and combine it with neural networks to handle Atari games. That’s a big jump from a small gridworld.

Tom: Right, and the results there show the TD errors dropping by orders of magnitude across several games like Pong and Boxing. It’s not perfect, but it’s a strong proof of concept that this can work beyond toy problems.

Jane: So the paper delivers on three fronts: a new formalism, a certified algorithm, and empirical validation that scales. That’s a complete package.

Tom: I want to bring in Meng here, because from an engineering standpoint, I’m wondering how hard it is to actually get that multi-action interface. The paper assumes you can query multiple actions under the same random seed. Is that realistic?

Meng: It is, actually, in a lot of simulators. If you have a physics engine or a financial market simulator, you can often just save the random seed, run action A, rewind, run action B. The paper is formalizing something that engineers have been doing informally for years. The value here is giving us a rigorous way to use that data.

Jane: And that’s a great segue to the bigger picture. We’ve got the formalism and the algorithms. What does this unlock for the field? What’s the vision?

Conclusion: Tom: So, as we wrap up our look at “Joint MDPs and Reinforcement Learning in Coupled-Dynamics Environments,” let’s pull it all together. Jane, what’s the one-line summary?

Jane: The paper says that the standard reinforcement learning framework ignores the hidden connections between what would happen if you took different actions at the same time, and it gives you a new framework, the JMDP, plus the math to actually compute those connections.

Tom: And it’s not just academic. They proved the algorithms converge, they showed the correlations exist in simple environments, and they demonstrated it scales to Atari. That’s a full arc from theory to practice.

Jane: The impact here is on decision-making under uncertainty. If you’re comparing a new drug to an old one, or a new trading strategy to a baseline, you want to know the probability that one is truly better. That’s a joint question, and this paper gives you the tools to answer it.

Tom: I also love that they’re honest about the limitations. The one-step coupling regime is a simplifying assumption. They’re not claiming to solve the fully coupled counterfactual tree problem, which would be exponentially hard.

Jane: Right, but they’ve carved out a tractable slice that’s still incredibly useful. And they’ve laid out a clear path for future work, like extending this to control, where you’re not just evaluating a fixed policy but actively improving it.

Tom: So we’re saying goodbye to this paper, but the ideas are going to stick with us. It’s a reminder that the way we model the world shapes the questions we can ask.

Jane: Well said, Tom. That’s a wrap on “Joint MDPs and Reinforcement Learning in Coupled-Dynamics Environments.” Thanks for joining us, and we’ll see you next time with a fresh paper to dig into.

Tom: Take care, everyone.

Ege C. Kaya, Mahsa Ghasemi, Abolfazl Hashemi

Purdue University

cs.LG, math.OC

Submitted: 2026-08-16

Updated: 2026-08-18

Comments: 25 pages, 7 figures

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 78/100

The gist: = Zπ(s, a) − Zπ(s, ã) and its distribution," "any tail functional of Gπ(s; a, ã), such as the quantile qα(Gπ(s; a, ã)) or the CVaRα(Gπ(s; a, ã))," and "the probability of superiority

Key concepts

Markov Decision Process (MDP): This is the standard mathematical framework used in reinforcement learning. It calculates rewards or new states based on taking one action at a time. However, it fails to account for how multiple actions interact when the environment shares a single random condition.
Joint MDP (JMDP): A new formalism that models the environment's hidden structure by allowing multiple actions simultaneously. It allows researchers to calculate the reward of action A and action B together, given a single shared random event, which is impossible with standard frameworks.
Bellman Operator: This is a dynamic programming tool used in policy evaluation. It calculates the long-term return by equating the current reward plus the discounted value of where you end up next. It is mathematically guaranteed to converge when applied repeatedly.
One-step Coupling Regime: A simplifying assumption where shared randomness between actions only affects the immediate next step. After this first step, future outcomes for each action unfold independently, making complex calculations tractable and manageable.

Terminology

Reinforcement Learning — 26× in this paper · explained in 4 episodes
Action Space — 6× in this paper · explained in 1 episode
Convergence Guarantees — 5× in this paper · explained in 3 episodes
Markov Decision Process (MDP) — 4× in this paper · explained in 11 episodes
Monte Carlo — 4× in this paper
Accuracy — 2× in this paper · explained in 1 episode
Generalization — 2× in this paper · explained in 13 episodes
Alignment — 1× in this paper · explained in 1 episode
Axioms — 1× in this paper · explained in 2 episodes
Generative Model — 1× in this paper · explained in 2 episodes
Objective Function — 1× in this paper · explained in 2 episodes
Q-learning — 1× in this paper · explained in 3 episodes
Reward Function — 1× in this paper · explained in 2 episodes
Trace — 1× in this paper · explained in 2 episodes

Summary

Summary

The paper introduces joint MDPs (JMDPs) as a formalism for coupled-dynamics environments, where the joint law of counterfactual one-step outcomes across multiple actions at a state is specified, going beyond the classical Markov decision process (MDP) formalism.

Motivation and Problem Setting. The authors note that many distributional quantities in reinforcement learning are intrinsically joint across actions, including the gap RV Gπ(s; a, ã):= Zπ(s, a) − Zπ(s, ã) and its distribution, any tail functional of Gπ(s; a, ã), such as the quantile qα(Gπ(s; a, ã)) or the CVaRα(Gπ(s; a, ã)), and "the probability of superiority P(Zπ(s, a) > Zπ(s, ã)). These quantities depend on the joint law of (Zπ(s, a), Zπ(s, ã)) and thus require knowledge of coupling structure. The paper states: A limitation of the Markov decision process (MDP) formalism. MDPs specify the marginal distributions of the reward and successor state under each action. However, they make no specification about the joint distribution of counterfactual, one-step outcomes across multiple actions at a state. Therefore, joint objects such as Law(Zπ(s, a) − Zπ(s, ã)) are not well-defined without adopting additional coupling conventions."

Coupled-Dynamics Environments. The paper formalizes environments where counterfactual one-step outcomes under multiple actions are operationally meaningful and sampleable, such as scenario-based simulation, where an exogenous disturbance U is realized at each decision point and one-step outcomes are functions of (s, a, U). Evaluating multiple candidate actions under the same realized disturbance is standard in simulation optimization and yields a multi-action generative interface.

Key Definitions.

m-JSTM (Definition 4.1): A kernel Jm(· s, a1:m) ∈ P(([0,1] × S)m) for m distinct actions such that each coordinate marginal matches the classical one-step marginals: if ((Ri, S′i))mi=1 ∼ Jm(· s, a1:m), then for each coordinate i, (Ri, S′i) ∼ (PR(· s, ai), PS(· s, ai)).
Coupled-dynamics environment (Definition 4.2): An environment with marginal kernels (PR, PS) is a coupled-dynamics environment if it admits an m-JSTM for some m ≥ 2 such that for some query, Jm is not the product coupling of its marginals.
JMDP (Definition 4.3): A quadruple (S, A, γ, J), where J(· s) ∈ P(([0,1] × S)N) is a Markov kernel over counterfactual one-step outcome tables ((R(a), S′(a)))a∈A. At time t, conditional on St = s, the environment samples a one-step outcome table; the agent selects action At to execute, and the realized transition is given by the executed action. The remaining coordinates are counterfactual.

One-Step Coupling Regime. The paper adopts Assumption 4.5 (Exogenous noise representation): there exist an i.i.d. exogenous sequence (Ut)t≥0 and measurable maps g: S × A × U → [0,1], h: S × A × U → S such that R(a)t = g(s, a, Ut) and S′(a)t+1 = h(s, a, Ut). The paper states: the dependence across different actions is entirely mediated by the current table draw, and does not persist through future shared noise. This regime avoids the exponential blow-up associated with fully-coupled counterfactual trees.

Main Theoretical Contributions.

2nd-order joint Bellman operator (Definition 5.1): The operator Tπ2 maps a moment collection M = (Mµ, MΣ) to Tπ2M coordinatewise. For first moments: (Tπ2M)µ(s, a):= E[R(a) + γMµ(S′(a), A′(a)) S = s]. For second moments: (Tπ2M)Σ(s, a, s̃, ã):= E[R(a)R̃(ã) + γR(a)Mµ(S̃′(ã), Ã′(ã)) + γR̃(ã)Mµ(S′(a), A′(a)) + γ2MΣ(S′(a), A′(a), S̃′(ã), Ã′(ã)) (S, S̃) = (s, s̃)].
Contraction and convergence (Lemma 5.2, Theorem 5.3): With norm ∥M∥λ:= max ∥Mµ∥∞, (1/λ)∥MΣ∥∞ where λ:= 2/(1−γ), Tπ2 is a γ-contraction. The iteration Mk+1 = Tπ2Mk converges geometrically with rate γ: ∥Mk − Mπ2∥λ ≤ γk∥M0 − Mπ2∥λ. The Bellman residual provides a computable stopping criterion: ∥M − Mπ2∥λ ≤ (1/(1−γ))∥M − Tπ2M∥λ.
Incremental JIPE-2 (Section 5.2): A stochastic approximation variant using one-sample backups T̂π2(M, i) for each index type (µ-index, diagonal Σ-index, off-diagonal Σ-index). Theorem 5.4 establishes almost-sure convergence under standard Robbins-Monro step-size conditions and infinite visitation of each coordinate.
Generalization to higher moments (Section 5.3): The construction extends to nth-order moments via the operator Tπn defined through the system of equations (TnπM)k(x1:k):= E[ΣI⊆[k] γI (Πi∉I Ri) MI(X′I) x1:k], with norm ∥M∥λ:= maxk∈[n] (1/λk)∥M(k)∥∞ for λk:= (2/(1−γ))k−1. The operator is a γ-contraction with unique fixed point the true collection of moments up to nth order.
Function approximation (Appendix B): For high-dimensional settings, the paper proposes a projected JIPE-2 with linear function approximation: M̂µ(s, a; θµ):= ϕµ(s, a)⊤θµ and M̂Σ(s, a, s̃, ã; θΣ):= ϕµ(s, a)⊤ΘΣϕµ(s̃, ã) with ΘΣ ∈ ΘPSD. Under Assumptions B.1–B.3 (including a concentration coefficient condition γ2√cρ < 1), Theorem B.4 establishes that the projected operator ΠTπ2 is a κ-contraction with κ < 1, yielding a unique fixed point and an approximation error bound.

Connection to Gap RV (Section 5.4). The paper shows that joint structure becomes relevant for comparative or risk-sensitive questions about gaps. The variance of the gap is given by: var(Gπ(s; a, ã)) = Σπ(s, a, s, a) − Σπ(s, ã, s, ã) − 2Σπ(s, a, s, ã) − (µπ(s, a) − µπ(s, ã))2, where all moments are computable through JIPE-2. The paper also derives a Chebyshev-based upper bound on the inferiority probability P(Gπ(s; a, ã) ≤ 0) ≤ σ2G/(σ2G + µ2G).

Experiments (Section 6). The paper validates the theory in four ways:

Tabular coupled-dynamics environments: Windy gridworld (WGW) with leftward wind gust coupling counterfactual next states, and coupled-reward chain (CRC) with anti-correlated rewards. Figure 2 shows linear decay of the Bellman residual on a log scale, consistent with γ-contraction.
Visualization of joint structure: Figure 3 shows per-state 4×4 action correlation matrices in WGW, indicating structured, state-dependent joint law across actions which is invisible to an MDP marginal description.
Gap validation: Figure 5 compares JIPE-2-derived predictions of E[Gπ] and var(Gπ) to Monte Carlo estimates, showing agreement. The empirical cdf of the ratio P̂(Gπ ≤ 0)/Chebyshev(Gπ) shows ratios bounded by 1, indicating empirical non-violation of the upper bound.
Large-scale experiments: Incremental JIPE-2 with neural function approximation in four ALE environments (Pong, BattleZone, Boxing, Atlantis) with a coupled-dynamics interface. Figure 7 shows moving-average TD errors decreasing by several orders of magnitude across all moment blocks (µ, σdiag, σsame, σcross).

Conclusion. The paper concludes: "We introduced JMDPs for coupled-dynamics environments that explicitly model the joint law of counterfactual outcomes across actions. Under a one-step coupling regime, we derived Bellman operators for return moments and obtained convergence guarantees for DP algorithms and their incremental variants, along with a projected function-approximation formulation. This makes intrinsically joint quantities such as gap moments and inferiority-probability bounds computable from multi-action simulator access. A natural next step is control: extending joint moment evaluation to policy improvement under joint distributional objectives in JMDPs."

Improvements for AI systems

Based on the paper, here are the specific improvements I can make to AI systems and what the improved systems can do:

Implementation: Extend standard RL environments (e.g., Gym, ALE) with a multi action sample transition(s, a list) interface that returns coupled counterfactual outcomes (R i, S'i) for multiple actions under shared exogenous noise U t.
What it enables: The agent can now query multiple actions at the same state under identical randomness, enabling estimation of joint return moments (e.g., E[Z(s,a)Z(s,ã)]) that are undefined in standard MDPs.
Implementation: Add a policy evaluation module that maintains moment collections M = (M(1),..., M(n)) and applies the operator T n π from Eq. (33) with the contraction norm ∥M∥ λ = max k ∥M(k)∥∞ / λ k where λ k = (2/(1-γ))(k-1).
What it enables: The system can compute exact mixed moments of the joint return vector Z π(s) up to arbitrary order n, with geometric convergence at rate γ and a computable Bellman residual certificate ∥M - T 2 π M∥ λ for stopping criteria.
Implementation: Replace exact DP with the recursion in Eq. (28) using one-sample backups T̂ 2 π(M, i) from Eqs. (25)–(27), with step sizes satisfying Σ α k(i) = ∞ and Σ α k(i) squared < ∞ per coordinate.
What it enables: The system can learn joint moments online from simulator queries without full state enumeration, with almost-sure convergence in ∥·∥ λ (Theorem 5.4).
Implementation: Parameterize M Σ(s,a,s̃,ã) = φ(s,a)T Θ Σ φ(s̃,ã) with Θ Σ ⪰ 0, and project via Eq. (71) onto the PSD cone after each Bellman backup.
What it enables: The system scales to high-dimensional/continuous state spaces while preserving valid second-moment geometry (positive semi-definite covariance matrices), with a contractive projected operator under Assumption B.3.
Implementation: Using learned mixed moments, compute gap variance via Eq. (35): var(G π(s;a,ã)) = Σ(s,a,s,a) + Σ(s,ã,s,ã) - 2Σ(s,a,s,ã) - (µ(s,a)-µ(s,ã))2.
What it enables: The system can produce Chebyshev-style upper bounds on P(G π ≤ 0) (inferiority probability) and estimate tail functionals (quantiles, CVaR) of gap distributions, which are impossible with marginal-only methods.

Answer joint queries: Given a state s and actions a, ã, return E[Z(s,a)Z(s,ã)], var(Z(s,a)-Z(s,ã)), and P(Z(s,a) > Z(s,ã)) with certified accuracy (via Bellman residual).
Perform risk-aware policy comparison: For a fixed policy, compute per-state action correlation matrices ρ s(a,ã) (as in Figure 3), revealing hidden coupling structure that marginal MDPs cannot see.
Provide safety certificates: For safety-critical decisions, output an upper bound on the probability that an alternative action outperforms the current policy's action, using moment-based inequalities (e.g., Chebyshev).
Scale to complex environments: In Atari-like settings with coupled dynamics (shared emulator randomness), the system can train neural networks to estimate joint moments with stable TD-error descent (as in Figure 7), enabling practical deployment.
Support counterfactual reasoning: Given a simulator with a multi-action generative interface, the system can estimate the joint law of counterfactual one-step outcomes and propagate this through time under the one-step coupling regime, enabling principled what-if analysis across actions.
Provide convergence guarantees: Unlike heuristic distributional RL, the system offers provable γ-contraction, geometric convergence, and residual-based stopping, making it suitable for high-assurance applications (e.g., robotics, finance).

Abstract

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.