Struct MultiDiscreteMlpBurnPolicy 

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

Multi-discrete MLP actor-critic policy on Burn.

Shared trunk built from the same MlpBurnConfig knobs the single-action MlpBurnPolicy consumes, plus one Linear action head per dimension. Per-step log-probs are summed across dims (treating the dims as conditionally independent given the state), and per-step entropies are averaged.

Implementations

impl<B: Backend> MultiDiscreteMlpBurnPolicy<B>

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

Build a fresh multi-discrete policy with the default 2-layer architecture (mirrors crate::policy::multi_discrete_mlp::MultiDiscreteMlpBurnPolicy::new).

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

Seeded variant of new: same default architecture (orthogonal init), but constructed deterministically from seed so two calls with the same seed produce bit-identical weights (issue #135). Convenience wrapper for PSRO/NFSP policy factories over the multi-discrete policy.

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

Build a fresh multi-discrete policy with custom configuration.

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

Shared-trunk features (mirrors crate::policy::multi_discrete_mlp::MultiDiscreteMlpBurnPolicy::encoder_features).

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

Forward pass: per-dim action logits plus value estimate.

Returns (Vec<logits_i>, value) where logits_i: [batch, action_dims[i]] and value: [batch].

pub fn num_action_dims(&self) -> usize

Number of action dimensions (heads).

pub fn action_dim_cardinalities(&self) -> Vec<usize>

Per-dimension action cardinalities (one entry per head).

Returns the same vector that was passed to MultiDiscreteMlpBurnPolicy::with_config / MultiDiscreteMlpBurnPolicy::new. Reads each action head’s weight tensor shape — Burn’s burn::nn::Linear stores weight: Param<Tensor<B, 2>> with shape [d_input, d_output], so d_output is the per-dim cardinality. Used by the multi-agent joint trainer’s crate::multi_agent::joint::JointPolicy::action_dims_joint impl to size the rollout action buffer.

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

Sample one action per row per dim from the per-dim categorical distributions and return (actions_host, log_probs_host, values_host) as plain Vecs.

Thin backwards-compat wrapper around MultiDiscreteMlpBurnPolicy::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).

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.

Mirrors crate::policy::mlp::MlpBurnPolicy::get_action_host_seeded — the trainer-side rollout loop does not need gradient flow through the sampled action (only the eventual MultiDiscreteMlpBurnPolicy::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.

Layout:

• actions[row * num_dims + d] is the action for dim d of row row. Length is batch * num_dims.
• log_probs[row] is the joint log-probability summed across dims.
• values[row] is the value estimate.

Bit-exactness contract: two calls with the same obs, same policy state, and same-seeded rng (StdRng::seed_from_u64) produce element-wise identical (actions, log_probs, values). The per-row outer loop and per-dim inner loop consume RNG draws in a fixed row major → dim major order.

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 → dim 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 (same weights in the single batched forward, RNG consumed one draw per dim per row, ascending). The crate::multi_agent::joint::JointPolicy-trait batched entry point that eliminates per-call batch-1 overhead 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, 2, Int>, ) -> (Tensor<B, 1>, Tensor<B, 1>, Tensor<B, 1>)

Evaluate given actions: per-step summed log-prob, per-step mean entropy (across dims), and value.

Arguments

• obs - [batch, obs_dim]
• actions - [batch, num_dims] int (one action per dim per row)

Returns

(log_probs [batch], entropy [batch], values [batch]). log_probs is summed across dims; entropy is averaged across dims (matching the tch convention so parity holds).

Trait Implementations

impl<B> AutodiffModule<B> for MultiDiscreteMlpBurnPolicy<B>

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

type InnerModule = MultiDiscreteMlpBurnPolicy<<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 MultiDiscreteMlpBurnPolicy<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 MultiDiscreteMlpBurnPolicy<B>

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

Formats the value using the given formatter.

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

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

Formats the value using the given formatter.

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

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

type TrainModule = MultiDiscreteMlpBurnPolicy<B>

The module with auto-differentiation.

impl<B: AutodiffBackend> JointPolicy<B> for MultiDiscreteMlpBurnPolicy<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 MultiDiscreteMlpBurnPolicy<B>

type Record = MultiDiscreteMlpBurnPolicyRecord<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 MultiDiscreteMlpBurnPolicy<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 MultiDiscreteMlpBurnPolicy<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.