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.
Laboratory for Information and Decision Systems, Massachusetts Institute of Technology · Department of Industrial Engineering and Management Sciences, Northwestern University