Algorithms

Core training algorithm, rollout logic, configuration, and utilities.

Types

Runtime types for the PPO algorithm.

class nnx_ppo.algorithms.types.EnvState(*args, **kwargs)

Bases: Protocol

Minimal environment state interface.

Satisfied by mujoco_playground.State and any compatible environment state. Fields are declared as read-only properties so that frozen dataclasses (which have immutable attributes) satisfy this protocol. Note: replace() is intentionally omitted — playground’s State has it at runtime (via @struct.dataclass) but not in its type stubs.

property obs: Any

property done: Shaped[Array, '...']

property reward: Any

property info: dict[str, Any]

property metrics: dict[str, Any]

__init__(*args, **kwargs)

class nnx_ppo.algorithms.types.RLEnv(*args, **kwargs)

Bases: Protocol

Minimal RL environment interface.

Satisfied by mujoco_playground.MjxEnv and any compatible environment.

reset(rng)

step(state, action)

__init__(*args, **kwargs)

class nnx_ppo.algorithms.types.TrainingState(networks: Any, network_states: Any, env_states: Any, optimizer: Any, rng_key: jaxtyping.Key[Array, ''] | jaxtyping.UInt32[Array, '2'] | jaxtyping.UInt32[Array, '4'], steps_taken: jaxtyping.Float[Array, ''])

Bases: JaxDataclass

networks: Any

network_states: Any

env_states: Any

optimizer: Any

rng_key: Key[Array, ''] | UInt32[Array, '2'] | UInt32[Array, '4']

steps_taken: Float[Array, '']

__init__(networks, network_states, env_states, optimizer, rng_key, steps_taken)

class nnx_ppo.algorithms.types.Transition(obs, network_output, rewards, done, truncated, next_obs, metrics, rollout_extras=None)

Bases: JaxDataclass

Environment state for training and inference.

Note: rewards, done, truncated use *time to allow both [batch] (single timestep) and [time, batch] (full rollout) shapes.

obs: PyTree[jaxtyping.Float[Array, '...']]

network_output: PPONetworkOutput

rewards: PyTree[jaxtyping.Float[Array, '*time batch']]

done: Bool[Array, '*time batch']

truncated: Bool[Array, '*time batch']

next_obs: PyTree[jaxtyping.Float[Array, '...']]

metrics: dict[str, Any]

rollout_extras: Any = None

__init__(obs, network_output, rewards, done, truncated, next_obs, metrics, rollout_extras=None)

class nnx_ppo.algorithms.types.DistillationTransition(obs, student_output, rewards, done, truncated, next_obs, metrics, student_rollout_extras=None, teacher_rollout_extras=None)

Bases: JaxDataclass

Rollout transition for distillation training.

Carries both student and teacher network outputs. The student’s actions drive the environment; the teacher’s outputs (pre-computed outside of gradient computation) serve as the distillation target.

obs: PyTree[jaxtyping.Float[Array, '...']]

student_output: PPONetworkOutput

rewards: PyTree[jaxtyping.Float[Array, '*time batch']]

done: Bool[Array, '*time batch']

truncated: Bool[Array, '*time batch']

next_obs: PyTree[jaxtyping.Float[Array, '...']]

metrics: dict[str, Any]

student_rollout_extras: Any = None

teacher_rollout_extras: Any = None

__init__(obs, student_output, rewards, done, truncated, next_obs, metrics, student_rollout_extras=None, teacher_rollout_extras=None)

class nnx_ppo.algorithms.types.DistillationState(student, student_states, teacher_states, env_states, optimizer, rng_key, steps_taken)

Bases: JaxDataclass

Training state for distillation.

Mirrors TrainingState but has separate student and teacher states. The teacher is passed as an external argument (like env) and is not stored here. Only the teacher’s per-env carry state is tracked.

student: Any

student_states: Any

teacher_states: Any

env_states: Any

optimizer: Any

rng_key: Key[Array, ''] | UInt32[Array, '2'] | UInt32[Array, '4']

