# How to Save and Load Model Checkpoints Using Orbax in EvoRL

> Learn to efficiently save and load model checkpoints using Orbax in EvoRL. Explore helper functions and the Checkpoint Manager for seamless training.

- Repository: [EMI-Group/evorl](https://github.com/emi-group/evorl)
- Tags: how-to-guide
- Published: 2026-03-01

---

**EvoRL uses Orbax as its checkpointing backend, providing both low-level helper functions in [`evorl/utils/orbax_utils.py`](https://github.com/emi-group/evorl/blob/main/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`](https://github.com/emi-group/evorl/blob/main/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.

```python
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`](https://github.com/emi-group/evorl/blob/main/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`](https://github.com/emi-group/evorl/blob/main/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.

```python

# 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.

```python
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`](https://github.com/emi-group/evorl/blob/main/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`](https://github.com/emi-group/evorl/blob/main/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`](https://github.com/emi-group/evorl/blob/main/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.