Struct PsroTrainer 

pub struct PsroTrainer<B, P, O, E, FP, FO, FE>
where
    B: AutodiffBackend,
    P: JointPolicy<B>,
    O: Optimizer<P, B>,
    E: JointEnv,
    FP: Fn(&B::Device, u64) -> P,
    FO: Fn() -> BurnOptimizer<B, P, O>,
    FE: Fn() -> E,
{ /* private fields */ }

PSRO outer-loop trainer for symmetric N-agent games (N ≥ 2).

Generic over the Burn backend B, policy module P, and Burn optimizer type O. The trainer owns:

• N populations of policies (one per agent role) under populations: Vec<Vec<P>>.
• A MetaSolver for the empirical meta-game. For N=2 the 2-player MetaSolver::solve path is used (any in-tree solver works); for N≥3 the trainer calls MetaSolver::solve_n_player and only AlphaRankMetaSolver provides a non-panicking override.
• A cached empirical-payoff N-tensor PayoffCache keyed by joint pure strategy.
• User-supplied factories for fresh policies + optimizers + envs.

Policy/optimizer factories

The trainer doesn’t know how to construct a Burn module of the caller’s chosen architecture, so we take closures:

• policy_factory: Fn(&B::Device, u64) -> P — fresh policy. The u64 is a per-construction seed the trainer derives from PsroConfig::seed via a monotonic init-counter. A reproducibility-aware factory threads it into MlpBurnPolicy::new_seeded / MlpBurnConfig::with_seed so that every agent’s initial policy and every per-iteration best-response gets distinct but deterministic weights (issue #135). Factories that don’t care about reproducibility may ignore the argument.
• optimizer_factory: Fn() -> BurnOptimizer<B, P, O> — fresh optimizer.
• env_factory: Fn() -> E — fresh env instance.

This keeps PSRO architecture-agnostic at the cost of slightly awkward generics at the call site (see the matching-pennies test).

Single-policy-class assumption

All agents in both populations share the same policy class P. For 2-agent symmetric games (matching pennies, homogeneous bucket brigade) this is exactly what we want — the symmetry lets us transpose the payoff matrix for the column player’s solve. For fully asymmetric games (different obs/action spaces per role), the trainer needs to be re-parameterized over (P_row, P_col); that’s out of scope for the first PR.

Implementations

impl<B, P, O, E, FP, FO, FE> PsroTrainer<B, P, O, E, FP, FO, FE>

where B: AutodiffBackend, P: JointPolicy<B>, O: Optimizer<P, B>, E: JointEnv, FP: Fn(&B::Device, u64) -> P, FO: Fn() -> BurnOptimizer<B, P, O>, FE: Fn() -> E,

pub fn new( config: PsroConfig, joint_config: JointTrainerConfig, meta_solver: Box<dyn MetaSolver>, device: B::Device, policy_factory: FP, optimizer_factory: FO, env_factory: FE, ) -> Result<Self>

Construct a PSRO trainer with one initial random policy per agent.

joint_config.num_agents must be ≥ 2. For num_agents == 2 the trainer accepts any MetaSolver implementation; for num_agents > 2 the meta-solver’s MetaSolver::solve_n_player is called — at the time of this PR only AlphaRankMetaSolver provides a non-panicking override for N>2.

pub fn populations(&self, agent: usize) -> &[P]

Borrow agent agent’s policy population.

pub fn population_row(&self) -> &[P]

Borrow the row-player (agent 0) population.

Backward-compat shim retained for callers that pre-date the N-tensor refactor (notably tests/test_psro_matching_pennies.rs). New N≥2 code should use PsroTrainer::populations.

pub fn population_col(&self) -> &[P]

Borrow the column-player (agent 1) population.

Backward-compat shim retained for callers that pre-date the N-tensor refactor. Panics for N=1 (which is rejected by new anyway). New N≥2 code should use PsroTrainer::populations.

pub fn payoff_cache(&self) -> &PayoffCache

Borrow the cached empirical N-tensor payoff cache.

pub fn run<F>(&mut self, on_iteration: F) -> Result<PsroStats>

where F: FnMut(&PsroIterationStats, &[&P]), P: Send + Sync, E: Send, FP: Sync, FO: Sync, FE: Sync, B::Device: Sync,

Run the PSRO outer loop and return the per-iteration history.

on_iteration is invoked once per outer iteration, immediately after that iteration’s PsroIterationStats is constructed and before it is pushed onto the returned history. This mirrors NfspTrainer::run and lets callers observe per-iteration progress during the run (live tracing logging, mid-run checkpoint triggers, etc.) rather than only inspecting the aggregate stats after run returns.

The callback receives two arguments:

1. &PsroIterationStats — this iteration’s stats, whose iteration field increases monotonically from 1 to config.max_iterations.
2. &[&P] — the newest best-response policy for each agent (brs[a] is agent a’s freshly-trained BR appended this iteration, i.e. populations(a).last()). This lets the callback persist per-agent BR policies to disk during the run (mid-run checkpointing, issue #204) without a borrow conflict against the &mut self held by run: the trainer cannot itself write files (it is backend/format-agnostic, the Recorder lives in the example), so it hands the closure the policy references it needs to checkpoint. Checkpointing is a pure side-effect read; it does not alter the trainer state or the deterministic training trajectory.

For the common case of “run with no per-iteration hook”, use Self::run_silent.

pub fn run_silent(&mut self) -> Result<PsroStats>

where P: Send + Sync, E: Send, FP: Sync, FO: Sync, FE: Sync, B::Device: Sync,

Convenience entry point: drives Self::run with a no-op iteration callback. Use this when per-iteration observation is not needed (mirrors NfspTrainer::run_silent).

pub fn current_meta_nash_per_agent(&self) -> Vec<Vec<f32>>

Most-recent per-agent meta-Nash distributions (one row per agent), or uniform over the initial population if run has not been called.

pub fn current_meta_nash(&self) -> Vec<f32>

Backward-compat shim returning agent 0’s meta-Nash marginal.