steps_taken: Float[Array, '']

__init__(student, student_states, teacher_states, env_states, optimizer, rng_key, steps_taken)

class nnx_ppo.algorithms.types.LoggingLevel(*values)

Bases: Flag

LOSSES = 1

CRITIC_EXTRA = 2

ACTOR_EXTRA = 4

ROLLOUT_STATS = 8

ROLLOUT_OBS = 16

ENV_METRICS = 32

NETWORK_METRICS = 64

GRAD_NORM = 128

WEIGHTS = 256

THROUGHPUT = 512

BASIC = 1

ALL = 1007

NONE = 0

Configuration

Configuration dataclasses for train_ppo and train_distillation.

class nnx_ppo.algorithms.config.PPOConfig(n_envs=256, rollout_length=20, total_steps=512000, gae_lambda=0.95, discounting_factor=0.99, clip_range=0.2, learning_rate=0.0001, normalize_advantages=True, combine_advantages=False, n_epochs=4, n_minibatches=4, critic_loss_weight=1.0, gradient_clipping=None, weight_decay=None, logging_level=<LoggingLevel.LOSSES: 1>, logging_percentiles=None)

Bases: object

Core PPO algorithm parameters.

n_envs: int = 256

rollout_length: int = 20

total_steps: int = 512000

gae_lambda: float = 0.95

discounting_factor: float = 0.99

clip_range: float = 0.2

learning_rate: float = 0.0001

normalize_advantages: bool = True

combine_advantages: bool = False

n_epochs: int = 4

n_minibatches: int = 4

critic_loss_weight: float = 1.0

gradient_clipping: float | None = None

weight_decay: float | None = None

logging_level: LoggingLevel = 1

logging_percentiles: tuple[int, ...] | None = None

__init__(n_envs=256, rollout_length=20, total_steps=512000, gae_lambda=0.95, discounting_factor=0.99, clip_range=0.2, learning_rate=0.0001, normalize_advantages=True, combine_advantages=False, n_epochs=4, n_minibatches=4, critic_loss_weight=1.0, gradient_clipping=None, weight_decay=None, logging_level=<LoggingLevel.LOSSES: 1>, logging_percentiles=None)

class nnx_ppo.algorithms.config.EvalConfig(enabled=True, every_steps=50000, n_envs=64, max_episode_length=1000, logging_level=<LoggingLevel.NONE: 0>, logging_percentiles=(0, 25, 50, 75, 100))

Bases: object

Evaluation rollout configuration.

enabled: bool = True

every_steps: int = 50000

n_envs: int = 64

max_episode_length: int = 1000

logging_level: LoggingLevel = 0

logging_percentiles: tuple[int, ...] | None = (0, 25, 50, 75, 100)

__init__(enabled=True, every_steps=50000, n_envs=64, max_episode_length=1000, logging_level=<LoggingLevel.NONE: 0>, logging_percentiles=(0, 25, 50, 75, 100))

class nnx_ppo.algorithms.config.VideoConfig(enabled=False, every_steps=200000, episode_length=1000, render_kwargs=<factory>)

Bases: object

Video recording configuration.

enabled: bool = False

every_steps: int = 200000

episode_length: int = 1000

render_kwargs: dict[str, Any]

__init__(enabled=False, every_steps=200000, episode_length=1000, render_kwargs=<factory>)

class nnx_ppo.algorithms.config.TrainConfig(ppo=<factory>, eval=<factory>, video=<factory>, seed=17, checkpoint_every_steps=500000)

Bases: object

Complete training configuration.

ppo: PPOConfig

eval: EvalConfig

video: VideoConfig

seed: int = 17

checkpoint_every_steps: int = 500000

__init__(ppo=<factory>, eval=<factory>, video=<factory>, seed=17, checkpoint_every_steps=500000)

class nnx_ppo.algorithms.config.DistillationConfig(n_envs=256, rollout_length=20, total_steps=512000, learning_rate=0.0001, n_epochs=4, n_minibatches=4, gradient_clipping=None, weight_decay=None, logging_level=<LoggingLevel.LOSSES: 1>, logging_percentiles=None)

