Struct PPOTrainerBurn 

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

Burn-backend PPO trainer.

Generic over:

• B: AutodiffBackend — the Burn backend (e.g. Autodiff<NdArray<f32>>, Autodiff<Wgpu>, etc.).
• P: AutodiffModule<B> — the policy module type.
• O: Optimizer<P, B> — the Burn optimizer (typically built from AdamConfig::new().init()).

The policy is held in Option<P> because Burn’s Optimizer::step consumes the module by value. We use .take() / put-back across each gradient step.

Implementations

impl<B, P, O> PPOTrainerBurn<B, P, O>

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

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

Build a new Burn PPO trainer.

Validates the config and stages the global gradient-norm cap (PPOConfig::max_grad_norm) on the optimizer wrapper; train_step applies it to the gradients of every minibatch step (issue #299).

pub fn config(&self) -> &PPOConfig

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.

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>, old_log_probs: Tensor<B, 1>, old_values: Tensor<B, 1>, advantages: Tensor<B, 1>, returns: Tensor<B, 1>, evaluate_fn: F, ) -> Result<TrainingStats>

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

Train for one PPO update.

Implements the same algorithm as PPOTrainer::train_step_with_policy on the tch path:

1. Normalize advantages to zero mean / unit variance.
2. For each of n_epochs epochs, shuffle the buffer and iterate minibatches.
3. Compute surrogate / value / entropy losses.
4. total_loss = policy + vf_coef * value + ent_coef * entropy.
5. Backprop, build GradientsParams, step the optimizer.
6. KL early stop if approx_kl > target_kl.
7. Entropy-collapse guard.

The evaluate_fn closure receives (&policy, obs, actions) and must return (log_probs, entropy, values) — exactly the same shape as the policy network’s evaluate_actions method.