jaxdem.rl.trainers#

Interface for defining reinforcement learning model trainers.

Classes

`Trainer`(env, graphdef, graphstate, key, ...)	Base class for reinforcement learning trainers.
`TrajectoryData`(*, obs, action, value, ...)	Container for rollout data (single step or stacked across time).

class jaxdem.rl.trainers.PPOTrainer(env: Environment, graphdef: nnx.GraphDef[Any], graphstate: nnx.GraphState, key: ArrayLike, advantage_gamma: jax.Array, advantage_lambda: jax.Array, advantage_rho_clip: jax.Array, advantage_c_clip: jax.Array, drip_decay: jax.Array, ppo_clip_eps: jax.Array, ppo_value_coeff: jax.Array, ppo_entropy_coeff: jax.Array, importance_sampling_alpha: jax.Array, importance_sampling_beta: jax.Array, anneal_importance_sampling_beta: jax.Array, num_epochs: int, stop_at_epoch: int, num_steps_epoch: int, num_minibatches: int, minibatch_size: int, skip_frames: int)#

Bases: Trainer

Proximal Policy Optimization (PPO) trainer in PufferLib style.

This trainer implements the PPO algorithm with clipped surrogate objectives, value-function loss, entropy regularization, and prioritized experience replay (PER).

Loss function

Given a trajectory batch with actions \(a_t\), states \(s_t\), rewards \(r_t\), advantages \(A_t\), and old log-probabilities \(\log \pi_{\theta_\text{old}}(a_t \mid s_t)\), we define:

Probability ratio:

\[\rho_t(\theta) = \exp\big( \log \pi_\theta(a_t \mid s_t) - \log \pi_{\theta_\text{old}}(a_t \mid s_t) \big)\]
Clipped policy loss:

\[L^{\text{policy}}(\theta) = - \mathbb{E}_t \Big[ \min\big( \rho_t(\theta) A_t,\; \text{clip}(\rho_t(\theta), 1-\epsilon, 1+\epsilon) A_t \big) \Big]\]

where \(\epsilon\) is the PPO clipping parameter.
Value-function loss (with clipping):

\[L^{\text{value}}(\theta) = \tfrac{1}{2} \mathbb{E}_t \Big[ \max\big( (V_\theta(s_t) - R_t)^2,\; (\text{clip}(V_\theta(s_t), V_{\theta_\text{old}}(s_t) - \epsilon, V_{\theta_\text{old}}(s_t) + \epsilon) - R_t)^2 \big) \Big]\]

where \(R_t = A_t + V_{\theta_\text{old}}(s_t)\) are return targets.
Entropy bonus:

\[L^{\text{entropy}}(\theta) = \mathbb{E}_t \big[ \mathcal{H}[\pi_\theta(\cdot \mid s_t)] \big]\]

which encourages exploration.
Total loss:

\[L(\theta) = L^{\text{policy}}(\theta) + c_v L^{\text{value}}(\theta) - c_e L^{\text{entropy}}(\theta)\]

where \(c_v\) and \(c_e\) are coefficients for the value and entropy terms.

Prioritized Experience Replay (PER)

This trainer uses a prioritized categorical distribution over segments (environments x agents) to form minibatches. For each segment index \(i \in \{1,\dots,N\}\), we define a priority from the trajectory advantages:

\[\tilde{p}_i \;=\; \Big\| A_{\cdot,i} \Big\|_1^{\,\alpha} \quad\text{with}\quad \Big\| A_{\cdot,i} \Big\|_1 \;=\; \sum_{t=1}^{T} \big|A_{t,i}\big|,\]

where \(\alpha \ge 0\) (importance_sampling_alpha) controls the strength of prioritization. We then form a categorical sampling distribution

\[P(i) \;=\; \frac{\tilde{p}_i}{\sum_{k=1}^{N} \tilde{p}_k},\]

and sample indices \(\{i\}\) to create each minibatch (jax.random.choice() with probabilities \(P(i)\)). This mirrors Prioritized Experience Replay (PER), where \(\tilde{p}\) comes from the TD-error magnitude. Here we use the per-trajectory advantage magnitude as a proxy for learning progress. Recent large-scale self-play systems for autonomous driving use the same design. We use the absolute value of the advantage so that the sampler keeps both the best and the worst samples.

