Struct MlpBurnPolicy 

pub struct MlpBurnPolicy<B: Backend> { /* private fields */ }

Two- or three-layer MLP actor-critic for discrete action spaces, ported to Burn.

Layout mirrors crate::policy::mlp::MlpBurnPolicy at a high level:

obs → fc1 →act→ fc2 →act→ (fc3 →act→)? policy_head (logits)
                                      └─ value_head  (V(s))

Both heads share the trunk activations — standard PPO actor-critic.

Numerical parity

When constructed with use_orthogonal_init = true (the default), the trunk uses [Initializer::Orthogonal { gain: sqrt(2) }] and the output heads use gain = 0.01. These match the tch policy’s init gains exactly (see crate::policy::mlp::MlpBurnPolicy::with_config), which is the necessary precondition for the phase-4 numerical-parity check called out on issue #81.

Implementations

impl<B: Backend> MlpBurnPolicy<B>

pub fn new( obs_dim: usize, action_dim: usize, hidden_dim: usize, device: &B::Device, ) -> Self

Backward-compatible 2-layer constructor (the phase 1 scout signature). Uses Burn’s default Kaiming-uniform init — kept so the existing bandit trainer and parity tests are not perturbed.

New call sites that want PPO-style orthogonal init should call MlpBurnPolicy::with_config instead.

pub fn new_seeded( obs_dim: usize, action_dim: usize, hidden_dim: usize, seed: u64, device: &B::Device, ) -> Self

Seeded variant of new: same 2-layer Kaiming architecture, but constructed deterministically from seed so two calls with the same seed produce bit-identical weights.

Convenience wrapper for callers (and the PSRO/NFSP policy factories) that want reproducible policies without assembling a full MlpBurnConfig. Equivalent to with_config(.., MlpBurnConfig { use_orthogonal_init: false, .., seed: Some(seed) }, ..).

pub fn with_config( obs_dim: usize, action_dim: usize, config: MlpBurnConfig, device: &B::Device, ) -> Self

Build a fresh policy on device with the given configuration.

This is the production constructor for phase 4 onwards. Mirrors crate::policy::mlp::MlpBurnPolicy::with_config.

pub fn forward(&self, obs: Tensor<B, 2>) -> (Tensor<B, 2>, Tensor<B, 1>)

Forward pass: returns (logits, value).

• obs is shape [batch, obs_dim].
• logits is shape [batch, action_dim] (pre-softmax).
• value is shape [batch] (squeezed from [batch, 1]).

pub fn encoder_features(&self, obs: Tensor<B, 2>) -> Tensor<B, 2>

Compute the shared-trunk feature representation for obs.

Mirrors crate::policy::mlp::MlpBurnPolicy::encoder_features — auxiliary regularizers (cross-agent redundancy penalties, behavioural-diversity bonuses) tap this directly.

Gradients flow back into the trunk.

pub fn policy_head_action_dim(&self) -> usize

Action-head output dimensionality (number of discrete actions).

Reads the policy_head weight tensor’s shape — Burn’s burn::nn::Linear stores weight: Param<Tensor<B, 2>> with shape [d_input, d_output], so d_output is the action cardinality. Used by the multi-agent joint trainer’s crate::multi_agent::joint::JointPolicy::action_dims_joint impl to size the rollout action buffer without consuming RNG draws.

pub fn fc1(&self) -> &Linear<B>

Borrow the first shared-trunk linear layer.

pub fn fc2(&self) -> &Linear<B>

Borrow the second shared-trunk linear layer.

pub fn policy_head(&self) -> &Linear<B>

Borrow the policy (action-logits) head.

pub fn value_head(&self) -> &Linear<B>

Borrow the value (V(s)) head.

pub fn get_action_host( &self, obs: Tensor<B, 2>, ) -> (Vec<i64>, Vec<f32>, Vec<f32>)

Sample one action per row from the policy’s categorical distribution and return (actions_host, log_probs_host, values_host) as plain Vecs.

