RolloutBufferMemoryDiagnostics

RolloutBufferMemoryDiagnostics(memory_bytes: int, observation_shape: tuple[int, ...], observation_dtype: np.dtype, total_size: int) -> None

Estimates CPU and GPU memory pressure for a rollout buffer's pre-allocated arrays.

Initialize the diagnostics helper.

Parameters:

• memory_bytes (int) – The rollout buffer's total pre-allocated CPU memory, in bytes.
• observation_shape (tuple[int, ...]) – Shape of a single observation.
• observation_dtype (dtype) – Dtype of observations.
• total_size (int) – Total number of transitions (rollout_length * num_envs).

Source

Methods

estimate_memory_pressure

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