jaxdem.rl.trainers.ppo_trainer#

Implementation of the PPO algorithm.

Classes

PPOTrainer(env, graphdef, graphstate, key, ...)

Proximal Policy Optimization (PPO) trainer in PufferLib style.

class jaxdem.rl.trainers.ppo_trainer.PPOTrainer(env: Environment, graphdef: nnx.GraphDef[Any], graphstate: ModelState, 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, vtrace: jax.Array, num_epochs: int, stop_at_epoch: int, num_steps_epoch: int, num_minibatches: int, minibatch_size: int, skip_frames: int, caps_temporal_coeff: float = 0.0, caps_spatial_coeff: float = 0.0, caps_noise_std: jax.Array = <factory>)#

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 sequential minibatches of fixed-horizon agent segments. A segment may contain episode boundaries; it need not contain a complete episode.

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 and \(A_t\) is detached. Throughout these losses, \(\mathbb{E}_t[f_t]\) means \(\sum_t m_t f_t / \max(1, \sum_t m_t)\) over all time/agent entries in the minibatch, with active-agent indicator \(m_t\).
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 = \operatorname{stop\_gradient}(A_t + V_\theta(s_t))\) are detached return targets. The clipping reference remains the immutable rollout value.
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) + \lambda_T L_T(\theta) + \lambda_S L_S(\theta)\]

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

Optional CAPS regularization

Both coefficients default to zero, removing CAPS computations from the compiled update. We use a squared-Euclidean variant of CAPS:

\[L_T = \mathbb{E}\!\left[\|u_\theta(o_{t+1},h_{t+1}) - u_\theta(o_t,h_t)\|_2^2\right],\qquad L_S = \mathbb{E}\!\left[\|u_\theta(o_t+\epsilon,h_t) - u_\theta(o_t,h_t)\|_2^2\right], \quad \epsilon\sim\mathcal{N}(0,\operatorname{diag}(\sigma^2)).\]

Continuous outputs are the Gaussian mean transformed through the action bijector; categorical outputs are probability vectors. No exploration samples or environment rewards enter these penalties. Distances are in action units and are not automatically normalized. Temporal pairs use consecutive policy decisions from the same active agent and episode; rollout-end pairs are omitted. Spatial comparisons share clean incoming recurrent state. Both branches remain differentiable.

References

Mysore, S., Mabsout, B., Mancuso, R., and Saenko, K. (2021). Regularizing Action Policies for Smooth Control with Reinforcement Learning. ICRA. https://arxiv.org/abs/2012.06644. The paper uses unsquared Euclidean distances; this implementation uses squared distances with finite gradients at identical outputs.

Minibatch updates

Transitions are selected in contiguous agent-major blocks, cycling through the rollout when num_minibatches exceeds one pass. Full fixed-horizon sequences provide context, with a loss mask for partial sequences. Each update runs the current model before computing detached advantages and returns. Policy advantages are used directly, without normalization or sampling weights. Rollout values and behavior log-probabilities remain fixed for PPO clipping.

With vtrace=False (default), targets use ordinary GAE. With vtrace=True, a V-trace-style lambda trace uses the clipped current importance ratios in Trainer.compute_advantages. This still uses the PPO actor objective, not the IMPALA actor objective. Final-horizon and truncation bootstraps retain their collection-time estimates; episode boundary handling is independent of this choice.

Distributed Reward Information Processing (DRIP)

Before target calculation, rewards are replaced by the backward sum

\[\widetilde r_t = r_t + d(1-\mathrm{done}_t)\widetilde r_{t+1}, \qquad \widetilde r_T = 0.\]

Here \(d\) is drip_decay and done includes termination and truncation. This is an unnormalized sum of future rewards within the current rollout and episode. It changes reward scale and the learning objective; it does not conserve total reward or remove horizon dependence. Use \(0 < d \leq 1\) to enable it; 0.0 leaves rewards unchanged. Rewards are not clipped or normalized by this trainer.

— 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.

drip_decay: jax.Array#: Decay factor \(\lambda_{DRIP}\) for Distributed Reward Information Processing (DRIP). Adds decayed future rewards backward within each rollout and episode. 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).

vtrace: jax.Array#: Use current policy ratios in advantage estimation; otherwise use GAE.

num_epochs: int#: Planned number of rollout-and-update iterations, also used by the LR schedule. This is not the number of replay passes over one rollout.

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

num_steps_epoch: int#: Rollout horizon \(T\) in policy decisions per outer iteration. The rollout contains \(ST\) slots, where \(S=\text{num\_envs}\times\text{max\_num\_agents}\), including padding.

num_minibatches: int#: Number of sequential minibatch iterations per rollout. With an explicit minibatch size, iterations beyond a complete pass replay from the start. Inactive-only blocks do not advance optimizer state or accumulation.

minibatch_size: int#: Number of transition slots selected per minibatch, including inactive slots. Must be between num_steps_epoch and num_envs * max_num_agents * num_steps_epoch. Selection wraps at the rollout end. Full sequences are evaluated for recurrent context and targets; only selected transitions contribute to the loss. If omitted, defaults to rollout size divided exactly by num_minibatches.

skip_frames: int#: Number of additional physics frames requested for each policy action. State advances by at most 1 + skip_frames frames, ending at a boundary; intermediate physics-frame rewards are not accumulated.

caps_temporal_coeff: float#

