Replay Buffer

A circular buffer of Transitions, used by the off-policy algorithms (DQN, SAC).

jaxnasium.algorithms.core.TransitionBuffer

A buffer for storing transitions and sampling contiguous sequences from them (or single transitions in the case of n_steps=1). Samples uniformly from valid sequence start positions in the buffer. The buffer is implemented as a circular buffer, where the oldest transitions are overwritten when the buffer is full.

Arguments: max_size: The maximum number of transitions the buffer can hold (across all envs). sample_batch_size: The number of sequences to sample from the buffer. data_sample: A sample Transition to initialize the buffer structure. vectorized_env: Whether the inserted transitions are vectorized environment rollouts. If True, inserted transitions are assumed to be of shape (insert_batch_size, num_envs, ...). data is stored as is, but at sampling time, the num_envs axis is collapsed into the batch dimension. n_steps: Length of each sampled sequence. 1 reproduces single-transition sampling with output shape (sample_batch_size, ...). For n_steps > 1, output shape is (sample_batch_size, n_steps, ...).

insert(transition: Transition) -> Self

Insert a batch of transitions into the buffer.

Possible vectorized environments are stored as is (to preserve sequences) but flattened into a single batch dimension at sampling time.

sample(key: PRNGKeyArray, with_replacement: bool = True, batch_size: int | None = None, with_next_obs: bool = False) -> Transition

Sample a batch of transitions from the buffer. Samples a batch of sequences of length n_steps when n_steps > 1, otherwise a batch of single transitions.

When n_steps=1, returns shape (sample_batch_size, ...). When n_steps > 1, returns shape (sample_batch_size, n_steps, ...).

When vectorized_env is True, the vectorized environment axis is collapsed into the batch dimension of the returned transitions.

Arguments: key: JAX PRNG key for sampling. with_replacement: Whether to sample with replacement. Batch size: Number of sequences to sample. Defaults to sample_batch_size. with_next_obs: sets the next_observation field of the returned Transitions to the observation of the next transition in the buffer. This only fires if the transition did not already have an explicit next_observation set. In case this fires, truncated transitions are not sampled such that bootstrapping is still correct.

Prioritized replay

jaxnasium.algorithms.core.PrioritizedTransitionBuffer

A circular buffer with Prioritized Experience Replay

Additional arguments: alpha: Exponent controlling how strongly priorities bias sampling. 0 recovers uniform sampling, 1 samples fully proportionally. beta: Exponent for the importance-sampling correction. May be annealed externally. eps: Small constant added to |TD-error| so no transition ever gets a zero probability of being sampled.

insert(transition: Transition) -> Self

Insert a transition into the buffer. Maximum priority is assigned to the new insertion such that it is sampled at least once.

sample(key: PRNGKeyArray, with_replacement: bool = False, batch_size: int | None = None, with_next_obs: bool = False) -> Transition

Sample a batch of transitions from the buffer. Samples a batch of sequences of length n_steps when n_steps > 1, otherwise a batch of single transitions.

When n_steps=1, returns shape (sample_batch_size, ...). When n_steps > 1, returns shape (sample_batch_size, n_steps, ...).

When vectorized_env is True, the vectorized environment axis is collapsed into the batch dimension of the returned transitions.

Alongside the Transition batch, this PER adds the PER_weights to the Transition batch and returns the indices of the sampled sequence starts as (Transition, flat_indices).

Arguments: key: JAX PRNG key for sampling. with_replacement: Whether to sample with replacement. Batch size: Number of sequences to sample. Defaults to sample_batch_size. with_next_obs: sets the next_observation field of the returned Transitions to the observation of the next transition in the buffer. This only fires if the transition did not already have an explicit next_observation set. In case this fires, truncated transitions are not sampled such that bootstrapping is still correct.

update_priorities(indices: Array, td_errors: Array) -> Self

Refresh the priorities of the given buffer indices from new TD errors.

Arguments: indices: Flat buffer start indices to update, i.e. the indices returned by sample. td_errors: TD errors (or any priority signal) for those indices.

update_beta(beta: float) -> Self

Update the beta parameter of the buffer.