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 nth-order return moments, providing dynamic programming and incremental algorithms with convergence guarantees.
Figures & tables
Figure 1: A 3×3 WGW environment and the evaluated policy.
Figure 2: Bellman residual convergence in tabular coupled-dynamics environments. Left: 5×5 WGW. Right: CRC with ∣S∣=25 . We plot the Bellman residual ∥Mk−T2πMk∥λ on logarithmic scale.
Figure 3: One CRC transition. Both actions reach the same next state but receive anti-correlated rewards.
Figure 4: Per-state action correlation matrices in WGW. Each tile corresponds to a state in a 3×3 WGW, and displays the 4×4 correlation matrix ρsπ computed through JIPE- 2 .
Figure 5: Correlation matrix ρs0π computed from the converged JIPE- 2 moments in CRC.
Figure 6: Gap validation in WGW. We evaluate a fixed goal-directed policy and consider gaps between the policy’s action and alternatives. Left: Predicted vs. MC gap means E[Gπ] . Middle: Predicted vs. MC gap variances var(Gπ) . Right: Empirical cdf of the ratio P^(Gπ≤0)/Chebyshev(Gπ) . Ratios closer to 1 indicate tighter upper bounds, while small ratios indicate looseness. The blue curve computes the denominator using JIPE- 2 -estimated moments, while the orange curve uses MC-estimated moments (a near-ground-truth proxy), separating intrinsic bound looseness from moment-estimation error.
Figure 7: Incremental JIPE- 2 with neural functional approximation on coupled ALE environments. Moving-average TD errors (log scale) during SGD-based incremental JIPE- 2 with neural function approximation, shown for Pong, BattleZone, Boxing, and Atlantis (left-to-right, top-to-bottom). The plotted coordinates corresponds to the JIPE- 2 blocks: μ is the 1 st-moment (mean-return) component, σdiag is the block where s=s~ and a=a~ . σsame is the block where s=s~ , a=a~ . σcross is the block where s=s~ , a=a~ . Bold curves show the across-seed mean TD error, and faint curves show the individual trajectories from three random seeds . Across all four ALE games, each block exhibits stable descent over training, supporting the practicality of the incremental JIPE- 2 updates under nonlinear function approximation. Gap moments are evaluated separately against recursive-pair MC in Appendix C .
Appendix figures & tables1 asset
Supplementary material from the paper’s appendix.
Appendix
Game
Coupled / MC
Independent / MC
Atlantis
1.03
83.3
BattleZone
3.87
115.7
Boxing
3.97
111.1
Pong
1.30
135.4
Appendix
Table 1: Reward-coupled ALE gap-variance validation. For each game, we sample emulator states, estimate the action-gap variance from recursive-pair MC, and report the ratios of the learned coupled prediction and an independent-marginals prediction to the same MC estimate. Values near 1.0 represent better calibrated estimates to the coupled MC gap variance.
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.
Ege C. Kaya, Aliasghar Pourghani, Mahsa Ghasemi +2
Elmore Family School of Electrical and Computer Engineering Purdue University
Reinforcement learning algorithms are typically designed for generic Markov Decision Processes (MDPs), where any state-action pair can lead to an arbitrary transition distribution. In many practical systems, however, only a subset of the state variables is directly influenced by the agent's actions, while the remaining components evolve according to exogenous dynamics and account for most of the stochasticity. In this work, we study a structured class of MDPs characterized by exogenous state components whose transitions are independent of the agent's actions. We show that exploiting this structure yields significantly improved learning guarantees, with only the size of the exogenous state space appearing in the leading terms of the regret bounds. We further establish a matching lower bound, showing that this dependence is information-theoretically optimal. Finally, we empirically validate our approach across classical toy settings and real-world-inspired environments, demonstrating substantial gains in sample efficiency compared to standard reinforcement learning methods.
Davide Maran, Davide Salaorni, Marcello Restelli
Department of DEIB, Politecnico di Milano, Milan, Italy.
We study a structured class of Markov Decision Processes, known as Exo-MDPs, in which the state space is partitioned into exogenous and endogenous components. Exogenous states evolve stochastically, independent of the agent's actions, while endogenous states evolve deterministically based on both state components and actions. Exo-MDPs capture many operations research settings, including inventory control, resource management, and ride-sharing. Our first contribution is structural: we establish a representational equivalence between discrete MDPs, Exo-MDPs, and discrete linear mixture MDPs. Our second contribution is statistical. We characterize the minimax regret of learning in Exo-MDPs when the effective dimension r is small relative to the endogenous state and action spaces. When the exogenous states are unobserved, we prove matching upper and lower regret bounds of order Θ(HrK) over K episodes of horizon H, where r is the effective dimension of the Exo-MDP. When exogenous states are observed, the minimax regret improves to Θ(HrK), revealing a Θ(r) statistical gap due to observation of the exogenous states. These results show that Exo-MDPs decouple sample complexity from action space and endogenous state space. We validate these insights with experiments on inventory control and resource allocation.
Jia Wan, Sean R. Sinclair, Devavrat Shah +1
Laboratory for Information and Decision Systems, Massachusetts Institute of Technology · Department of Industrial Engineering and Management Sciences, Northwestern University