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.TrajectoryData(*, obs: Array, action: Array, value: Array, log_prob: Array, ratio: Array, reward: Array, done: Array, terminated: Array, truncated: Array, agent_mask: Array, bootstrap_value: Array, latent_action: Array | None = None, latent_log_prob: Array | None = None)#

Bases: object

Container for rollout data (single step or stacked across time).

obs: Array#: Observations.

action: Array#: Actions sampled from the policy; inactive agent entries are zeroed.

value: Array#: Immutable behavior-policy value estimates \(V_b(s_t)\) used as the PPO value-clipping reference. Current learner values are computed separately.

log_prob: Array#: Behavior-policy log-probabilities \(\log \pi_b(a_t \mid s_t)\) at collection time.

ratio: Array#: Collection-time probability ratio (one). Retained for trajectory consumers; PPO computes current importance ratios locally without overwriting this field.

reward: Array#: Post-action rewards \(r_t\) during collection. PPO epoch output contains DRIP-transformed rewards when that option is enabled.

done: Array#: Episode-boundary flags (terminated | truncated).

terminated: Array#: Terminal flags; these suppress value bootstrapping.

truncated: Array#: Time-limit/external boundary flags; these retain bootstrapping unless the same transition is also terminal.

agent_mask: Array#: Boolean flags selecting active agents in padded environments.

bootstrap_value: Array#: Collection-time value of the final pre-reset observation for truncations.

latent_action: Array | None#: Pre-bijector sample for transformed continuous policies, otherwise None.

latent_log_prob: Array | None#

Behavior base-distribution log-probability of latent_action.

Ratios for a fixed bijector are evaluated in base coordinates to avoid inverting saturated actions. log_prob still records transformed density.

class jaxdem.rl.trainers.Trainer(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)#

Bases: Factory, ABC

Base class for reinforcement learning trainers.

This class holds the environment and model state (Flax NNX GraphDef/State). It provides rollout utilities (step(), trajectory_rollout()) and a general advantage method (compute_advantages()). Subclasses must implement algorithm-specific training logic in epoch().

Example

To define a custom trainer, inherit from Trainer and implement its abstract methods:

>>> @Trainer.register("myCustomTrainer")
... @jax.tree_util.register_dataclass
... @dataclass(slots=True)
... class MyCustomTrainer(Trainer):
...     ...

env: Environment#: Environment object.

graphdef: nnx.GraphDef[Any]#: Static graph definition of the model/optimizer.

graphstate: ModelState#: Mutable state (parameters, optimizer state, RNGs, etc.).

key: ArrayLike#: PRNGKey used to sample actions and for other stochastic operations.

advantage_gamma: jax.Array#: Discount factor \(\gamma \in [0, 1]\).

advantage_lambda: jax.Array#: Generalized Advantage Estimation parameter \(\lambda \in [0, 1]\).

advantage_rho_clip: jax.Array#: V-trace \(\bar{\rho}\) (importance weight clip for the TD term).

advantage_c_clip: jax.Array#: V-trace \(\bar{c}\) (importance weight clip for the recursion/trace term).

property model: Any[source]#: Return the live model rebuilt from (graphdef, graphstate).

static step(env: Environment, graphdef: nnx.GraphDef[Any], graphstate: ModelState, key: jax.Array, skip_frames: int = 0) → tuple[tuple[Environment, ModelState, jax.Array], TrajectoryData][source]#

Take one environment step and record a single-step trajectory. Repeats the action for skip_frames extra frames when set.

Parameters:

env (Environment) – The (vectorized) environment to step.
graphdef (nnx.GraphDef) – Python part of the nnx model.
graphstate (nnx.State) – State of the nnx model.
key (jax.Array) – Jax random key.
skip_frames (int) – Number of additional frames to repeat the action.

Returns:

Tuple[Tuple[Environment, nnx.State, jax.Array], TrajectoryData] – Updated state and the new single-step trajectory. Trajectory data is shaped (N_envs, N_agents, …).
Finished environments are reset before the next transition. Frame
skipping stops each environment at its first boundary. Truncations
retain a value estimate of their final observation for bootstrapping.
advance_action calls checkpoint(env, action) once with the
masked policy action before its physics loop. Action wrappers apply
the same transformation at the checkpoint and at each physics step.
Reward and truncation bootstrap observations read the live endpoint
before reset, with no endpoint checkpoint.
Intermediate rewards are not accumulated.

static trajectory_rollout(env: Environment, graphdef: nnx.GraphDef[Any], graphstate: ModelState, key: jax.Array, num_steps_epoch: int, unroll: int = 8, skip_frames: int = 0) → tuple[Environment, ModelState, jax.Array, TrajectoryData][source]#

Roll out \(T = \text{num\_steps\_epoch}\) policy decisions using jax.lax.scan().

Parameters:

