Struct A2cTrainer 

pub struct A2cTrainer<B, P, O>
where
    B: AutodiffBackend,
    P: AutodiffModule<B>,
    O: Optimizer<P, B>,
{ /* private fields */ }

Burn-backend A2C trainer.

Generic over:

• B: AutodiffBackend — the Burn backend (e.g. Autodiff<NdArray<f32>>).
• P: AutodiffModule<B> — the shared actor-critic policy module.
• O: Optimizer<P, B> — the Burn optimizer (typically AdamConfig::new().init()).

The policy is held in Option<P> because Burn’s Optimizer::step consumes the module by value; we .take() it and put back the updated copy across the single gradient step.

Implementations

impl<B, P, O> A2cTrainer<B, P, O>

where B: AutodiffBackend, P: AutodiffModule<B> + Clone, O: Optimizer<P, B>,

pub fn new( config: A2cConfig, policy: P, optimizer: BurnOptimizer<B, P, O>, ) -> Result<Self>

Build a new Burn A2C trainer.

Validates the config and stages the global gradient-norm clip (A2cConfig::max_grad_norm) on the optimizer wrapper.

pub fn config(&self) -> &A2cConfig

Borrow the configuration.

pub fn policy(&self) -> &P

Borrow the policy. Panics if the trainer is mid-step (the policy has been moved into the optimizer); only safe to call between train_step invocations.

pub fn total_steps(&self) -> usize

Total completed gradient updates (one per train_step).

pub fn total_episodes(&self) -> usize

Total completed episodes (caller increments).

pub fn increment_episodes(&mut self, n: usize)

Increment the episode counter.

pub fn train_step<F>( &mut self, observations: Tensor<B, 2>, actions: Tensor<B, 1, Int>, advantages: Tensor<B, 1>, returns: Tensor<B, 1>, evaluate_fn: F, ) -> Result<A2cStats>

where F: FnMut(&P, Tensor<B, 2>, Tensor<B, 1, Int>) -> (Tensor<B, 1>, Tensor<B, 1>, Tensor<B, 1>),

Train for one A2C update.

Performs exactly one gradient step over the whole rollout:

1. Optionally normalize advantages to zero-mean/unit-variance when A2cConfig::normalize_advantages is set.
2. Evaluate the policy to get (log_probs, entropy, values).
3. policy_loss = -mean(log_prob * advantage) (no ratio, no clip).
4. value_loss = mean((V(s) - returns)^2) (plain MSE, no clip).
5. entropy_loss = -mean(entropy).
6. `total = policy_loss + value_coef * value_loss
   • entropy_coef * entropy_loss`.
7. Backprop, build GradientsParams, step the optimizer once.
8. Entropy-collapse guard (shared with PPO).

Note the absence of old_log_probs / old_values: A2C is on-policy with a single update, so there is no behaviour policy to form an importance ratio against, and no old value baseline to clip against.

The evaluate_fn closure receives (&policy, obs, actions) and must return (log_probs, entropy, values) — exactly the shape of MlpBurnPolicy::evaluate_actions.