Bases: object

Core distillation algorithm parameters.

n_envs: int = 256

rollout_length: int = 20

total_steps: int = 512000

learning_rate: float = 0.0001

n_epochs: int = 4

n_minibatches: int = 4

gradient_clipping: float | None = None

weight_decay: float | None = None

logging_level: LoggingLevel = 1

logging_percentiles: tuple[int, ...] | None = None

__init__(n_envs=256, rollout_length=20, total_steps=512000, learning_rate=0.0001, n_epochs=4, n_minibatches=4, gradient_clipping=None, weight_decay=None, logging_level=<LoggingLevel.LOSSES: 1>, logging_percentiles=None)

class nnx_ppo.algorithms.config.DistillationTrainConfig(distillation=<factory>, eval=<factory>, video=<factory>, seed=17, checkpoint_every_steps=500000)

Bases: object

Complete training configuration for distillation.

distillation: DistillationConfig

eval: EvalConfig

video: VideoConfig

seed: int = 17

checkpoint_every_steps: int = 500000

__init__(distillation=<factory>, eval=<factory>, video=<factory>, seed=17, checkpoint_every_steps=500000)

class nnx_ppo.algorithms.config.VideoData(frames, step, episode_reward, episode_length)

Bases: object

Data passed to video callback.

frames: ndarray

step: int

episode_reward: float

episode_length: int

__init__(frames, step, episode_reward, episode_length)

class nnx_ppo.algorithms.config.TrainResult(training_state, final_metrics, eval_history, total_steps, total_iterations)

Bases: object

Result of train_ppo containing final state and summary.

training_state: TrainingState

final_metrics: dict[str, Any]

eval_history: list[dict[str, Any]]

total_steps: int

total_iterations: int

__init__(training_state, final_metrics, eval_history, total_steps, total_iterations)

class nnx_ppo.algorithms.config.DistillationTrainResult(training_state, final_metrics, eval_history, total_steps, total_iterations)

Bases: object

Result of train_distillation containing final state and summary.

training_state: DistillationState

final_metrics: dict[str, Any]

eval_history: list[dict[str, Any]]

total_steps: int

total_iterations: int

__init__(training_state, final_metrics, eval_history, total_steps, total_iterations)

Training

nnx_ppo.algorithms.ppo.default_config()

Return default training configuration.

nnx_ppo.algorithms.ppo.train_ppo(env, networks, config=None, *, total_steps=None, seed=None, log_fn=None, video_fn=None, checkpoint_fn=None, eval_env=None, initial_state=None)

Train a PPO agent.

Parameters:

• env (RLEnv) – Training environment (MjxEnv).

• networks (StatefulModule) – PPO network (actor-critic).

• config (TrainConfig | None) – Training configuration. If None, uses default_config().

• total_steps (int | None) – Override config.ppo.total_steps (convenience parameter).

• seed (int | None) – Override config.seed (convenience parameter).

• log_fn (Callable[[dict[str, Any], int], None] | None) – Called with (metrics_dict, step) after each PPO step. If None, no logging is performed.

• video_fn (Callable[[VideoData], None] | None) – Called with VideoData after rendering eval episodes. If None, no videos are recorded even if config.video.enabled.

• checkpoint_fn (Callable[[TrainingState, int], None] | None) – Called with (training_state, step) at config.checkpoint_every_steps intervals. If None, no checkpointing is performed.

• eval_env (RLEnv | None) – Environment for evaluation rollouts. If None, uses env.

• initial_state (TrainingState | None) – Resume training from an existing TrainingState. If None, creates a new TrainingState.

Returns:

TrainResult containing final TrainingState, metrics, and eval history.

Return type:

TrainResult

nnx_ppo.algorithms.ppo.ppo_step(env, training_state, n_envs, rollout_length, gae_lambda, discounting_factor, clip_range, normalize_advantages, combine_advantages, n_epochs, n_minibatches, critic_loss_weight=1.0, logging_level=<LoggingLevel.LOSSES: 1>, logging_percentiles=None)