env (Environment) – The (vectorized) environment to roll out.
graphdef (nnx.GraphDef) – Python part of the nnx model.
graphstate (nnx.State) – State of the nnx model.
key (jax.Array) – Jax random key.
num_steps_epoch (int) – Number of policy decisions to roll out per agent slot.
unroll (int) – Number of loop iterations to unroll for compilation speed.
skip_frames (int) – Number of additional physics frames requested per policy action.

Returns:

The final environment, graph state, and PRNG key, plus a TrajectoryData instance whose fields are stacked along time (leading dimension \(T = \text{num\_steps\_epoch}\)).

Return type:

Tuple[Environment, nnx.State, jax.Array, TrajectoryData]

static compute_advantages(value: Array, reward: Array, ratio: Array, done: Array, advantage_rho_clip: Array, advantage_c_clip: Array, advantage_gamma: Array, advantage_lambda: Array, last_value: Array | None = None, terminated: Array | None = None, truncated: Array | None = None, bootstrap_value: Array | None = None, agent_mask: Array | None = None, unroll: int = 8) → tuple[Array, Array][source]#

Return detached targets and GAE/V-trace-style advantages.

All transition arrays have leading time dimension \(T\). For active-agent indicator \(m_t\), terminal flag \(z_t\), and episode-boundary flag \(d_t\), the recurrence is

\[\begin{split}\widehat\rho_t &= \min(\mathrm{ratio}_t, \bar\rho), \qquad \widehat c_t = \min(\mathrm{ratio}_t, \bar c),\\ \delta_t &= m_t\widehat\rho_t [r_t+\gamma(1-z_t)B_t-V_t],\\ A_t &= \delta_t+\gamma\lambda m_t(1-d_t)\widehat c_t A_{t+1}, \qquad A_T=0,\\ R_t &= V_t + m_t A_t.\end{split}\]

Here \(B_t\) is bootstrap_value[t] at a truncation, last_value at the final non-truncated transition, and value[t+1] otherwise. A terminal flag suppresses bootstrapping even if the same transition is also truncated. Inputs must be finite, including inactive padding; multiplication by zero does not mask NaNs.

Unit ratios and unit clipping caps give ordinary GAE. Otherwise this is a clipped-importance lambda trace. PPO uses the resulting detached \(A_t\) in its clipped actor objective, rather than the separate policy-gradient advantage from the IMPALA algorithm.

Parameters:

value (jax.Array) – Per-transition values, post-action rewards, and importance ratios. PPO supplies its current learner values and, with vtrace=True, current-to-behavior policy ratios. Otherwise it supplies unit ratios and unit caps, regardless of the configured trace caps.
reward (jax.Array) – Per-transition values, post-action rewards, and importance ratios. PPO supplies its current learner values and, with vtrace=True, current-to-behavior policy ratios. Otherwise it supplies unit ratios and unit caps, regardless of the configured trace caps.
ratio (jax.Array) – Per-transition values, post-action rewards, and importance ratios. PPO supplies its current learner values and, with vtrace=True, current-to-behavior policy ratios. Otherwise it supplies unit ratios and unit caps, regardless of the configured trace caps.
done (jax.Array) – Episode boundaries, consistent with terminated | truncated. These cut the trace after the current transition.
advantage_rho_clip (jax.Array) – Upper caps on the TD residual and trace importance weights.
advantage_c_clip (jax.Array) – Upper caps on the TD residual and trace importance weights.
advantage_gamma (jax.Array) – Discount and trace decay per recorded policy transition.
advantage_lambda (jax.Array) – Discount and trace decay per recorded policy transition.
last_value (jax.Array | None) – Bootstrap on the post-rollout observation, without a time axis. If omitted, value[-1] is reused; this is generally biased at a continuing horizon boundary. PPO provides a collection-time estimate and holds it fixed during minibatch updates.
terminated (jax.Array | None) – Supply both to distinguish terminal from truncated transitions. If terminated is omitted it defaults to done; if truncated is omitted it defaults to false. Passing only truncated therefore does not enable time-limit bootstrapping.
truncated (jax.Array | None) – Supply both to distinguish terminal from truncated transitions. If terminated is omitted it defaults to done; if truncated is omitted it defaults to false. Passing only truncated therefore does not enable time-limit bootstrapping.
bootstrap_value (jax.Array | None) – Values of final pre-reset observations at truncations, zero if omitted. PPO records these during collection.
agent_mask (jax.Array | None) – Active-agent mask, all true if omitted. Inactive entries have zero advantage and return target equal to their input value.
unroll (int) – Unroll factor for the reverse scan.

Returns:

(returns, advantages), both with gradients stopped and the same shape as value.

Return type:

Tuple[jax.Array, jax.Array]

References

Schulman et al., High-Dimensional Continuous Control Using Generalized Advantage Estimation, 2015/2016.
Espeholt et al., IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures, 2018.

Run one training epoch.

Subclasses implement this method with their algorithm-specific logic.

abstractmethod static train(tr: Trainer, *args: Any, **kwargs: Any) → Any[source]#

Training loop.

Subclasses implement this method with their algorithm-specific logic.

Modules

ppo_trainer

Implementation of the PPO algorithm.