Transitions

jaxnasium.algorithms.core.Transition

Container for (possibly batches of) transitions Comes with functionality of creating minibatches and some utilities for multi-agent reinforcement learning.

structure: PyTreeDef property

Returns the top-level structure of the transition objects (using reward as a reference). This is either PyTreeDef() for single agents or PyTreeDef((, x num_agents)) for multi-agent environments. usefull for unflattening Transition.flat.properties back to the original structure.

replace(**updates)

normalize(normalizer: Normalizer)

Normalizes the observation and rewards in the transition based on the given normalizer.

view_transposed: PyTree[Transition] property

For single-agent settings, this will do nothing and return the original transition.

For multi-agent settings: The original transition is a Transition of PyTrees e.g. Transition(observation={a1: ..., a2: ...}, action={a1: ..., a2: ...}, ...) The transposed transition is a PyTree of Transitions e.g. {a1: Transition(observation=..., action=..., ...), a2: Transition(observation=..., action=..., ...), ...} This is useful for multi-agent environments where we want to have a single Transition object per agent.

from_transposed(transposed: PyTree[Transition]) -> Transition classmethod

scan(fn, init, *, reverse=False, unroll=1, **scan_kwargs) -> Any

Scan an arbitrary function over the time axis of this transition batch.

The user writes fn as if operating on single-agent (plain array) data. In multi-agent settings the transition is automatically transposed into per-agent Transitions, scanned individually, and the results are merged back.

Arguments:

• fn: (carry, step: Transition) -> (carry, output) – written for a single agent.
• init: Initial carry value. If the transition is multi-agent and init's pytree structure already matches the agent structure, each agent gets its own init leaf. If any of init's first-level children match the agent structure, these are forwarded to the respective agent. Otherwise init is broadcast (replicated) to every agent.
• reverse, unroll, **scan_kwargs: forwarded to jax.lax.scan.

Returns: (final_carry, outputs) where in multi-agent mode every output has been merged back into the original per-agent pytree structure.

make_minibatches(key: PRNGKeyArray, n_minibatches: int, n_epochs: int = 1, n_batch_axis: int = 1) -> Transition

Creates shuffled minibatches from the transition. Returns a copy of the transition with each leaf reshaped to (num_minibatches, ...),

This function first flattens the transition over the leading n_batch_axis. This is useful if your data hasn't been flattened yet and may be structured as (rollout_length, num_envs, ...), where num_envs is the number of parallel environments.

If n_epochs > 1, it will create n_epochs copies of the minibatches. and stack these such that there is a single leading axis to scan over for training.

If the batch size is not divisible by the number of minibatches, the remainder is truncated so that each minibatch has the same size.

Arguments: - key: JAX PRNG key for randomization. - num_minibatches: Number of minibatches to create. - n_epochs: Number of copies the minibatches should be stacked. - n_batch_axis: Number of leading batch axes to flatten over. Default is 1 (already flattened).

Utilities

jaxnasium.algorithms.core.n_step_to_cumulative_single_step(transition: Transition, n_step: int, gamma: float) -> Transition

Converts an n-step transition into a single step transition with the cumulative reward. terminated, truncated, and next_observation are set to the boundary index of the first terminated or truncated step.

n_step_axis is assumed to be the leading axis of the transition.

Operates on single transitions, use jax.vmap to apply to a batch of transitions.

Scanning over transitions

jaxnasium.algorithms.core.scan_transitions(f: Callable[[Carry, Transition], tuple[Carry, Y]], init: Carry, xs: Transition, reverse: bool = False, unroll: int | bool = 1, *, key: PRNGKeyArray | None = None, num_epochs: int = 1, num_minibatches: int = 1) -> tuple[Carry, Y]

A wrapper around jax.lax.scan that works with Transition PyTrees.

Parameters:

Name	Type	Description	Default
f	Callable[[Carry, Transition], tuple[Carry, Y]]	A function of the form (carry, x) -> (carry, y)	required
init	Carry	The initial carry value	required
xs	Transition	A Transition PyTree containing the inputs to scan over	required
reverse	bool	Whether to scan in reverse order	False
unroll	int | bool	Unroll factor for the scan	1
key	PRNGKeyArray | None	A PRNG key for shuffling the data when using epochs or minibatches	None
num_epochs	int	Number of epochs in the Transition	1
num_minibatches	int	Number of minibatches in the Transition	1

Returns:

Type	Description
tuple[Carry, Y]	A tuple of (final carry, stacked outputs)

jaxnasium.algorithms.core.scan_minibatch_epoch(f: Callable[[Carry, Transition], tuple[Carry, Y]], init: Carry, xs_batch: Transition, *, minibatch_rng: PRNGKeyArray, num_epochs: int, num_minibatches: int, unroll: int | bool | tuple[int | bool, int | bool] = (False, False))

A simple double jax.lax.scan loop over num_epochs x num_minibatches.

Each epoch reshuffles the xs_batch into num_minibatches. Follows the same pattern as jax.lax.scan, but has a nested loop.

Arguments:

• f: A function of the form (carry, xs) -> (carry, ys).
• init: The initial carry value.
• xs_batch: The batch of inputs to scan over.
• key: A PRNG key for shuffling the data when using epochs or minibatches.
• num_epochs: Number of epochs in the batch.
• num_minibatches: Number of minibatches in the batch.
• unroll: Unroll for both scans. Can be a tuple, where the the first element is the outer (epoch) scan and the second element is the inner (minibatch) scan.