Checkpointing

nnx-ppo can periodically write the full TrainingState to disk and reload it later to resume training (or to load a trained policy for inference). The checkpoint format is split across orbax (for the bulk array state) and pickle (for new-style PRNG-key variables, which orbax does not handle).

What’s saved

Each checkpoint persists every field of TrainingState:

• Network params (nnx.Param arrays for the actor, critic, samplers, heads, and any custom submodules) — via orbax.

• Non-param NNX variables — e.g. Normalizer’s running mean / M2 / counter, any custom nnx.Variable subclasses your code defines — via orbax.

• RNG variables on the network (nnx.RngKey, nnx.RngCount). The count goes through orbax; the key goes through pickle (because orbax does not yet support the key<…> dtype).

• Optimizer state — Adam moments, learning rate schedule state — via orbax.

• Carry state for the network and the env (network_states, env_states).

• The top-level training RNG key, the cumulative step count, and optionally the TrainConfig that produced the run.

Disk layout per checkpoint:

{directory}/step_{step:010d}/
    networks/          # orbax: non-key network variables
    optimizer/         # orbax: optimizer state
    metadata.pkl       # RngKey vars + TrainingState fields + step + config

Saving during training

Pass a checkpoint_fn to train_ppo(). The make_checkpoint_fn() helper builds one for you:

from nnx_ppo.algorithms.ppo import train_ppo
from nnx_ppo.algorithms.checkpointing import make_checkpoint_fn

result = train_ppo(
    env=env,
    networks=nets,
    config=config,
    checkpoint_fn=make_checkpoint_fn("/tmp/my_run", config=config),
)

Checkpoints are written every checkpoint_every_steps cumulative env steps. The config keyword is optional but recommended: it stores the TrainConfig alongside the checkpoint so the run is self-describing.

A checkpoint_fn is just a callable (training_state: TrainingState, step: int) -> None — you can plug in your own (e.g. to write to a remote object store) instead of using make_checkpoint_fn().

Resuming a run

Restoring a checkpoint requires a network instance with the same architecture as the one that was saved. The network’s params are overwritten in-place by the checkpoint values; the architecture itself is not reconstructed from disk.

from nnx_ppo.algorithms import ppo
from nnx_ppo.algorithms.checkpointing import load_checkpoint
from nnx_ppo.networks.factories import make_mlp_actor_critic

# 1. Rebuild the same network architecture you trained with.
nets = make_mlp_actor_critic(...)  # same kwargs as the original run
training_state = ppo.new_training_state(env, nets, n_envs, seed)

# 2. Load: weights / optimizer / carry / RNGs are restored in place.
ckpt = load_checkpoint(
    "/tmp/my_run/step_0000500000",
    training_state.networks,
    training_state.optimizer,
)

# 3. Continue training from where the checkpoint left off.
result = ppo.train_ppo(
    env, nets, ckpt["config"],
    initial_state=ckpt["training_state"],
)

The returned dict also contains ckpt["step"] (int) and ckpt["config"] (the persisted TrainConfig, or None if none was stored).

Loading for inference only

To load a trained policy without resuming training, build a network + optimizer template, call load_checkpoint(), and then ignore the optimizer / training-state fields:

nets = make_mlp_actor_critic(...)
training_state = ppo.new_training_state(env, nets, n_envs, seed)
load_checkpoint("/tmp/my_run/step_0000500000",
                training_state.networks, training_state.optimizer)
nets.eval()       # deterministic action sampler
# ... call nets(...) directly on observations from the env.