Source-linked AI summary
Learning to Control Coupled-Dynamics Environments with Joint Markov Decision Processes
Ege C. Kaya, Aliasghar Pourghani, Mahsa Ghasemi, Vijay Gupta, Abolfazl Hashemi
TL;DR
Ordinary MDPs discard dependence among counterfactual action outcomes, motivating control methods that preserve this joint information in JMDPs. The paper defines a distributional Bellman optimality operator, proves convergence under unique optimal policies and a weaker second-moment condition for ties, and derives sampled targets for neural approximation. The results support recovering joint quantities from coupled outcome tables.
Problem
Ordinary MDPs retain marginal action laws but discard dependence among counterfactual outcomes, although that dependence is needed for joint return quantities.
Method
The paper defines a nonparametric JMDP Bellman optimality operator and sampled first- and second-moment targets that retain successor state–action pairs.
Results
The distributional iterates converge under a unique optimal policy, while first two moments converge under a weaker common second-moment fixed-point condition for optimal ties.
Takeaways & Limitations
JMDP control preserves the mean-optimal policy of the induced MDP while determining how counterfactual action returns move together.
Takeaways & Limitations
Moment convergence with multiple optimal actions requires every optimal selector to share the same second-moment fixed point; differing variances or covariances can violate this condition.
Abstract
from arXiv · showhide
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.
1 Introduction
The paper extends JMDPs from fixed-policy evaluation to optimal control in coupled-dynamics environments, preserving dependence among counterfactual action outcomes. It introduces distributional and moment-level control methods, convergence results, sampled targets, and empirical validation.
- Coupled-dynamics environments expose several counterfactual action outcomes under shared randomness, whereas ordinary MDPs retain only action-wise marginal laws.
- The paper defines a nonparametric JMDP Bellman optimality operator for joint return laws.
- Convergence holds for risk-neutral iterates under a unique optimal policy and for second moments under a weaker common fixed-point condition.
- Sampled first- and second-moment targets retain both successor state–action pairs in each cross term.
- Experiments validate exact distributional recursion, finite-dimensional moment updates, covariance-dependent allocation, and neural estimation from coupled samples.
2 Related work
The related work distinguishes marginal return modeling from joint counterfactual return information. It positions JMDPs as a formal framework for preserving dependence needed for cross-action comparisons and allocation analysis.
- Distributional reinforcement learning estimates marginal return distributions and can derive risk-aware policies from marginal quantiles.
- Prior JMDP work formalized coupled-dynamics environments, defined joint return laws, and proved fixed-policy finite-order moment evaluation.
- Expected-return and risk-aware objectives based on one executed action can be evaluated from marginal return laws.
- Comparing two actions requires their joint return law, while allocation variance requires cross-action covariances.
- In JMDPs, return-vector components index mutually exclusive counterfactual actions rather than objectives jointly produced by one executed action.
3 Preliminaries
The preliminaries define finite discounted MDPs, their return and value quantities, and the distributional Bellman optimality operator. Distributional control convergence relies on stabilization of the greedy policy despite non-contraction.
- A finite discounted MDP consists of finite state and action spaces, a joint reward-transition kernel, and discount factor γ ∈ (0, 1).
- A policy induces returns from state–action pairs, whose expectation is the action-value function Qπ(s, a).
- The optimal value Q* is the fixed point of the expected Bellman optimality operator.
- The distributional Bellman optimality operator selects actions greedily by mean return and is not contractive.
- With a unique optimal policy, distributional control can converge because the greedy policy stabilizes once value estimates become sufficiently accurate.
4 JMDPs and joint moment evaluation
A JMDP samples a full counterfactual outcome table, whose coordinate marginals induce an ordinary MDP while cross-action dependence supports joint return and moment evaluation. The recursion couples same-state branches through shared outcome tables and advances successor branches jointly.
- A finite JMDP augments states, actions, and discounting with a Markov kernel over full one-step counterfactual outcome tables.
- Each action-indexed branch contains a reward and successor state; one branch is executed while the others remain counterfactual observations.
- The environment draws a full outcome table at each state, and coordinate marginals define the induced marginal MDP.
- At each recursion depth, same-state branches query one shared fresh outcome table, whereas branches at distinct states query independent fresh tables.
- The first moment is the expected return, while uncentered second moments track dependence between returns from two state–action pairs.
- Cross-moment continuation terms evaluate successor state–action pairs jointly, with shared outcomes for the same branch and JMDP coupling when current states coincide.
- Cross-action dependence governs counterfactual comparisons and simultaneous allocations beyond the realized return law of single-action execution.
5 Risk-neutral JMDP optimality
The paper extends JMDPs to risk-neutral optimal control by defining joint distributional and moment backups over counterfactual action returns. Convergence follows under a unique mean-optimal action, while second moments also converge under a weaker common-fixed-point condition for optimal ties.
- Joint distributional representation: JMDP distributional families assign a joint return law to each state-action configuration while preserving consistent scalar marginals.The usual action-return vector at a state contains one coordinate per action, and its full law captures cross-action dependence.
- Fixed-policy operator: The fixed-policy Bellman operator is Wasserstein-contracting, with updates coupling shared one-step outcome tables and continuation laws.The contraction is established by using the same one-step tables and optimal continuation couplings for corresponding successor configurations.
- Unique-policy convergence: Under a unique optimal action at every state, distributional optimality iterates converge to the joint return law of the optimal policy.Ordinary value iteration stabilizes the greedy selector once estimates enter the action-gap neighborhood, after which the fixed-policy contraction applies.
- Moment convergence: With p = 2, coordinate means and pairwise second moments converge uniformly to those of the optimal joint return law.The same argument extends to mixed moments of any fixed integer degree n by applying the distributional theorem with p = n.
- Optimal ties: Second-moment convergence permits multiple mean-optimal actions when every optimal selector has the same second-moment fixed point.Different variances, covariances, or pairwise second moments across tie resolutions can violate this condition.
6 Approximate joint moment control
The paper derives sampled targets for learning first and second joint moments from branch tables in large state spaces. The second-moment target preserves continuation dependence across both successor state-action pairs.
- Parametric approximation: Sampled branch tables provide training data for parametric first- and second-moment equations in large state spaces.The approximation uses a mean head µθ and an uncentered second-moment head, with a positive-semidefinite covariance parameterization available.
- Mean target: Mean targets use one sampled reward-successor outcome and a greedy continuation action selected from fixed target mean parameters.The greedy action is chosen by maximizing the target mean head at the sampled successor state.
- Second-moment target: Second-moment targets sample state-action pairs jointly, using one coupled table for same-state pairs and independent tables for distinct states.This sampling scheme matches the JMDP coupling structure across counterfactual branches.
- Cross-moment dependence: The continuation term in the second-moment target depends on both successor states and can evaluate a moment between two different successor state-action pairs.A covariance block local to either successor would discard this cross-successor dependence.
- Optimization: Training minimizes a weighted sampled objective combining first- and second-moment errors, with gradients passing through both mean and covariance parameters.The weighting parameter λ controls the balance between the two moment terms.
7 Experiments
The experiments validate distributional and moment control in exact tabular environments, then examine covariance-sensitive decisions and neural recovery of joint return information from coupled outcome tables.
- Exact distributional control: Exact tabular backups converge to the ground-truth joint law, mean vector, and covariance blocks when the optimal action is unique.With optimal ties, alternating mean-greedy resolutions switch the distributional iterates between two optimal return laws.
- Exact distributional control: The coupled reward chain's shared outcomes shift return coordinates oppositely, while differing successor depths create distinct continuation ranges and offset bands.
- Joint moment control: Moment iteration in the windy gridworld reaches residual 10^-8 and agrees with ordinary value iteration, fixed-policy moment evaluation, and 8000 jointly propagated Monte Carlo rollouts.Shared gusts lower several counterfactual returns together, producing nonnegative action correlations.
- Anti-correlated route choice: The route-choice experiment uses full cross-action covariance to allocate traffic across safe and risky routes, unlike diagonal and single-route baselines.The negative covariance between ridge routes makes equal allocation free of catastrophic loss in this example, though the return remains variable.
- Neural approximation: Neural estimators trained on coupled tables recover joint moments and particle laws better than diagonal or shuffled baselines in the coupled inverted-pendulum task.The particle model improves on its shuffled counterpart and attains lower covariance error than the moment model in this experiment.
8 Conclusion
The paper extends JMDPs from fixed-policy evaluation to control, establishing convergence results and sampled targets that preserve counterfactual dependence. Experiments support recovery of joint quantities from coupled outcome tables.
- The work extends JMDPs from fixed-policy evaluation to control while preserving the induced MDP's mean-optimal policy and counterfactual action-return dependence.
- Distributional iterates converge under a unique mean-optimal policy, while first two moments converge under a weaker common-fixed-point condition for optimal ties.
- Sampled cross-moment targets retain both successor state–action pairs, and tabular and neural experiments support these recursions.