To correct sampling bias we apply PER-style importance weights (importance_sampling_beta with optional linear annealing):

\[w_i(\beta_t) \;=\; \Big(N \, P(i)\Big)^{-\beta_t}, \qquad \beta_t \in [0,1].\]

Classical PER often normalizes \(w_i\) by \(\max_j w_j\) to keep the scale bounded. This implementation omits that normalization and uses \(w_i\) directly. The minibatch advantages are standardized and reweighted with these IS weights before the PPO loss:

\[\hat{A}_{t,i} \;=\; w_i(\beta_t)\; \frac{A_{t,i} - \mu_{\text{mb}}(A)}{\sigma_{\text{mb}}(A)+\varepsilon}.\]

If importance_sampling_alpha = 0, we get uniform sampling. If importance_sampling_beta = 1 we get full PER correction.

Off-policy correction of advantages (V-trace)

The trainer recomputes the advantages on each minibatch iteration with fresh values and probability ratios. When a minibatch reuses a sample, the V-trace off-policy correction in Trainer.compute_advantages() recomputes the advantages. This matters because the policy changes during each minibatch update, which makes the rollout off-policy and the value stale.

Distributed Reward Information Processing (DRIP)

DRIP addresses the credit assignment problem in environments with sparse or delayed feedback. It runs a recursive backward pass over the trajectory that distributes terminal or delayed rewards back to the past causal states. The implementation uses an exponential smoothing filter bounded by episode terminations, so rewards do not bleed across episodes. - To activate: Set drip_decay to a value between (0.0, 1.0] (e.g., 0.8). - To deactivate: Set drip_decay to 0.0 (the default behavior).

— References

Schulman et al., Proximal Policy Optimization Algorithms, 2017.
Espeholt et al., IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures, ICML 2018.
Schulman et al., High-Dimensional Continuous Control Using Generalized Advantage Estimation, 2015/2016.
Schaul et al., Prioritized Experience Replay, ICLR 2016.
Cusumano-Towner et al., Robust Autonomy Emerges from Self-Play, ICML 2025.

drip_decay: jax.Array#: Decay factor \(\lambda_{DRIP}\) for Distributed Reward Information Processing (DRIP). Drips delayed/sparse rewards backward through time to assign credit to past actions. Set to 0.0 to disable (default).

ppo_clip_eps: jax.Array#: PPO clipping parameter \(\epsilon\) used for both the policy ratio clip and the value-function clip.

ppo_value_coeff: jax.Array#: Coefficient \(c_v\) scaling the value-function loss term in the total loss.

ppo_entropy_coeff: jax.Array#: Coefficient \(c_e\) scaling the entropy bonus (encourages exploration).

importance_sampling_alpha: jax.Array#: Prioritization strength \(\alpha \ge 0\) for minibatch sampling. Higher values put more probability mass on envs with larger advantages.

importance_sampling_beta: jax.Array#: Initial PER importance-weight exponent \(\beta \in [0,1]\) used in \(w_i(\beta) = (N P(i))^{-\beta}\). It compensates the sampling bias.

anneal_importance_sampling_beta: jax.Array#: If True, the trainer linearly anneals \(\beta\) toward 1 during training (more correction later in training).

num_epochs: int#: Number of PPO training epochs (outer loop count).

stop_at_epoch: int#: Stop after this epoch. Must satisfy 1 ≤ stop_at_epoch ≤ num_epochs.

num_steps_epoch: int#: Rollout horizon \(T\) per epoch. Total collected steps = \(N \times T\).

num_minibatches: int#: Number of minibatches per epoch used for PPO updates.

minibatch_size: int#: Minibatch size (number of env indices sampled per update). Typically \(N / \text{num\_minibatches}\).

skip_frames: int#: Number of frames to skip (repeat action) for each observation.

