Skip to content

Make evaluation checkpoint writes atomic with required interface metadata #18

Description

@NeoLorenzo

Audit Priority: P2
Audit Type: FIX

Problem / Opportunity

PolyVision treats a .cleanrl_model checkpoint and its adjacent .action_interface.json sidecar as one evaluation artifact, but the trainer does not persist them as one successful unit. Both periodic and final save paths write model weights first, then write the sidecar inside try/except Exception; if the metadata write fails, the trainer only prints a warning and still prints checkpoint_saved=... or model_saved=....

Maintained evaluators fail closed when that sidecar is missing or unreadable, so the trainer can currently report a checkpoint as successfully saved even though the repository's supported evaluation path cannot load it.

Why This Matters

Long Phase 1 runs rely on periodic model snapshots for later fixed-protocol evaluation. A filesystem, serialization, interrupted-write, or metadata-write failure should not leave an apparently successful but incomplete checkpoint artifact. This is particularly confusing for long runs because the weights file may exist and look valid while the compatibility metadata required to establish its action/observation contract is absent or stale.

The integrity of evaluation-only checkpoints is independent of whether full trainer state is resumable: issue #4 owns resumable PPO checkpoints, while this issue protects the lightweight model-plus-interface artifact that current evaluation already depends on.

Evidence

  • py_rl/cleanrl/cleanrl/ppo.py periodically calls torch.save(agent.state_dict(), checkpoint_path) and then writes checkpoint_path + ".action_interface.json" inside a broad try/except; metadata failures are reduced to warning: failed to save action interface metadata, after which checkpoint_saved=... is still printed.
  • The final --save-model path repeats the same pattern and still prints model_saved=... after a sidecar failure.
  • pol_env/Tribes/py/environment_contract.py::read_checkpoint_metadata() raises CheckpointCompatibilityError when the adjacent sidecar is missing, unreadable, or not a JSON object.
  • tools/evaluate_phase1.py reads checkpoint metadata before constructing the PPO evaluation policy, and evaluate_brain.py uses the same fail-closed compatibility path.
  • docs/training.md states that every saved model has an adjacent .action_interface.json sidecar and that evaluators require it.
  • docs/reproducibility.md likewise requires preserving the checkpoint sidecar as part of run provenance.
  • Open issue Add resumable full-state checkpoints for long PPO training runs #4 is about full-state training continuation and explicitly retains lightweight evaluation snapshots; it does not define failure semantics for writing the existing model-plus-sidecar pair.

Proposed Outcome

Treat the model weights and required action-interface metadata as one checkpoint publication operation. A checkpoint should be announced/left as successfully published only when both files have been written successfully and the sidecar satisfies the maintained metadata contract. Failure to write either component should fail visibly and must not masquerade as a valid completed checkpoint.

Use an implementation that avoids exposing a newly published half-written pair where practical, for example temporary files plus atomic replacement in the target directory.

Definition of Done

  • Periodic and final model-save paths use one maintained checkpoint-writing path for weights plus action-interface metadata.
  • A sidecar write/serialization failure is no longer swallowed while the corresponding checkpoint is reported as saved.
  • A successful checkpoint publication guarantees that both the .cleanrl_model file and adjacent .action_interface.json exist and the metadata is readable under the maintained compatibility contract.
  • Failed or interrupted publication does not leave a newly advertised orphan weights file or silently reuse an incompatible stale sidecar as the new checkpoint's metadata.
  • Success logging occurs only after the complete model-plus-sidecar artifact has been published.
  • Regression coverage simulates a metadata-write failure and verifies fail-closed behavior, plus a successful save that can be read through read_checkpoint_metadata().
  • Existing evaluator compatibility checks and the current lightweight checkpoint format remain usable.

Constraints / Non-Goals

This issue does not add optimizer/RNG/environment state or make training resumable; that work remains in #4. It also does not change checkpoint compatibility semantics or require a new model serialization format beyond what is needed to publish the existing evaluation artifact safely.

Validation

Exercise the checkpoint writer in a temporary directory. Verify a normal save produces a readable weights/sidecar pair, then inject a sidecar-write failure and confirm the operation fails without emitting a successful checkpoint message or leaving a newly valid-looking partial artifact.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions