diff --git a/config/training/dqn-smoketest.toml b/config/training/dqn-smoketest.toml index c03875057..ec7924ab0 100644 --- a/config/training/dqn-smoketest.toml +++ b/config/training/dqn-smoketest.toml @@ -37,7 +37,7 @@ count_bonus_coefficient = 0.1 q_gap_threshold = 0.05 [replay_buffer] -buffer_size = 1024 +buffer_size = 256 min_replay_size = 64 regime_replay_decay = 0.3 diff --git a/crates/ml-dqn/src/replay_buffer_type.rs b/crates/ml-dqn/src/replay_buffer_type.rs index 66304bd9d..a690001f1 100644 --- a/crates/ml-dqn/src/replay_buffer_type.rs +++ b/crates/ml-dqn/src/replay_buffer_type.rs @@ -176,10 +176,11 @@ impl ReplayBufferType { batch_size: usize, stream: &Arc, ) -> Result { - const MIN_GPU_CAPACITY: usize = 1024; + // Minimum capacity = batch_size (can't sample more than buffer holds) + let min_capacity = batch_size.max(1); let mut try_cap = capacity; - while try_cap >= MIN_GPU_CAPACITY { + while try_cap >= min_capacity { match Self::new_gpu_prioritized(try_cap, state_dim, alpha, beta, beta_max, beta_annealing_steps, max_memory_bytes, batch_size, stream) { Ok(buf) => { if try_cap < capacity { @@ -201,14 +202,17 @@ impl ReplayBufferType { try_cap, e, try_cap / 2 ); try_cap /= 2; - if try_cap < MIN_GPU_CAPACITY { + if try_cap < min_capacity { return Err(e); } } } } - Err(MLError::ModelError("GPU PER capacity below floor".into())) + Err(MLError::ModelError(format!( + "GPU PER allocation failed: capacity {} below minimum {} (batch_size)", + try_cap, min_capacity + ))) } /// Sample a batch from the buffer