thrust_rl/multi_agent/
mod.rs

1//! Multi-agent training infrastructure for Thrust (Burn backend).
2//!
3//! PR #98 (Burn phase 5) deleted the previous `tch`-coupled `multi_agent`
4//! module wholesale. This module is the Burn-native rebuild called out by
5//! issue #100, sitting on top of the Burn policy networks
6//! ([`crate::policy::mlp::MlpBurnPolicy`],
7//! [`crate::policy::multi_discrete_mlp::MultiDiscreteMlpBurnPolicy`]) and the
8//! backend-agnostic [`crate::train::optimizer::BurnOptimizer`].
9//!
10//! # Scope of the initial Burn port
11//!
12//! This first cut restores the pieces of the multi-agent surface that the
13//! resurrected trainer examples (`train_snake_multi_v2`,
14//! `train_pong_self_play`, the Slepian-Wolf bucket-brigade experiments)
15//! and the multi-discrete integration tests will need:
16//!
17//! - [`crate::multi_agent::environment::MultiAgentEnvironment`] — extension to
18//!   the base [`crate::env::Environment`] trait for per-agent observation /
19//!   action / reward routing. No tensor dependency.
20//! - [`crate::multi_agent::messages`] — Burn-native experience / policy-update
21//!   / control payload types. Pre-Burn these carried `tch::Tensor`; post-Burn
22//!   they carry plain `Vec<f32>` host buffers so the producer and consumer can
23//!   pick different Burn backends.
24//! - [`crate::multi_agent::joint::JointMultiAgentTrainer`] — synchronized joint
25//!   trainer used when a loss term depends on **all** agents' parameters
26//!   evaluated on the **same minibatch** (the Slepian-Wolf cross-agent
27//!   redundancy penalty is the canonical motivation). See the module doc for
28//!   the Burn-specific single-graph-multiple-optimizer pattern.
29//!
30//! # What's intentionally NOT here yet
31//!
32//! The pre-Burn module also shipped:
33//!
34//! - `learner.rs` / `population.rs` / `simulator.rs` / `matchmaking.rs` —
35//!   per-thread independent learners with shared experience channels. The
36//!   threading layer is orthogonal to the Burn migration but its old
37//!   implementation was tightly coupled to `tch`'s shared `VarStore` model; a
38//!   Burn-native port should design around Burn's move-through optimizer
39//!   semantics and is left to a follow-up.
40//! - `centralized_critic.rs` — optional centralized critic for MAPPO-style
41//!   training. The Burn analog already lives at [`crate::train::optimizer`] one
42//!   level up; threading it through the joint trainer is a small follow-up
43//!   extension.
44//!
45//! Issue #100's definition of done only requires the joint trainer plus
46//! its synthetic-env smoke test; those follow-ups are explicitly out of
47//! scope for the first port.
48
49/// Bucket-brigade reference baseline policies (specialist, cell enum).
50#[cfg(feature = "env-bucket-brigade")]
51pub mod bucket_brigade_baselines;
52/// Bucket-brigade-specific scoring metrics (`gap_closed`).
53#[cfg(feature = "env-bucket-brigade")]
54pub mod bucket_brigade_metrics;
55/// Bucket-brigade best-response improvability oracle (issue #259).
56#[cfg(feature = "env-bucket-brigade")]
57pub mod bucket_brigade_oracle;
58/// Fixed-vocab agent-to-agent communication (comms) surface (issue #274).
59pub mod comms;
60/// Multi-agent environment trait.
61pub mod environment;
62/// Synchronized joint multi-agent PPO trainer.
63pub mod joint;
64/// Cross-thread message payloads.
65pub mod messages;
66/// Neural Fictitious Self-Play (NFSP) trainer (issue #106).
67pub mod nfsp;
68/// Policy-Space Response Oracles (PSRO) meta-game trainer (issue #107).
69pub mod psro;
70
71pub use comms::{AgentMessage, CommunicatingEnvironment, Delivery};
72pub use environment::{MultiAgentEnvironment, MultiAgentResult};
73pub use joint::{
74    JointEnv, JointMultiAgentTrainer, JointPolicy, JointRollout, JointStats, JointStepResult,
75    JointTrainerConfig,
76};
77pub use messages::{
78    AgentId, ControlMessage, Experience, PolicyBroadcast, PolicyUpdate, TrainingStats,
79};
80pub use nfsp::{NfspConfig, NfspIterationStats, NfspStats, NfspTrainer, ReservoirBuffer};
81pub use psro::{
82    AlphaRankMetaSolver, FictitiousPlayMetaSolver, MetaSolver, PayoffCache, PsroConfig,
83    PsroIterationStats, PsroStats, PsroTrainer, ReplicatorDynamicsMetaSolver, UniformMetaSolver,
84};
thrust_rl/multi_agent/mod.rs

thrust_rl/multi_agent/
mod.rs
