Add resumable full-state checkpoints for long PPO training runs

[NeoLorenzo]

Audit Priority: P2
Audit Type: BUILD

Problem / Opportunity

PolyVision's periodic training checkpoints are inference/evaluation snapshots rather than resumable training checkpoints. py_rl/cleanrl/cleanrl/ppo.py writes agent.state_dict() plus an action-interface sidecar, but does not persist optimizer state, trainer progress, RNG state, or environment/map-stream state. The periodic save occurs during rollout collection, which also makes the named checkpoint step an unsafe continuation boundary without additional state capture.

Why This Matters

The active reference run is 16,000,000 environment steps and took about 7.7 hours, and future multi-seed experiments will multiply that cost. A process, machine, or JVM failure currently forces the affected training run to restart from scratch if scientific continuity is required. Resumable checkpoints would reduce wasted compute and make longer future curricula materially safer to run.

Evidence

• py_rl/cleanrl/cleanrl/ppo.py performs periodic checkpointing inside the rollout-collection loop and saves only agent.state_dict() to model_checkpoint_<step>.cleanrl_model, plus an action-interface JSON sidecar.
• Because a save-frequency milestone can be crossed by a vectorized environment step, the filename's checkpoint_step can describe a milestone crossed during the current rollout rather than a complete trainer/update boundary.
• The same trainer's final save also writes only model weights plus interface metadata.
• docs/training.md explicitly states: 'The trainer does not implement checkpoint resume or optimizer-state restoration.'
• docs/reproducibility.md likewise records that the trainer saves model weights only and does not save optimizer/RNG state for exact training continuation.
• The current v5 reference run documented in docs/training.md uses 16M steps across 20 environments and reports a multi-hour training run, so interruption recovery is already relevant at the project's present scale.

Proposed Outcome

Add a distinct resumable training-checkpoint format and resume path that preserves the complete state required to continue a run under the same experiment contract. Keep lightweight model-only checkpoints available for evaluation, but make the difference explicit. Resumable checkpoints should be written only at a well-defined safe trainer boundary and should fail closed when environment/interface, map-pool, reward, or trainer configuration is incompatible.

Definition of Done

• The trainer can save and load a resumable checkpoint containing model weights, optimizer state, global step/iteration, learning-rate schedule position, relevant Python/NumPy/PyTorch RNG states, and any environment/map-stream state required by the chosen continuation contract.
• Resumable checkpoints are written at a documented safe boundary rather than being ambiguous mid-rollout snapshots.
• Resume validates the same environment/action-interface and pool identities used for evaluation compatibility, plus training-critical reward and PPO configuration.
• A deterministic small-run regression test compares uninterrupted training with save/resume training under a controlled configuration and verifies the documented continuation guarantee.
• CLI and documentation clearly distinguish evaluation-only model snapshots from resumable training checkpoints and show how to resume an interrupted run.

Constraints / Non-Goals

Existing lightweight model-only snapshots should remain usable for evaluation. This issue does not require arbitrary mid-step recovery; the continuation guarantee may deliberately use a safe rollout/update boundary as long as that contract is explicit and tested.

Validation

Run a controlled short experiment twice from the same seed/configuration: once uninterrupted and once with save/exit/resume at the supported checkpoint boundary. Compare the trainer state and resulting trajectory according to the documented continuation guarantee, and separately verify that incompatible environment/interface or training configuration is rejected.