nnx_ppo.algorithms.ppo.gae(rewards, values_excl_last, last_value, done, truncation, lambda_, gamma)

nnx_ppo.algorithms.ppo.ppo_loss(networks, network_state, rollout_data, clip_range, normalize_advantages, combine_advantages, discounting_factor, gae_lambda, critic_loss_weight, logging_level)

nnx_ppo.algorithms.ppo.new_training_state(env, networks, n_envs, seed, learning_rate=0.0001, gradient_clipping=None, weight_decay=None)

Rollout

nnx_ppo.algorithms.rollout.single_transition(env, networks, carry, rng_keys_for_env_reset)

nnx_ppo.algorithms.rollout.unroll_env(env, env_state, networks, network_state, unroll_length, rng_key_for_env_reset)

nnx_ppo.algorithms.rollout.eval_rollout(env, networks, n_envs, max_episode_length, key, logging_percentiles=None, logging_level=<LoggingLevel.NONE: 0>)

nnx_ppo.algorithms.rollout.record_activations_rollout(env, networks, n_envs, max_episode_length, key)

Roll out a deterministic episode and return per-step unit activations.

Convenience wrapper for activation analysis. networks is made recordable with with_recording() (the original is left untouched) and run in eval mode, so action samplers are deterministic. Each step’s activations are stacked over time via the scan’s ys return — this is not eval_rollout, which reduces metrics to scalars and so cannot expose per-unit values.

Unlike eval_rollout(), environments are not reset on termination: each env runs a single episode for max_episode_length steps. The returned dones flag is the pre-step “already terminated” mask (matching eval_rollout’s accounting); callers should mask out steps where it is set.

Memory cost: max_episode_length × n_envs × Σ units is materialised on the device. On a small GPU prefer a modest n_envs (e.g. a handful) and/or a shorter max_episode_length.

Returns:

A pytree mirroring the network’s container structure, each

leaf an array with leading dims [max_episode_length, n_envs, ...], keyed by module path (e.g. action/0, value/...); the population graph contributes one entry per population.

dones: [max_episode_length, n_envs] pre-step termination mask.

Return type:

activations

class nnx_ppo.algorithms.rollout.SlimData(qpos, qvel, time, mocap_pos, mocap_quat, xfrc_applied)

Bases: NamedTuple

Minimal mjx.Data fields needed for rendering.

qpos: Any

Alias for field number 0

qvel: Any

Alias for field number 1

time: Any

Alias for field number 2

mocap_pos: Any

Alias for field number 3

mocap_quat: Any

Alias for field number 4

xfrc_applied: Any

Alias for field number 5

class nnx_ppo.algorithms.rollout.SlimState(data, done, info, metrics)

Bases: NamedTuple

Minimal env state for rendering, avoiding large mjx.Data contact buffers.

data: SlimData

Alias for field number 0

done: Any

Alias for field number 1

info: Any

Alias for field number 2

metrics: Any

Alias for field number 3

nnx_ppo.algorithms.rollout.eval_rollout_for_render_scan(env, networks, max_episode_length, key)

JIT-compatible scan-based rollout that returns stacked slim states.

Returns stacked SlimState (qpos/qvel/time/info/done/metrics only) rather than the full env state, avoiding large MuJoCo Warp contact buffers.

Returns:

SlimState pytree with leading dimension of max_episode_length. final_state: The final environment slim state. total_reward: Total reward accumulated during the episode.

Return type:

stacked_states

nnx_ppo.algorithms.rollout.unstack_trajectory(stacked_states, final_state, max_episode_length)

Convert stacked states from scan to a list for rendering.

This must be called outside of JIT since it creates a Python list.

nnx_ppo.algorithms.rollout.tree_where(cond, on_true, on_false)

Checkpointing

Checkpointing utilities for saving and loading training state.

class nnx_ppo.algorithms.checkpointing.CheckpointCallback(*args, **kwargs)

