How to Save and Load Model Checkpoints Using Orbax in EvoRL

EvoRL uses Orbax as its checkpointing backend, providing both low-level helper functions in evorl/utils/orbax_utils.py for direct pytree serialization and a managed CheckpointManager for automatic interval-based saving during training workflows.

When working with evolutionary reinforcement learning experiments in EvoRL, persisting model state is critical for long-running training jobs and experiment reproducibility. The framework leverages Orbax (the official JAX checkpointing library) to handle serialization of JAX pytrees, offering both simple save/load utilities and a higher-level manager that integrates directly with training workflows.

Understanding EvoRL's Checkpoint Architecture

EvoRL implements a two-tiered approach to checkpointing. The low-level API provides direct control over serialization through save and load functions, while the managed API offers automated checkpointing through a CheckpointManager subclass that handles scheduling, retention policies, and metadata tracking.

Both approaches handle zero-size JAX arrays specially—filtering them out during save operations and restoring placeholder arrays during load to maintain pytree structure integrity.

Low-Level Checkpoint Operations with Orbax

The save and load Functions

Located in evorl/utils/orbax_utils.py, the save and load functions provide the simplest interface for persisting arbitrary Chex ArrayTree structures.

The save function normalizes the target path to an absolute location, filters zero-size arrays to prevent serialization errors, and delegates to ocp.StandardCheckpointer(). The load function reverses this process, requiring a dummy state (a pytree with the same structure as the saved data) to serve as a template for restoration.

from evorl.utils.orbax_utils import save, load
import chex

# Prepare model state (any Chex ArrayTree)

state = {"params": model_params, "opt_state": optimizer_state}

# Save checkpoint

checkpoint_path = "./experiments/run_001/checkpoint"
save(checkpoint_path, state)

# Load checkpoint (requires dummy structure for restoration)

dummy_state = {
    "params": chex.tree_util.tree_map(lambda _: 0, model_params),
    "opt_state": chex.tree_util.tree_map(lambda _: 0, optimizer_state)
}
restored_state = load(checkpoint_path, dummy_state)

Managed Checkpointing During Training

Configuring the CheckpointManager

For automated checkpointing during training loops, EvoRL provides setup_checkpoint_manager in evorl/utils/orbax_utils.py. This factory function constructs a custom CheckpointManager (subclassing ocp.CheckpointManager) that respects experiment configuration parameters.

The manager reads configuration fields from the experiment config (typically defined in YAML files under configs/agent/*.yaml):

• checkpoint.enable: Boolean toggle for checkpointing functionality
• checkpoint.save_interval_steps: Frequency of automatic saves (e.g., every 1000 steps)
• checkpoint.max_to_keep: Retention policy limiting stored checkpoints (e.g., last 5 checkpoints)
• output_dir: Root directory where checkpoints are stored in <output_dir>/checkpoints

Integration with Workflows

Every EvoRL workflow inherits from evorl/workflows/workflow.py, which automatically attaches a checkpoint manager via self.checkpoint_manager = setup_checkpoint_manager(config). During training, workflows check manager.should_save(step) to determine if the current step matches the configured interval, then call manager.save(step, state) to persist the full training state.

The manager automatically handles metadata serialization, including the resolved experiment configuration, alongside the model parameters.

# Example workflow integration

from evorl.workflows import Workflow

class CustomRLWorkflow(Workflow):
    def learn(self, state):
        for step in range(self.config.training.total_steps):
            # Perform training step

            metrics, state = self.step(state)
            self.recorder.record(metrics)
            
            # Automatic checkpointing

            if self.checkpoint_manager.should_save(step):
                self.checkpoint_manager.save(step, state)
        
        return state

Restoring from Checkpoints

To resume training or evaluate a saved model, use the checkpoint manager's restore method. The manager handles zero-size array restoration automatically.

from evorl.utils.orbax_utils import setup_checkpoint_manager
import omegaconf

# Load configuration

cfg = omegaconf.OmegaConf.load("configs/agent/ppo.yaml")
manager = setup_checkpoint_manager(cfg)

# Find latest checkpoint

latest_step = manager.latest_step()
if latest_step is not None:
    # Restore state (provide dummy structure for shape reference)

    restored_state = manager.restore(latest_step, dummy_state)

Summary

• EvoRL uses Orbax as its checkpointing backend, with utilities located in evorl/utils/orbax_utils.py.
• Low-level operations use save() and load() for direct pytree serialization, with automatic handling of zero-size JAX arrays.
• Managed checkpointing uses setup_checkpoint_manager() to create a CheckpointManager that automates save intervals, retention policies, and metadata tracking.
• Workflows automatically receive a checkpoint manager via self.checkpoint_manager and should call should_save(step) and save(step, items) during training loops.
• Checkpoints are stored in <output_dir>/checkpoints and can be restored using manager.restore(step, items) with a dummy state template.

Frequently Asked Questions

How do I enable automatic checkpointing in my EvoRL experiment?

Enable checkpointing by setting checkpoint.enable: true in your configuration file (typically under configs/agent/*.yaml). Specify the save frequency with checkpoint.save_interval_steps (e.g., 1000 for every 1000 steps) and set checkpoint.max_to_keep to control how many recent checkpoints to retain. The workflow will automatically create the checkpoint manager and handle saving during the training loop.

What is the difference between the save/load functions and the CheckpointManager?

The save and load functions in evorl/utils/orbax_utils.py provide immediate, one-off serialization—useful for debugging or manual saves where you specify the exact path and timing. The CheckpointManager provides automated, policy-driven checkpointing that handles scheduling (via save_interval_steps), retention (via max_to_keep), and metadata tracking automatically during long training runs. The manager also integrates seamlessly with EvoRL's workflow system.

How does EvoRL handle zero-size JAX arrays during checkpointing?

EvoRL filters out zero-size JAX arrays before saving to prevent serialization errors and storage inefficiency. In evorl/utils/orbax_utils.py, the save function removes these arrays using a filter function, and the load function restores placeholder zero-size arrays with the correct shapes when reconstructing the pytree. This ensures that model states containing empty arrays (common in certain neural network configurations) checkpoint correctly without losing structural information.

Can I restore a checkpoint from a different experiment configuration?

Yes, but you must ensure the pytree structure of the saved state matches what the current code expects. When calling manager.restore() or load(), provide a dummy_state (or items parameter) that has the same structure as the original saved state. The checkpoint manager will restore the arrays into this template. If the model architecture has changed (different layer shapes or names), you may need to manually map or reshape the restored parameters before use.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}