# Joint-Particle Distributional Critic

- ID: 74
- Canonical URL: https://synthcore.org/idea/74/joint-particle-distributional-critic
- API JSON: https://synthcore.org/api/idea/74.json
- API Markdown: https://synthcore.org/api/idea/74.md
- Verification status: audited
- Source: [arXiv:2608.22765](https://arxiv.org/abs/2608.22765)
- Category: architecture
- Solves: accuracy, stability, sample-efficiency
- ML areas: rl, loss, world-model, training-dynamics
- Math tags: probability, optimal-transport, dynamical-systems, statistics
- Ratings: usefulness 7/10; difficulty 6/10; novelty 7/10

## Idea description

Replace independent per-action distributional value heads with a critic whose shared latent particle produces a vector of return samples for all actions simultaneously. Train the predicted joint return vector against a Bellman target vector formed from coupled counterfactual reward-transition samples, using a sliced Wasserstein loss. The greedy action is selected by the mean of the corresponding marginal particles, while shared particles retain cross-action dependence for learning and uncertainty estimation.

## Mathematical statement

Definition 1 defines the discounted return as $Z^{\pi}(s,a):=\sum_{t=0}^{\infty}\gamma^{t}R_t$, where $s\in\mathcal{S}$ is a state, $a\in\mathcal{A}$ is an action, $R_t$ is the reward at time $t$, and $\gamma\in(0,1)$ is the discount factor. In a JMDP, the one-step variables $(R_a,S'_a)_{a\in\mathcal{A}}$ are sampled jointly from one coupled kernel, so different counterfactual actions share exogenous randomness. The adapted joint Bellman target for particle $k$ is $Y_k=(Y_{a,k})_{a\in\mathcal{A}}$ with $Y_{a,k}=R_{a,k}+\gamma f_{\bar\theta}(S'_{a,k},\pi_{\bar\theta}(S'_{a,k}),\xi'_k)_a$, where $\bar\theta$ are target-network parameters, $\xi'_k$ is a shared latent particle, and $\pi_{\bar\theta}(s)=\arg\max_b K^{-1}\sum_{j=1}^{K}f_{\bar\theta}(s,\xi_j)_b$ is the mean-greedy policy. The critic predicts $X_k=f_\theta(s,\xi_k)\in\mathbb{R}^{A}$, where $A=|\mathcal{A}|$ and $X_{a,k}$ is a sample from the return distribution for action $a$. Use the sliced Wasserstein discrepancy $\operatorname{SW}_1(X,Y)=\mathbb{E}_{u\sim\operatorname{Unif}(\mathbb{S}^{A-1})}[W_1(\{u^\top X_k\}_{k=1}^{K},\{u^\top Y_k\}_{k=1}^{K})]$, approximated with random directions and sorted projected samples. The paper's convergence result states that iterates of the distributional Bellman optimality operator converge in Wasserstein distance to the optimal joint return law when the induced marginal MDP has a unique optimal policy; first and second moments may converge under weaker tie conditions. The implementation uses the shared coupling and Wasserstein geometry to preserve cross-action return dependence.

## Key formulas

- $$Z^{\pi}(s,a):=\sum_{t=0}^{\infty}\gamma^{t}R_t$$
- $$Y_{a,k}=R_{a,k}+\gamma f_{\bar\theta}\!\left(S'_{a,k},\pi_{\bar\theta}(S'_{a,k}),\xi'_k\right)_a,\qquad \pi_{\bar\theta}(s)=\arg\max_{b\in\mathcal A}\frac{1}{K}\sum_{j=1}^{K}f_{\bar\theta}(s,\xi_j)_b$$
- $$\operatorname{SW}_1(X,Y)=\mathbb{E}_{u\sim\operatorname{Unif}(\mathbb S^{A-1})}\left[W_1\!\left(\{u^\top X_k\}_{k=1}^{K},\{u^\top Y_k\}_{k=1}^{K}\right)\right]$$
- $$\mathcal L(\theta)=\frac{1}{M}\sum_{m=1}^{M}\frac{1}{K}\sum_{k=1}^{K}\left|\operatorname{sort}_{j}\big(u_m^\top X_j\big)_k-\operatorname{sort}_{j}\big(u_m^\top Y_j\big)_k\right|$$

## Implementation notes

(1) Integration point: implement this as a distributional RL critic replacing a scalar Q-network or independent quantile heads. Given a state embedding $h(s)$ and a latent particle $\xi_k\sim\mathcal N(0,I)$, an MLP outputs an $A$-dimensional vector $X_k=f_\theta(s,\xi_k)$, one return sample for every action. Use the same particle index across all action coordinates. The replay buffer must store coupled transitions $(R_a,S'_a)_{a=1}^{A}$ generated by one shared simulator random seed; ordinary single-action environments cannot test the joint advantage directly. (2) Pseudocode: sample a minibatch of coupled transitions and latent particles; compute $X_k=f_\theta(s,\xi_k)$; compute $\mu_b=K^{-1}\sum_k f_{\bar\theta}(s'_a,\xi_k)_b$ and choose $b^*(s'_a)=\arg\max_b\mu_b$; construct $Y_{a,k}=R_a+\gamma f_{\bar\theta}(S'_a,\xi'_k)_{b^*(S'_a)}$ using shared $\xi'_k$; draw $M$ random unit vectors $u_m$; project both $X_k$ and $Y_k$ onto every direction, sort the $K$ projections, and minimize the mean absolute sorted difference in the sliced-Wasserstein loss. Update $\theta$ with Adam and periodically copy it to $\bar\theta$. (3) Mathematical versus empirical quantities: the return recursion, shared-vector target, and Wasserstein objective are mathematical components; $K$, $M$, target-update period, and $\gamma$ are hyperparameters. Estimate action means by particle averages. Log empirical cross-action covariance $\widehat{\operatorname{Cov}}(X_{\cdot,a},X_{\cdot,b})$ and compare it with covariance from held-out coupled rollouts. (4) First cheap experiment: create a stochastic two-action bandit or gridworld in which both counterfactual rewards depend on a controllable shared Bernoulli or Gaussian exogenous variable. Compare scalar DQN, independent quantile regression, and the joint critic at equal network size and environment interactions. Then test a small vectorized control task with common-randomness simulator seeds. Measure held-out joint-return Wasserstein error, mean return, action-selection accuracy when means are close, and cross-action covariance calibration. Success means lower joint-law error or faster return improvement at equal interaction count; if shared-particle training gives no improvement over independently resampled transitions, reject the added architecture.

## Verification

- Status: audited
- Mechanism evidence: yes
- Mechanism confirmed: no
- Verdict: Built a shared-latent joint-particle critic with sliced-Wasserstein training and an independent marginal-head baseline on a coupled two-action Gaussian terminal task. The math check passed: identical particle sets had SW1=0, and translation error was 2.98e-8. The joint critic achieved lower held-out joint SW1 (0.1413 vs 0.2709) and reproduced cross-action covariance (0.750 vs true 0.706), while the independent baseline produced incorrect covariance (-0.092), demonstrating the intended dependence-preservation effect.

### Mechanism check

- Verdict: Built a shared-latent joint-particle critic with sliced-Wasserstein training and an independent marginal-head baseline on a coupled two-action Gaussian terminal task. The math check passed: identical particle sets had SW1=0, and translation error was 2.98e-8. The joint critic achieved lower held-out joint SW1 (0.1413 vs 0.2709) and reproduced cross-action covariance (0.750 vs true 0.706), while the independent baseline produced incorrect covariance (-0.092), demonstrating the intended dependence-preservation effect.
- Confidence: 8/10
- Limitations: This is a terminal coupled Gaussian bandit, so gamma=0 and no multi-step Bellman backup or learned mean-greedy control policy was evaluated. Results use one fixed seed and finite particles/directions; the independent baseline is designed to match marginals but does not receive coupled-pair supervision.

## Artifacts

- [joint_particle_experiment.py](https://synthcore.org/code/64/joint_particle_experiment.py)
- [make_report.py](https://synthcore.org/code/64/make_report.py)
- [report.md](https://synthcore.org/code/64/report.md)
- [report_super_2026-08-30T104107.md](https://synthcore.org/code/64/report_super_2026-08-30T104107.md)
- [results.json](https://synthcore.org/code/64/results.json)
- [super_results.json](https://synthcore.org/code/64/super_results.json)
- [supercheck.py](https://synthcore.org/code/64/supercheck.py)
- [supercheck_report.md](https://synthcore.org/code/64/supercheck_report.md)
- [Download all files as ZIP](https://synthcore.org/download/64)

## Disclaimer

AI-generated research hypothesis, automatically tested. Not peer-reviewed.