Skip to content

Verification Rollout Performance

Approximate NashConv and final-satisfaction evaluation reuse the rollout helpers in src/verify/nash_conv/. For neural policies, those helpers must use jitted, batched action selectors rather than scalar per-action network.apply calls.

The slow path was especially visible for RNN IQL on native gridworlds: each evaluation episode can run for max_cycles steps, and calling Flax/JAX inference from a Python scalar helper once per agent per step made NashConv far slower than training. The fixed IQL and PR2-IQL rollouts now reuse the same jitted selector shape as the trainers:

  • IQL uses src.rl.iql.make_action_selection_fn.
  • PR2-IQL uses src.rl.pr2_iql.make_pr2_action_selection_fn.

IPPO and PR2-IPPO did not show the same severe slowdown in the observed logs, but their verifier rollouts had the same unjitted scalar inference pattern. They now build local jitted batched selectors for actor sampling/greedy action selection. The legacy scalar helpers remain for compatibility and focused unit tests, but rollout code should not call them.

Nash-Q is tabular and does not use JAX/Flax inference during rollout. Its NashConv evaluator already caches solved stage-game policies within each rollout call. The main Nash-Q cost risk is still the number of distinct tabular states and associated pygambit solves, not per-step neural inference dispatch.

Regression coverage in tests/test_ippo_nashconv.py patches the scalar neural action helpers to fail and then exercises IPPO, PR2-IPPO, IQL, and PR2-IQL rollouts. This is intended to catch accidental reintroduction of the slow neural NashConv path.

After the fix, a CPU smoke run using the saved SafeHarvest 4-agent IQL-CER checkpoint from exports/gridworlds/harvest/safe_harvest/4/iql_cer completed one 1000-cycle evaluation episode in about 10 seconds. The pre-fix Slurm log for the same family was spending roughly 30 minutes per displayed NashConv episode unit, so this is the practical regression signal to preserve.