Struct JointTrainerConfig 

pub struct JointTrainerConfig {

    pub num_agents: usize,
    pub rollout_steps: usize,
    pub gamma: f64,
    pub gae_lambda: f64,
    pub clip_range: f64,
    pub clip_range_vf: f64,
    pub vf_coef: f64,
    pub ent_coef: f64,
    pub n_epochs: usize,
    pub minibatch_size: usize,
    pub max_grad_norm: f64,
    pub normalize_advantages: bool,
    pub iterate_all_minibatches: bool,
    pub critic_lr: Option<f64>,
    pub max_minibatches_per_epoch: Option<usize>,
    pub comms_coef: f64,
}

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

Fields

num_agents: usize

Number of agents trained jointly. Must match the length of JointMultiAgentTrainer’s policy / optimizer slots.

rollout_steps: usize

Steps collected per rollout before each PPO update.

gamma: f64

Discount factor γ ∈ [0, 1].

gae_lambda: f64

GAE smoothing parameter λ ∈ [0, 1].

clip_range: f64

PPO policy-ratio clip range ε.

clip_range_vf: f64

PPO value-function clip range. Use f64::INFINITY to fall back to plain MSE (matches the Burn compute_value_loss contract).

vf_coef: f64

Weight on the value-function loss term inside the joint loss.

ent_coef: f64

Weight on the entropy bonus.

n_epochs: usize

Number of PPO epochs per update.

minibatch_size: usize

Minibatch size for SGD within each PPO epoch.

max_grad_norm: f64

Global gradient-norm clip applied through each per-policy optimizer.

normalize_advantages: bool

Standardize advantages to zero mean / unit variance per minibatch before computing the surrogate.

iterate_all_minibatches: bool

Iterate over all minibatches per epoch instead of a single truncated minibatch_size draw.

The historical Burn-native cut took one shuffled minibatch per epoch (issue #100 simplification), which discards ~97% of a large rollout and starves the best-response (issue #239). When true, each epoch shuffles the rollout once and walks every minibatch_size chunk, the conventional PPO pattern.

Defaults to false to keep existing NFSP/PSRO determinism tests bit-identical; callers that want the full-rollout update (e.g. the bucket-brigade examples) opt in explicitly.

critic_lr: Option<f64>

Optional separate critic learning rate (issue #239, ranked fix #4).

The shared actor-critic trunk is trained with one optimizer and one combined policy + vf_coef·value − ent_coef·entropy backward. On the bucket-brigade cells the critic never fits (explained-variance pinned at ~0 — see #239), so the normalized advantages are noise and the policy gradient stays ~0 (entropy stuck at the uniform-max floor). PR #240’s grad-clip / all-minibatch / vf_coef knobs did not move ev off 0, and dropping BR_REWARD_SCALE to 0.001 collapsed value loss to ~8 yet ev still stayed flat — i.e. the critic is not fitting even tiny well-scaled targets at the shared 3e-4 LR.

When Some(lr), the joint update splits the combined backward into two passes per minibatch:

1. actor pass — policy − ent_coef·entropy stepped through the usual per-agent optimizer at its construction LR;
2. critic pass — value_loss alone stepped through a dedicated per-agent critic optimizer at critic_lr.

Both passes update the full module (shared trunk + their respective heads), but the critic gets its own Adam moment state and (typically higher) LR, so it can fit the value target without the policy LR dragging it. The critic optimizers are supplied via JointMultiAgentTrainer::with_critic_optimizers; if critic_lr is Some but no critic optimizers were supplied the trainer falls back to the single combined backward (logged once at construction).

Defaults to None (single combined backward — the historical behaviour), so existing NFSP/PSRO runs and determinism tests are bit-identical unless a caller opts in.

max_minibatches_per_epoch: Option<usize>

Optional cap on the number of minibatch gradient steps per epoch (issue #251, throughput lever).

The #239 best-response fix sets iterate_all_minibatches = true, which walks every minibatch_size chunk of the rollout each epoch. Combined with br_train_steps_per_iteration = 8 over the un-batchable (issue #235) bucket-brigade rollout, the per-iteration BR update became the dominant outer-loop cost (>1 h/iter at 2048 rollout — see issue #251). At rollout_steps = 2048, minibatch_size = 256 that is 8 minibatches × n_epochs × br_train_steps gradient steps per outer iteration.

When Some(cap), after each epoch’s minibatch index-sets are built (and globally shuffled), the set is truncated to at most cap minibatches — a uniformly-random subsample of the rollout, since the indices were shuffled before chunking. This trades a bounded amount of BR fit per epoch for throughput without reverting the #239 learning behaviour: grad-clip, vf_coef, iterate_all_minibatches, and br_train_steps_per_iteration are all unchanged, and each capped minibatch is still a full forward+backward over minibatch_size samples. cap is clamped to at least 1 so a capped update never degenerates to zero gradient steps.

None (default) preserves the full all-minibatch coverage exactly, so existing NFSP/PSRO runs and determinism tests are bit-identical unless a caller opts in. With the default single-minibatch path (iterate_all_minibatches == false) the per-epoch set already has one entry, so any cap >= 1 is a no-op there.

comms_coef: f64

Weight on the optional comms-regularization term (issue #275, Phase 2 of docs/COMMS_DESIGN.md).

When nonzero, each minibatch update folds a message-entropy penalty — the Shannon entropy (nats) of every agent’s sampled action-token distribution over the minibatch, summed across agents and scaled by comms_coef — into JointStats::aux_loss and the joint backward. This reuses the existing aux-loss plumbing (the design’s “reuse the aux-loss slot, don’t add a new one”) rather than changing the aux_fn signature that PSRO/NFSP depend on.

The term is computed from the discrete sampled tokens, so under the Phase 2 (non-differentiable) comms surface it contributes no gradient and acts as a monitored regularizer; differentiable message gradients are Phase 3 (issue #276, Gumbel-softmax).

Defaults to 0.0 (disabled — no comms term added), so existing NFSP/PSRO runs and determinism tests are bit-identical unless a caller opts in.

Trait Implementations

impl Clone for JointTrainerConfig

fn clone(&self) -> JointTrainerConfig

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for JointTrainerConfig

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Default for JointTrainerConfig

fn default() -> Self

Returns the “default value” for a type.