Module trainer 

Burn-backend PPO trainer (phase 3 of the Burn migration, #80).

Sibling to crate::train::ppo::PPOTrainerBurn (tch path). Both trainers implement the same clipped-surrogate, value-clip, entropy-bonus, KL-early-stop recipe; the only difference is the tensor backend and the optimizer ownership model.

Ownership model

Burn’s Optimizer<M, B> is move-through: every gradient step consumes the module by value and returns the updated copy. Phase 1’s scout (#78) confirmed this is the single biggest structural divergence between the two backends (Burn-migration friction point #1).

The Burn trainer therefore owns the policy module via an Option<P> field and swaps it through the optimizer on every step:

let module = self.policy.take().unwrap();
let grads = loss.backward();
let grads = GradientsParams::from_grads(grads, &module);
let module = self.optimizer.inner_mut().step(lr, module, grads);
self.policy = Some(module);

The tch trainer is struct PPOTrainer<P> with policy: P; the Burn trainer is struct PPOTrainerBurn<B, P, O> with the policy held in Option<P>. Phase 5 (#82) collapses the two when the ownership-model asymmetry goes away (only Burn remains).

Evaluating the policy

The trainer takes a closure evaluate_fn(&P, observations, actions) that returns (log_probs, entropy, values) exactly as the tch trainer does (see PPOTrainer::train_step_with_policy). This keeps the loss math identical and lets the caller plug in any module whose forward pass yields the right tensor shapes — including, for phase 4, the proper MlpBurnPolicy/SnakeCnnBurn ports.

Structs

PPOTrainerBurn

Burn-backend PPO trainer.