Skip to content

RL Audit Fixes: 2026-04-29

Dated audit record. It explains why the listed fixes were made; it is not a complete description of the current trainer surface.

Confirmed Fixes

  • IPPO and PR2-IPPO learning-rate annealing must keep schedule arithmetic in JAX arrays. Calling Python float() on the traced Optax schedule counter raises ConcretizationTypeError under JIT.
  • _coerce_seed accepts integer seeds, legacy jax.random.PRNGKey(...) arrays, and typed jax.random.key(...) values. Legacy keys encode the seed as [0, seed], so reading the first word collapses normal keys to 0.
  • IPPO and PR2-IPPO counterfactual replay pads relabeled PPO batches to a stable per-agent monitor-state upper bound. Synthetic pad slots carry zero-valued masks, and PPO losses/advantage normalization use masked means so padding does not change sample weighting.

Audit Notes

  • There is no CPO or IPPO-Lagrangian trainer in src/rl or the temporal-game experiment registry. CPO, ICPO, and PPO-LAG currently appear only as visualization labels.
  • Counterfactual replay still enumerates all DFA states. A future large-automata performance pass may need sampling or a configurable cap.