OakEmu-RL v0.1.0

Appearance

PPOMinibatchUpdater ​

python
PPOMinibatchUpdater(config: PPOExperimentConfiguration, optimizer: torch.optim.Optimizer, device: torch.device, scaler: torch.amp.GradScaler, action_mask: torch.Tensor | None = None) -> None

Performs one PPO minibatch update, including PPO loss, BC loss, and optimizer step.

Initialize the minibatch updater.

Parameters:

  • config (PPOExperimentConfiguration) – PPO experiment configuration.
  • optimizer (Optimizer) – Learner optimizer.
  • device (device) – Learner device.
  • scaler (GradScaler) – Mixed-precision gradient scaler.
  • action_mask (Tensor | None) – Optional learner action mask.

Source

Methods ​

update_minibatch ​

python
update_minibatch(ppo_minibatch: RolloutMinibatch | None, bc_minibatch: BehaviorCloningMinibatch | None, model: PPOModel, weights: LossWeights) -> PPOMinibatchMetrics

Perform a single PPO + optional replay-BC update.

Parameters:

  • ppo_minibatch (RolloutMinibatch | None) – Optional on-policy PPO minibatch.
  • bc_minibatch (BehaviorCloningMinibatch | None) – Optional replay-only minibatch for behaviour cloning.
  • model (PPOModel) – The model to update.
  • weights (LossWeights) – The PPO/BC loss coefficients for this rollout.

Returns:

  • PPOMinibatchMetrics – Loss components, diagnostics, and gradient norm.

Source

update_value_head ​

python
update_value_head(ppo_minibatch: RolloutMinibatch, model: PPOModel) -> PPOMinibatchMetrics

Take one optimizer step on the value head alone, with the shared trunk features detached.

The BC bootstrap's critic-calibration step: the value head learns to predict returns from the features BC is imprinting, without writing to them. Letting the value loss train the trunk during the bootstrap destroyed the policy's near-manifold generalisation (2026-08-07): the cold critic's gradient re-shaped the trunk at the full clip budget every minibatch, BC re-pinned the policy on the exact demo observations but nothing defended slightly-off-demo inputs (for example the tester's initial frame-stack padding), so midpoint (exact-state) wins climbed while the from-start tester collapsed to zero.

Parameters:

  • ppo_minibatch (RolloutMinibatch) – An on-policy rollout minibatch.
  • model (PPOModel) – The model whose value head to update.

Returns:

  • PPOMinibatchMetrics – The update metrics; policy-side fields are zero.

Source