Module joint 

Synchronized joint multi-agent PPO trainer. Synchronized joint multi-agent PPO trainer (Burn backend).

Burn-native rebuild of the pre-Burn JointMultiAgentTrainer. The previous tch-coupled implementation was deleted in PR #98 along with the rest of src/multi_agent/; this module re-establishes the synchronized (one shared backward pass) joint trainer on top of the Burn policy networks and the crate::train::optimizer::BurnOptimizer wrapper.

When to use this module

Use the joint trainer when you need a loss term that depends on all agents’ parameters evaluated on the same minibatch at the same optimization step. The canonical motivating example is the Slepian-Wolf MARL P3 cross-agent representational redundancy penalty

L_red   = λ * Σ_{i<j} || corr(Z_i(obs), Z_j(obs)) ||_F² / d²
L_total = Σ_i L_ppo[i](obs, a_i, ...) + L_red

where Z_i = encoder_features_i(obs). The penalty couples every policy’s encoder through one shared backward pass; per-thread learners cannot express this without heavy synchronization.

Single-graph, multiple-optimizer semantics under Burn

Burn’s burn::optim::Optimizer::step consumes the module by value and returns the updated copy. That makes the tch-style “one .backward() plus N independent Optimizer::step calls touching disjoint var-stores” pattern slightly different in Burn:

1. Each policy computes its own (policy_loss, value_loss, entropy) on the minibatch.
2. The caller-supplied aux_fn is invoked on every policy’s encoder features and may return an additional scalar loss.
3. All per-agent losses plus the aux loss are summed into one scalar joint_loss; we call .backward() once.
4. For each policy i, we slice the joint gradients down to policy i’s parameters with [burn::optim::GradientsParams::from_grads(grads.clone(), &policies[i])], then call optimizer_i.step(lr, policies[i], slice_i).

That last step is the Burn analog of “every optimizer reads only its own var-store’s .grad” on the tch path — gradient flow stays parameter-isolated because each from_grads slice only contains the ids of one policy’s params. The aux term’s contribution flows into every policy’s slice through the shared backward pass, by construction.

Minibatch sampling

By default (and to keep the smoke test deterministic) this trainer takes one minibatch per epoch, sampled via crate::train::ppo::loss::generate_minibatch_indices and truncated to config.minibatch_size. Setting JointTrainerConfig::iterate_all_minibatches switches to the conventional PPO pattern of walking every minibatch_size chunk per epoch, so a large rollout is fully consumed instead of ~97% discarded (issue #239). The default stays single-minibatch so existing NFSP/PSRO determinism tests are bit-identical.

Gradient clipping

JointTrainerConfig::max_grad_norm is applied as a global L2-norm clip on each policy’s gradient slice before the optimizer step (issue #239): if the joint gradient norm exceeds the cap, every gradient is scaled down by max_norm / ‖g‖, preserving direction. This stops the raw-scale value-loss gradient from swamping the policy heads on the shared actor-critic trunk.

Structs

JointMultiAgentTrainer

Synchronized joint multi-agent PPO trainer (Burn backend).

JointRollout

Synchronized rollout buffer (host-side).

JointStats

Per-update training statistics for the joint trainer.

JointStepResult

Per-step result returned by a JointEnv implementation.

JointTrainerConfig

Trainer configuration. Plain data; defaults match the tch-era JointTrainerConfig field-for-field so the smoke-test parameters stay portable.

Traits

JointEnv

Minimal joint-environment surface needed by JointMultiAgentTrainer::collect_rollout.

JointPolicy

Capabilities a policy must expose to participate in JointMultiAgentTrainer.