Struct JointMultiAgentTrainer 

pub struct JointMultiAgentTrainer<B, P, O>
where
    B: AutodiffBackend,
    P: JointPolicy<B>,
    O: Optimizer<P, B>,
{ /* private fields */ }

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

Generic over:

• B: AutodiffBackend — the Burn backend.
• P: JointPolicy<B> — the per-agent policy module type.
• O: Optimizer<P, B> — the Burn optimizer type (typically built from AdamConfig::new().init()).

The trainer owns N policies and N optimizers; gradient flow is parameter-isolated because each burn::optim::GradientsParams::from_grads slice extracts only one policy’s parameters from the shared autograd gradients.

Implementations

impl<B, P, O> JointMultiAgentTrainer<B, P, O>

where B: AutodiffBackend, P: JointPolicy<B>, O: Optimizer<P, B>,

pub fn new( policies: Vec<P>, optimizers: Vec<BurnOptimizer<B, P, O>>, config: JointTrainerConfig, device: B::Device, ) -> Result<Self>

Construct a trainer from a fully-initialized set of policies and optimizers.

optimizers[i] is paired with policies[i] and only ever updates policies[i]’s parameters.

pub fn with_critic_optimizers( policies: Vec<P>, optimizers: Vec<BurnOptimizer<B, P, O>>, critic_optimizers: Vec<BurnOptimizer<B, P, O>>, config: JointTrainerConfig, device: B::Device, ) -> Result<Self>

Construct a trainer with a dedicated per-agent critic optimizer (issue #239, ranked fix #4).

critic_optimizers[i] is paired with policies[i] and steps only the value-loss gradient (a second backward over value_loss alone) at its own learning rate. This decouples the critic’s effective LR from the actor’s so the critic can fit the bucket-brigade value target (whose explained-variance was pinned at ~0 under the single shared optimizer — see JointTrainerConfig::critic_lr).

Requires JointTrainerConfig::critic_lr to be Some; the value is the LR the supplied critic optimizers were constructed with (recorded for diagnostics — the actual step uses each critic optimizer’s own construction LR). The same max_grad_norm cap is staged on the critic optimizers too.

pub fn device(&self) -> &B::Device

Device the trainer (and all its policies) live on.

pub fn config(&self) -> &JointTrainerConfig

Trainer configuration.

pub fn policy(&self, i: usize) -> &P

Borrow agent i’s policy. Panics if the trainer is mid-update.

pub fn collect_rollout<E: JointEnv>( &self, env: &mut E, last_obs: &mut [Vec<f32>], rng: &mut StdRng, ) -> JointRollout

Drive a JointEnv for config.rollout_steps and return the synchronized rollout buffer.

last_obs is the persistent “next observation per agent” handed in across iterations: pass the per-agent observations from the most recent env.reset_joint() or step (i.e. the env’s full Vec<Vec<f32>> shape, length = num_agents). The trainer updates it in place so callers can keep the rollout stream stitched across iterations.

rng is consumed by each per-step JointPolicy::get_action_host_seeded call. Pass the trainer-owned StdRng (e.g. PsroTrainer::self.rng) for PsroConfig::seed / NfspConfig::seed to produce bit-identical rollouts (issue #114).

pub fn update<F>( &mut self, rollout: &JointRollout, rng: &mut StdRng, aux_fn: F, ) -> Result<JointStats>

where F: FnMut(&[Tensor<B, 2>]) -> Option<Tensor<B, 1>>,

Joint PPO update.

aux_fn receives a slice of per-agent encoder-feature tensors for the current minibatch (one entry per policy, shape [mb, hidden_dim]) and returns an optional pre-scaled scalar loss (e.g. the cross-agent redundancy penalty). One .backward() flows through every encoder when aux_fn returns Some.

Minibatch sampling

One shuffled minibatch of size config.minibatch_size is drawn per epoch, truncated against the rollout length. The order of indices within the minibatch is irrelevant because every loss is a mean / sum reduction over the minibatch dim and therefore permutation-invariant.

pub fn update_with_active_agents<F>( &mut self, rollout: &JointRollout, active: &[bool], rng: &mut StdRng, aux_fn: F, ) -> Result<JointStats>

where F: FnMut(&[Tensor<B, 2>]) -> Option<Tensor<B, 1>>,

Joint PPO update with per-agent active mask — the freeze-N-1 primitive used by PSRO’s best-response step.

Identical to Self::update except that frozen agents (active[i] == false) skip the optimizer step. Their loss is still summed into the joint backward so the shared autograd graph remains balanced, but their parameters are guaranteed unchanged: we put the original policy back in its slot without calling optimizer.step. Per-agent stats for frozen agents are still recorded in the returned JointStats so callers can monitor the mixture’s behaviour on the rollout.

Use case

PSRO’s outer loop trains one best-response policy at a time against a meta-Nash mixture over the rest of the population (see crate::multi_agent::psro). Passing active = [false, ..., true (active idx), ..., false] here is the canonical freeze-N-1 pattern.

Panics

Returns Err if active.len() != config.num_agents.