OakEmu-RL v0.1.0

Appearance

RolloutBuffer ​

python
RolloutBuffer(rollout_length: int, num_envs: int, observation_shape: tuple[int, ...], observation_dtype: np.dtype) -> None

Fixed-size rollout buffer for PPO.

Pre-allocates numpy arrays for efficient storage. After a rollout is complete, computes GAE advantages and provides minibatch iteration for PPO updates.

Initialize the RolloutBuffer.

Parameters:

  • rollout_length (int) – Number of steps per environment per rollout.
  • num_envs (int) – Number of parallel environments.
  • observation_shape (tuple[int, ...]) – Shape of a single observation.
  • observation_dtype (dtype) – Dtype of observations.

Source

Properties ​

is_full ​

python
is_full: bool

Check if the rollout buffer is full.

Source

memory_bytes ​

python
memory_bytes: int

Get the total pre-allocated CPU memory of this buffer in bytes.

Returns:

  • int – Memory usage in bytes.

Source

memory_gb ​

python
memory_gb: float

Get the total pre-allocated CPU memory of this buffer in GiB.

Returns:

  • float – Memory usage in GiB.

Source

num_envs ​

python
num_envs: int

Get the number of environments.

Source

observation_dtype ​

python
observation_dtype: np.dtype

Get the dtype of observations stored in the rollout buffer.

Source

observation_shape ​

python
observation_shape: tuple[int, ...]

Get the shape of observations stored in the rollout buffer.

Source

rollout_length ​

python
rollout_length: int

Get the rollout length.

Source

step ​

python
step: int

Get the current step index within the rollout.

Source

total_size ​

python
total_size: int

Get the total number of transitions (rollout_length * num_envs).

Source

trainable_size ​

python
trainable_size: int

Get the number of rollout steps available to PPO updates.

Source

Methods ​

add ​

python
add(observations: np.ndarray, actions: np.ndarray, log_probs: np.ndarray, values: np.ndarray, rewards: np.ndarray, dones: np.ndarray, trainable: np.ndarray | None = None) -> None

Add a step of data from all environments.

Parameters:

  • observations (ndarray) – Observations with shape (num_envs, *obs_shape).
  • actions (ndarray) – Actions with shape (num_envs,).
  • log_probs (ndarray) – Log probabilities with shape (num_envs,).
  • values (ndarray) – Value estimates with shape (num_envs,).
  • rewards (ndarray) – Rewards with shape (num_envs,).
  • dones (ndarray) – Done flags with shape (num_envs,).
  • trainable (ndarray | None) – Optional mask indicating which transitions should be used for PPO learning, for example random-controlled steps are stored for bookkeeping but marked non-trainable so PPO stays on-policy.

Source

compute_gae ​

python
compute_gae(last_values: np.ndarray, last_dones: np.ndarray, gamma: float, gae_lambda: float) -> None

Compute Generalized Advantage Estimation (GAE).

Delegates to GAE for the actual computation.

Parameters:

  • last_values (ndarray) – Value estimates for the state after the last step, shape (num_envs,).
  • last_dones (ndarray) – Done flags for the state after the last step, shape (num_envs,).
  • gamma (float) – Discount factor.
  • gae_lambda (float) – GAE lambda parameter.

Source

estimate_memory_pressure ​

python
estimate_memory_pressure(num_minibatches: int) -> list[RolloutBufferMemoryDiagnostic]

Estimate CPU and GPU memory pressure for this rollout buffer.

Parameters:

  • num_minibatches (int) – Number of PPO minibatches per learner update.

Returns:

  • list[RolloutBufferMemoryDiagnostic] – Diagnostics describing projected memory pressure.

Raises:

  • ValueError – If num_minibatches is not positive.

Source

get_minibatches ​

python
get_minibatches(num_minibatches: int, device: torch.device) -> list[RolloutMinibatch]

Create shuffled minibatches from the rollout data.

Observations are uploaded to the device once per rollout (cached until the buffer changes) and minibatch rows are gathered on-device, so repeated epochs cost no further observation transfers. Scalar tensors (actions, log_probs, values, advantages, returns) are gathered to their valid rows on the CPU and uploaded whole once per call; minibatch slices are then indexed on-device.

Parameters:

  • num_minibatches (int) – Number of minibatches to split the data into.
  • device (device) – The device to move tensors to.

Returns:

  • list[RolloutMinibatch] – A list of RolloutMinibatch instances.

Source

reset ​

python
reset() -> None

Reset the buffer for a new rollout.

Source

trainable_diagnostics ​

python
trainable_diagnostics() -> RolloutBufferTrainableDiagnostics | None

Return flattened trainable rollout arrays for learner diagnostics.

Returns:

  • RolloutBufferTrainableDiagnostics | None – Flattened trainable values, returns, and advantages, or None when no trainable rollout items are currently available.

Source