changing it recompiles updates.

Type:: Nonnegative temporal CAPS weight. Static

caps_spatial_coeff: float#: Nonnegative spatial CAPS weight. Zero skips perturbations and extra forwards.

caps_noise_std: jax.Array#: Spatial noise standard deviation in observation units, scalar or per feature.

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, vtrace: bool = False, drip_decay: float = 0.0, num_envs: int = 64, 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), caps_temporal_coeff: float = 0.0, caps_spatial_coeff: float = 0.0, caps_noise_std: ArrayLike = 0.05) → Self[source]#

Construct a PPO trainer from an environment and a model.

Vectorizes the environment, builds the optimizer chain, and initializes the model carry. When enabled, learning-rate annealing is cosine decay over the planned optimizer updates in num_epochs. With no inactive-only blocks it remains constant within an epoch; skipped empty blocks also pause the optimizer’s schedule counter. Gradient accumulation averages raw minibatch gradients before clipping and the stateful optimizer. See the class-level field docstrings for parameter descriptions. Minibatch size counts transition slots; set minibatch_size explicitly to control replay independently of num_minibatches. The former PER constructor arguments were removed.

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).
caps_temporal_coeff (float) – Nonnegative CAPS loss weights, disabled by default. These are static configuration: changing either recompiles training.
caps_spatial_coeff (float) – Nonnegative CAPS loss weights, disabled by default. These are static configuration: changing either recompiles training.
caps_noise_std (ArrayLike) – Nonnegative finite Gaussian perturbation standard deviation, scalar or shape (observation_space_size,). Expressed in the model’s input units; use zero for features that must not vary. See the class docstring for the CAPS equations and paper citation.

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) – Positive integer interval for syncing metrics and logging iterations when verbose output or logging is enabled. Does not save checkpoints.
start_epoch (int) – Resume epoch counter for logging and rollout numbering. Exact learning-rate and momentum continuation also requires the restored trainer graphstate; changing this label alone does not restore optimizer state. Exact trajectory continuation also requires the environment, PRNG key, and recurrent carry. Must lie in [0, stop_at_epoch]. At stop_at_epoch, returns the trainer unchanged without opening a writer or running an epoch.
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, ppo_clip_eps: jax.Array, ppo_value_coeff: jax.Array, ppo_entropy_coeff: jax.Array, advantage_gamma: jax.Array, advantage_lambda: jax.Array, advantage_rho_clip: jax.Array, advantage_c_clip: jax.Array, last_value: jax.Array, vtrace: jax.Array, initial_carry: Any | None = None, loss_mask: jax.Array | None = None, caps_temporal_coeff: float = 0.0, caps_spatial_coeff: float = 0.0, caps_noise_std: ArrayLike = 0.05, caps_key: jax.Array | None = None) → tuple[jax.Array, dict[str, jax.Array]][source]#

Compute the clipped PPO loss for a minibatch.

Evaluate the current policy, form detached GAE/V-trace targets, and compute the composite loss. td.value and td.log_prob are frozen behavior-policy data. This function does not mutate the trajectory. Boundary bootstrap estimates come from the rollout. Transformed continuous policies use stored latent actions and base log-probabilities for the ratio; this assumes the bijector is fixed during replay.

Parameters:

model (Model) – Actor–critic model (called with sequence=True).
td (TrajectoryData) – Minibatch trajectory slice [T, M, ...].
loss_mask (jax.Array, optional) – Selected transitions, shape [T, M]. Only these contribute to losses and metrics; all active context still participates in RNN evaluation and GAE/V-trace target calculation.
last_value (jax.Array) – Rollout-end bootstrap values, shape [M].
vtrace (jax.Array) – Whether to use current importance ratios in target calculation.
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\).
caps_temporal_coeff (float) – Static weights of the squared-distance CAPS losses. Both zero preserves the original forward pass without extra RNG use.
caps_spatial_coeff (float) – Static weights of the squared-distance CAPS losses. Both zero preserves the original forward pass without extra RNG use.
caps_noise_std (ArrayLike) – Scalar or per-feature perturbation standard deviation.
caps_key (jax.Array, optional) – Required when spatial CAPS is enabled; fresh for each update. For the method and citation, see Mysore et al. (ICRA 2021), https://arxiv.org/abs/2012.06644, and the class docstring.

Returns:

Scalar total loss and diagnostics. ratio, value, target_values, and advantages are detached [T, M] arrays; the other entries are scalar minibatch statistics.

Return type:

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

Collect a rollout and train sequential minibatches with fresh targets.

Steps: 0. Save initial recurrent carry (resets occur inside each rollout step). 1. Collect a trajectory of length num_steps_epoch. 2. Flatten the agent axis and apply DRIP if enabled. 3. Visit contiguous segment blocks, wrapping for additional passes. 4. Run the current model, calculate targets, and update parameters. Rollout values and log-probabilities remain unchanged.

Parameters:

tr (PPOTrainer) – Current trainer state.
epoch (ArrayLike) – Zero-based outer iteration index (retained for caller compatibility).

Returns:

Updated trainer, trajectory data shaped [T, S, ...] with DRIP-transformed rewards, and metrics averaged equally over nonempty minibatches. An all-empty rollout yields zero metrics. Parameter updates occur according to gradient accumulation; empty minibatches do not advance the optimizer.

Return type:

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