classmethod Create(env: Environment, model: Model, seed: int | None = None, key: ArrayLike | None = None, optimizer: Any = <function muon>, learning_rate: float = 0.01, anneal_learning_rate: bool = True, max_grad_norm: float = 1.5, accumulate_n_gradients: int = 1, ppo_clip_eps: float = 0.2, ppo_value_coeff: float = 2.0, ppo_entropy_coeff: float = 0.001, advantage_gamma: float = 0.99, advantage_lambda: float = 0.95, advantage_rho_clip: float = 1.0, advantage_c_clip: float = 1.0, drip_decay: float = 0.0, importance_sampling_alpha: float = 0.8, importance_sampling_beta: float = 0.2, anneal_importance_sampling_beta: bool = True, num_envs: int = 1024, num_steps_epoch: int = 64, num_minibatches: int = 4, minibatch_size: int | None = None, skip_frames: int = 0, num_epochs: int = 1000, total_timesteps: int | None = None, stop_at_epoch: int | None = None, clip_actions: bool = False, clip_range: tuple[float, float]=(-0.2, 0.2)) → Self[source]#

Construct a PPO trainer from an environment and a model.

Vectorizes the environment, builds the optimizer chain, and initializes the model carry. See the class-level field docstrings for parameter descriptions.

Parameters:

env (Environment) – A single (non-vectorised) environment instance.
model (Model) – An actor–critic model whose observation_space_size and action_space_size match env.
seed (int, optional) – Integer seed used for random number generation. Used only when key is not provided.
key (jax.Array, optional) – PRNG key. When provided, it takes precedence over seed (same rule as jaxdem.System.create()).

Returns:

Ready-to-train trainer instance.

Return type:

PPOTrainer

static train(tr: Trainer, verbose: bool = True, log: bool = True, directory: Path | str = 'runs', save_every: int = 2, start_epoch: int = 0, debug_overflow_checks: bool = False, **kwargs: Any) → PPOTrainer[source]#

Run the full PPO training loop.

Parameters:

tr (Trainer) – Trainer instance (will be cast to PPOTrainer).
verbose (bool) – If True, display a tqdm progress bar.
log (bool) – If True, write TensorBoard scalars to directory.
directory (Path | str) – Root directory for TensorBoard logs.
save_every (int) – Sync metrics and log every save_every epochs.
start_epoch (int) – Resume epoch counter (useful after checkpoint restore).
debug_overflow_checks (bool) – If True, check the collider overflow flag after every epoch. This forces a host synchronization per epoch, which defeats async dispatch, so it is off by default. The flag is always checked once at the end of training.

Returns:

Trainer with updated parameters after training.

Return type:

PPOTrainer

static loss_fn(model: Model, td: TrajectoryData, returns: jax.Array, advantage: jax.Array, ppo_clip_eps: jax.Array, ppo_value_coeff: jax.Array, ppo_entropy_coeff: jax.Array, initial_carry: Any | None = None) → tuple[jax.Array, dict[str, jax.Array]][source]#

Compute the clipped PPO loss for a minibatch.

Runs a forward pass through model and returns the composite loss (policy + value + entropy) together with diagnostic scalars.

Parameters:

model (Model) – Actor–critic model (called with sequence=True).
td (TrajectoryData) – Minibatch trajectory slice [T, M, ...].
returns (jax.Array) – Return targets, shape [T, M].
advantage (jax.Array) – Normalized, IS-weighted advantages, shape [T, M].
ppo_clip_eps (jax.Array) – Clipping parameter \(\epsilon\).
ppo_value_coeff (jax.Array) – Value-loss coefficient \(c_v\).
ppo_entropy_coeff (jax.Array) – Entropy-bonus coefficient \(c_e\).

Returns:

Scalar total loss and a dictionary of diagnostic metrics.

Return type:

tuple[jax.Array, dict[str, jax.Array]]

Execute one PPO epoch: rollout → advantage → minibatch updates.

Steps: 0. Reset done environments and split PRNG keys. 1. Collect a trajectory of length num_steps_epoch. 2. Flatten the agent axis and apply DRIP if enabled. 3. Compute PER priorities. 4. Scan over num_minibatches updates, each recomputing V-trace advantages and applying the clipped PPO loss.

Parameters:

tr (PPOTrainer) – Current trainer state.
epoch (ArrayLike) – Zero-based epoch index (used for \(\beta\) annealing).

Returns:

Updated trainer, full trajectory data, and epoch-averaged metrics.

Return type:

Tuple[PPOTrainer, TrajectoryData, dict[str, jax.Array]]

Modules

ppo_trainer

Implementation of the PPO algorithm.