Thin backwards-compat wrapper around MlpBurnPolicy::get_action_host_seeded that constructs a thread-local RNG. Not deterministic across calls — use get_action_host_seeded and pass a seeded rand::rngs::StdRng when reproducibility is required (PSRO/NFSP/joint trainer rollouts call the seeded form via the crate::multi_agent::joint::JointPolicy trait so that PsroConfig::seed / NfspConfig::seed produce bit-identical rollouts; see issue #114).

Retained for example-driver convenience where the caller does not need bit-exact reproducibility and would otherwise have to thread an &mut StdRng through bespoke rollout loops.

pub fn get_action_host_seeded( &self, obs: Tensor<B, 2>, rng: &mut StdRng, ) -> (Vec<i64>, Vec<f32>, Vec<f32>)

Same contract as get_action_host but the host-side categorical draws consume rng instead of the thread-local generator.

The trainer-side rollout loop does not need gradient flow through the sampled action (only the eventual MlpBurnPolicy::evaluate_actions call on the stored transitions matters for the PPO surrogate). We therefore do the categorical draw on the host with rand, sidestepping Burn 0.21’s lack of a first-class multinomial op.

Bit-exactness contract: two calls with the same obs, same policy state, and same-seeded rng (StdRng::seed_from_u64) must produce element-wise identical (actions, log_probs, values). This is the load-bearing guarantee PsroConfig::seed / NfspConfig::seed rely on after issue #114.

pub fn get_actions_host_seeded_batched( &self, obs: Tensor<B, 2>, rng: &mut StdRng, ) -> (Vec<i64>, Vec<f32>, Vec<f32>)

Batched seeded sampler: one forward over [N, obs_dim], then N host-side categorical draws in row-major order.

Bit-identical to calling Self::get_action_host_seeded once per row on [1, obs_dim] slices provided the rows are drawn from the same policy — the forward is a single matmul over all N rows (same weights), and the RNG is consumed one draw per row, ascending. This is the crate::multi_agent::joint::JointPolicy-trait batched entry point that eliminates per-call batch-1 overhead on the NdArray backend wherever many observations are scored through the same model in one step (issue #235).

pub fn evaluate_actions( &self, obs: Tensor<B, 2>, actions: Tensor<B, 1, Int>, ) -> (Tensor<B, 1>, Tensor<B, 1>, Tensor<B, 1>)

Evaluate a batch of (obs, actions) pairs.

Returns (action_log_probs, entropy_per_row, values) — the quantities the PPO surrogate loss needs. Entropy is per-row here (not the mean): the caller decides how to aggregate. This matches the tch policy’s contract (the tch evaluate_actions returns a scalar mean; the trainer reduces per-row entropy on the Burn path inside crate::train::ppo::trainer::PPOTrainerBurn::train_step).

Trait Implementations

impl<B> AutodiffModule<B> for MlpBurnPolicy<B>

where B: AutodiffBackend + Backend, <B as AutodiffBackend>::InnerBackend: Backend,

type InnerModule = MlpBurnPolicy<<B as AutodiffBackend>::InnerBackend>

Inner module without auto-differentiation.

fn valid(&self) -> Self::InnerModule

Returns the same module, but on the inner backend without auto-differentiation.

fn from_inner(module: Self::InnerModule) -> Self

Wraps an inner module back into an auto-diff module.

impl<B: Backend> Clone for MlpBurnPolicy<B>

fn clone(&self) -> Self

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl<B: Debug + Backend> Debug for MlpBurnPolicy<B>

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

Formats the value using the given formatter.

impl<B: Backend> Display for MlpBurnPolicy<B>

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

Formats the value using the given formatter.

impl<B> HasAutodiffModule<B> for MlpBurnPolicy<B::InnerBackend>

where B: AutodiffBackend + Backend, <B as AutodiffBackend>::InnerBackend: Backend,

type TrainModule = MlpBurnPolicy<B>

The module with auto-differentiation.

impl<B: AutodiffBackend> JointPolicy<B> for MlpBurnPolicy<B>

where Self: AutodiffModule<B> + Clone,

fn get_action_host_seeded( &self, obs: Tensor<B, 2>, rng: &mut StdRng, ) -> (Vec<i64>, Vec<f32>, Vec<f32>)

Sample actions for a single rollout step using a caller-supplied seeded RNG.

fn get_actions_host_seeded_batched( &self, obs: Tensor<B, 2>, rng: &mut StdRng, ) -> (Vec<i64>, Vec<f32>, Vec<f32>)

Batched seeded sampler: obs carries N rows scored through this one policy in a single forward, returning N actions.

fn evaluate_actions_joint( &self, obs: Tensor<B, 2>, actions: Tensor<B, 2, Int>, ) -> (Tensor<B, 1>, Tensor<B, 1>, Tensor<B, 1>)

Re-evaluate the policy on previously-sampled actions.

fn encoder_features_joint(&self, obs: Tensor<B, 2>) -> Tensor<B, 2>

Shared-trunk feature representation; gradients flow back into the encoder. The natural quantity to feed into cross-agent regularizers.

fn action_dims_joint(&self) -> Vec<i64>

Per-dimension action cardinalities.

impl<B: Backend> Module<B> for MlpBurnPolicy<B>

type Record = MlpBurnPolicyRecord<B>

Type to save and load the module.

fn load_record(self, record: Self::Record) -> Self

Load the module state from a record.

fn into_record(self) -> Self::Record

Convert the module into a record containing the state.

fn num_params(&self) -> usize

Get the number of parameters the module has, including all of its sub-modules.

fn visit<Visitor: ModuleVisitor<B>>(&self, visitor: &mut Visitor)

Visit each tensor parameter in the module with a visitor.

fn map<Mapper: ModuleMapper<B>>(self, mapper: &mut Mapper) -> Self

Map each tensor parameter in the module with a mapper.

fn collect_devices(&self, devices: Devices<B>) -> Devices<B>

Return all the devices found in the underneath module tree added to the given vector without duplicates.

fn to_device(self, device: &B::Device) -> Self

Move the module and all of its sub-modules to the given device.

fn fork(self, device: &B::Device) -> Self

Fork the module and all of its sub-modules to the given device.

fn devices(&self) -> Vec<<B as BackendTypes>::Device>

Return all the devices found in the underneath module tree without duplicates.

fn no_grad(self) -> Self

Each tensor in the module tree will not require grad.

fn train<AB>(self) -> Self::TrainModule

where AB: AutodiffBackend<InnerBackend = B>, Self: HasAutodiffModule<AB>,

Move the module and all of its sub-modules to the autodiff backend.

fn save_file<FR, PB>( self, file_path: PB, recorder: &FR, ) -> Result<(), RecorderError>

where FR: FileRecorder<B>, PB: Into<PathBuf>,

Save the module to a file using the provided file recorder.

fn load_file<FR, PB>( self, file_path: PB, recorder: &FR, device: &<B as BackendTypes>::Device, ) -> Result<Self, RecorderError>

where FR: FileRecorder<B>, PB: Into<PathBuf>,

Load the module from a file using the provided file recorder.

fn quantize_weights(self, quantizer: &mut Quantizer) -> Self

Quantize the weights of the module.

impl<B: Backend> ModuleDisplay for MlpBurnPolicy<B>

fn format(&self, passed_settings: DisplaySettings) -> String

Formats the module with provided display settings.

fn custom_settings(&self) -> Option<DisplaySettings>

Custom display settings for the module.

fn custom_content(&self, _content: Content) -> Option<Content>

Custom attributes for the module.

impl<B: Backend> ModuleDisplayDefault for MlpBurnPolicy<B>

fn content(&self, content: Content) -> Option<Content>

Attributes of the module used for display purposes.

fn num_params(&self) -> usize

Gets the number of the parameters of the module.