Skip to content

Conversation

@thopkins32
Copy link
Contributor

@thopkins32 thopkins32 commented Jan 13, 2026

Summary

Adds the ability save an optimizer checkpoint for the optimizers that implement the Checkpointable protocol.

AxOptimizer changes

  • Saves checkpoint to a JSON file (configurable path)
  • Can load a new AxOptimizer instance from a checkpoint

Agent changes

  • Saves optimizer state to a JSON file (configurable path)
  • Can load a new Agent instance from a checkpoint
  • Removed a lot of unused internal state (and public views of it) to simplify loading from a checkpoint
  • Rename evaluation argument to evaluation_function

optimize changes

  • Can choose a checkpoint interval to save the optimizer state periodically

Other comments

Needed to simplify the internal state of the Agent in order to reconstruct it from a checkpoint. The alternative was trying to read private Ax attributes to reconstruct Blop's DOFs, Objectives, constraints, etc. which is not trivial.

We had to use private member access once to get the parameter names of the loaded Ax experiment. Apparently there is no good alternative to this.

@thopkins32 thopkins32 merged commit b52aae1 into main Jan 13, 2026
15 checks passed
@thopkins32 thopkins32 deleted the model-checkpoints branch January 13, 2026 20:58
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants