configuring-checkpoint-strategy

name	Configuring Checkpoint Strategy
description	This skill should be used when the user asks to "configure checkpointing", "implement checkpoint strategy", "save best model", "set up EvalCallback", "implement periodic snapshots", "save replay buffer to disk", "resume training from checkpoint", or "implement EMA weights". Do NOT hallucinate parameters outside the boundaries of Configure Checkpoint Strategy.
version	0.1.0

Configuring Checkpoint Strategy

Checkpoint strategies preserve training state to enable crash recovery, model selection, and training trajectory reconstruction. Three concerns must be addressed: best-model selection via evaluation callbacks, periodic milestone snapshots, and replay buffer persistence for off-policy algorithms.

EvalCallback: Best-Model Selection

Implement a background thread EvalCallback that pauses training every eval_freq timesteps, runs the policy deterministically for 10 episodes, and saves if the current mean reward exceeds the best seen so far.

from stable_baselines3.common.callbacks import EvalCallback

eval_callback = EvalCallback(
    eval_env,
    best_model_save_path="./checkpoints/best/",
    log_path="./logs/eval/",
    eval_freq=50_000,          # timesteps between evaluations
    n_eval_episodes=10,        # deterministic rollouts
    deterministic=True,
    render=False,
)

Internally: if current_mean_reward > best_mean_reward, save best_model.zip. This is the model to deploy — never the final model.

Periodic Milestone Snapshots

Regardless of performance, save named snapshots at fixed timestep milestones to reconstruct the full training trajectory for later evaluation or distillation:

from stable_baselines3.common.callbacks import CheckpointCallback

checkpoint_callback = CheckpointCallback(
    save_freq=10_000_000,      # every 10M timesteps
    save_path="./checkpoints/snapshots/",
    name_prefix="model",       # produces model_10000000_steps.zip
    save_replay_buffer=False,  # handle replay buffer separately
    save_vecnormalize=True,
)

Snapshots named model_10M_steps, model_20M_steps allow offline comparison of policy quality across the training timeline.

Replay Buffer Persistence (Off-Policy Only)

For off-policy algorithms (SAC, TD3, DQN), model weights alone cannot resume a crashed run — the replay buffer must also be saved. The buffer often represents gigabytes of experience.

# Save buffer explicitly — can be a massive .pkl or .hdf5 file
model.save_replay_buffer("./checkpoints/replay_buffer.pkl")

# Load for resumption
model = SAC.load("./checkpoints/best/best_model.zip", env=env)
model.load_replay_buffer("./checkpoints/replay_buffer.pkl")

Memory management during saves: Large buffer writes can spike RAM. Release references explicitly and call gc.collect() after writing to prevent memory leaks during massive sequential saves.

import gc

model.save_replay_buffer("./checkpoints/replay_buffer.pkl")
gc.collect()

On-policy algorithms (PPO, TRPO) do not maintain long-term replay buffers — skip this step for them.

Combining Callbacks

Stack multiple callbacks with CallbackList:

from stable_baselines3.common.callbacks import CallbackList

callback = CallbackList([eval_callback, checkpoint_callback])
model.learn(total_timesteps=50_000_000, callback=callback)

Decision Checklist

Requirement	Solution
Deploy best policy	EvalCallback → best_model.zip
Reconstruct training trajectory	CheckpointCallback at milestone steps
Resume crashed off-policy run	Save + reload replay buffer
Compare sample efficiency	Load milestone snapshots, evaluate offline

Additional Resources

Scripts

• scripts/checkpoint_manager.py — EvalCallback setup, periodic snapshot scheduler, replay buffer save/load with memory cleanup

References

• references/checkpoint_patterns.md — EMA weight averaging, VecNormalize persistence, distributed checkpoint coordination