Bases: Protocol

Protocol for checkpoint callbacks with named parameters.

__call__(training_state, step)

Call self as a function.

nnx_ppo.algorithms.checkpointing.make_checkpoint_fn(directory, config=None)

Create a checkpoint callback that saves TrainingState to disk.

Each checkpoint is written to {directory}/step_{step:010d}/, containing:

• networks/ — orbax checkpoint with all non-PRNG-key network variables (Param, RngCount, NormalizerStatistics, etc.)

• optimizer/ — orbax checkpoint with all optimizer state arrays

• metadata.pkl — pickle file with network RngKey variables, all remaining TrainingState fields (network_states, env_states, rng_key, steps_taken), the step count, and the optional TrainConfig.

To resume training from a checkpoint, use load_checkpoint().

Parameters:

• directory (str) – Base directory under which checkpoint subdirectories are created.

• config (TrainConfig | None) – Optional TrainConfig to store alongside each checkpoint, useful for reproducing training runs.

Returns:

A callback compatible with train_ppo’s checkpoint_fn parameter.

Return type:

CheckpointCallback

Example

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

nnx_ppo.algorithms.checkpointing.load_checkpoint(path, networks, optimizer)

Load a checkpoint saved by make_checkpoint_fn().

The networks and optimizer arguments serve as structural templates: their architecture must match the checkpoint, but their current parameter values are irrelevant and will be overwritten in-place by the checkpoint values.

Parameters:

• path (str) – Path to the step checkpoint directory, e.g. /tmp/my_run/step_0000500000.

• networks (Any) – Network instance with the same architecture as the checkpoint. Weights are updated in-place.

• optimizer (Optimizer) – Optimizer instance with the same structure as the checkpoint. State is updated in-place.

Returns:

• "training_state" — restored TrainingState

• "step" — training step at which the checkpoint was saved (int)

• "config" — TrainConfig if one was stored, else None

Return type:

A dict with the following keys

Example

>>> networks = factories.make_mlp_actor_critic(...)
>>> training_state = ppo.new_training_state(env, networks, n_envs, seed)
>>> ckpt = load_checkpoint(
...     "/tmp/my_run/step_0000500000",
...     training_state.networks,
...     training_state.optimizer,
... )
>>> result = train_ppo(
...     env, networks, ckpt["config"],
...     initial_state=ckpt["training_state"],
... )

Callbacks

Callback helpers for training logging.

nnx_ppo.algorithms.callbacks.wandb_video_fn(key='eval_video', fps=30)

Create a video callback that logs to wandb.

Parameters:

• key (str) – The wandb log key for the video.

• fps (int) – Frames per second for the video.

Returns:

A callback function compatible with train_ppo’s video_fn parameter.

Return type:

Callable[[VideoData], None]

Example

>>> result = train_ppo(
...     env, networks,
...     video_fn=wandb_video_fn(fps=50),
... )

Metrics

Metrics computation and logging utilities for PPO training.

nnx_ppo.algorithms.metrics.compute_metrics(loss_metrics, rollout_data, logging_level, percentile_levels=None)

Compute training metrics from loss metrics and rollout data.

Parameters:

• loss_metrics (dict[str, PyTree]) – Dictionary of loss values from ppo_loss.

• rollout_data (Transition) – Transition data from the rollout.

• logging_level (LoggingLevel) – Which metrics to include.

• percentile_levels (tuple[int, ...] | None) – Percentiles to compute (e.g., (0, 25, 50, 75, 100)). If None, uses mean/std instead.

Returns:

Dictionary of computed metrics.

Return type:

dict[str, Any]

nnx_ppo.algorithms.metrics.log_weight_stats(metrics, networks, percentile_levels=None)

Log weight statistics over all of the network’s parameters.

Parameters:

• metrics (dict[str, Any]) – Dictionary to add metrics to (mutated in place).

• networks (StatefulModule) – The PPO network to extract weights from.

• percentile_levels (tuple[int, ...] | None) – Percentiles to compute. If None, uses mean/std.