Quickstart

This tutorial trains a plain MLP actor-critic on CartpoleSwingup, a small mujoco_playground environment.

Setup

The tutorial assumes nnx-ppo, flax, jax, and mujoco_playground are importable from the active environment.

import mujoco_playground
from flax import nnx

from nnx_ppo.algorithms import ppo
from nnx_ppo.algorithms.config import TrainConfig, PPOConfig, EvalConfig
from nnx_ppo.networks.factories import make_mlp_actor_critic

Build the environment

Pick any registered Playground env. CartpoleSwingup is small enough to train in a few minutes on a laptop GPU.

env_name = "CartpoleSwingup"
env = mujoco_playground.registry.load(env_name)

The full list of registered envs lives in mujoco_playground.registry.

Build the network

make_mlp_actor_critic() returns a Sequential — an optional observation normalizer followed by a PPOAdapter whose two ports own the actor + sampler chain and the critic chain respectively.

SEED = 0
rngs = nnx.Rngs(SEED)
nets = make_mlp_actor_critic(
    obs_size=env.observation_size,
    action_size=env.action_size,
    actor_hidden_sizes=[64, 64],
    critic_hidden_sizes=[64, 64],
    rngs=rngs,
    activation=nnx.swish,
    normalize_obs=True,
)

The factory creates two MLPs (actor and critic), an action sampler, and — because normalize_obs=True — an observation normalizer that running-standardises each observation dimension. rngs is an flax.nnx.Rngs keychain used for parameter initialisation and for the stochastic action sampler. See Custom networks 1: composing with containers for how to build non-trivial networks by hand instead of using the factory.

Note that activation and normalize_obs are network-factory knobs (they shape how the network is built); they do not appear in the PPOConfig below, which only configures the training algorithm that operates on whatever network you hand it.

Configure training

PPO runs in iterations. In each iteration, the current policy (actor network) is first rolled out in the environment to collect a batch of experience. The rollout is parallelised across multiple environments to take advantage of GPU batching. Second, this batch of experience is used as a dataset to update both the actor and the critic using backpropagation and gradient descent.

The PPO algorithm has a number of hyperparameters that control this process. In this library, these parameters are specified in a PPOConfig.

During training we typically want to monitor progress. We can do this through regular evaluation runs (which are also batched across multiple environments). Configuration for these evaluation runs is specified in EvalConfig.

Finally, these two configurations are combined into a TrainConfig, along with a seed.

config = TrainConfig(
    ppo=PPOConfig(
        n_envs=512,
        rollout_length=20,
        total_steps=2_000_000,
        n_epochs=4,
        n_minibatches=8,
        learning_rate=3e-4,
    ),
    eval=EvalConfig(
        enabled=True,
        every_steps=100_000,
        n_envs=64,
        max_episode_length=1000,
    ),
    seed=SEED,
)

The PPOConfig fields used above:

n_envs

Number of environments rolled out in parallel per training iteration. Larger values give lower-variance gradient estimates at the cost of more GPU memory. The training loop applies jax.vmap() to the env across this batch for you — you do not need to wrap the env yourself.

rollout_length

Number of environment steps collected per environment, per iteration. The total batch size per iteration is n_envs * rollout_length.

total_steps

Stop training once the total number of environment steps taken across all parallel envs reaches this. Rollout length and n_envs are not adjusted — training simply ends after the iteration that pushes the cumulative step count past this threshold.

n_epochs

Number of passes over each rollout batch during the gradient phase.

n_minibatches

Number of minibatches each rollout batch is split into per epoch. Total gradient steps per iteration is n_epochs * n_minibatches.

learning_rate

Adam step size for the actor and critic.

Many more knobs (clip_range, gae_lambda, discounting_factor, entropy_weight, weight decay, gradient clipping, …) are available with sensible defaults. See Training parameters for the full list with explanations.

The EvalConfig fields used above:

enabled

Whether to run periodic evaluation rollouts. With enabled=False no evaluation is performed and the every_steps / n_envs / max_episode_length settings are ignored.

every_steps

Approximate interval between eval runs, measured in cumulative environment steps.

n_envs

Number of environments stepped in parallel during an eval rollout.

max_episode_length

Each eval episode is cut off at this many steps. Useful for envs that don’t terminate on their own.

The top-level seed controls JAX RNG initialisation throughout the training loop. Use the same seed as the nnx.Rngs(SEED) you passed to the network factory if you want fully reproducible runs.

Train

train_ppo() runs the whole loop: JIT-compiled rollouts, gradient updates, periodic eval, optional video and checkpointing. A log_fn receives the per-iteration metrics dict — useful for plotting eval curves without re-running eval manually.

history = []

def log_fn(metrics, steps):
    # train_ppo merges eval metrics into the per-iteration metrics
    # dict only on iterations where an eval ran. Use the presence of
    # episode_reward_mean as the "this step carried an eval" signal.
    if "episode_reward_mean" in metrics:
        history.append((steps, float(metrics["episode_reward_mean"])))
        print(f"step={steps}: reward={metrics['episode_reward_mean']:.2f}")

result = ppo.train_ppo(env=env, networks=nets, config=config, log_fn=log_fn)

print(f"Done: {result.total_steps} steps, "
      f"{result.total_iterations} PPO iterations")

log_fn is called once per training iteration with (metrics: dict, steps: int). Metrics include the loss components for that iteration; eval keys (episode_reward_mean, episode_reward_std, …) are merged in only on iterations where eval ran.

Inspecting the result

TrainResult exposes the final TrainingState, the eval history, and the final metrics:

for entry in result.eval_history:
    print(entry["step"], entry.get("episode_reward_mean"))

The network itself is still nets — train_ppo() mutates it in place. You can call nets.eval() to switch the action sampler to deterministic (use the mean instead of sampling) and then run the network forward yourself for ad-hoc inference.

To save and resume training across runs, pass a checkpoint_fn to train_ppo() and reconstruct the matching network on the way back — see Checkpointing for the details.

Next steps

If you need more complicated networks: