Checkpointing

Agents are saved with their trainer, so a checkpoint contains everything needed to keep training or to evaluate:

agent.save("ppo_cartpole.eqx")

from jaxnasium.algorithms import RLAgent

agent = RLAgent.load("ppo_cartpole.eqx")
agent, metrics = agent.train(key, env)

This requires jaxon, which is not installed by default (pip install jaxon). Alternatively, agents are plain Equinox modules and can be serialized like any other.

jaxnasium.algorithms.core.save_agent(file_path: str, agent: Any) -> None

Writes agent -- parameters, optimizer state and trainer -- to file_path using jaxon.

jaxnasium.algorithms.core.load_agent(file_path: str) -> Any

Reads a checkpoint written by save_agent using jaxon.