diff --git a/crates/ml-dqn/src/seg_tree_kernel.cu b/crates/ml-dqn/src/seg_tree_kernel.cu
index d49a3789d..8649fb3a5 100644
--- a/crates/ml-dqn/src/seg_tree_kernel.cu
+++ b/crates/ml-dqn/src/seg_tree_kernel.cu
@@ -56,8 +56,7 @@ extern "C" __global__ void seg_tree_update(
     // ancestor. atomicAdd is commutative+associative, so concurrent threads
     // produce correct sums without synchronization barriers.
     int leaf = capacity + (int)idx;
-    float old_pa = tree[leaf];
-    tree[leaf] = pa;
+    float old_pa = atomicExch(&tree[leaf], pa);
     float delta = pa - old_pa;
 
     int node = leaf >> 1;
@@ -90,8 +89,7 @@ extern "C" __global__ void seg_tree_insert(
 
     // Write leaf and propagate delta to root via atomicAdd (race-free).
     int leaf = capacity + (int)idx;
-    float old_pa = tree[leaf];
-    tree[leaf] = pa;
+    float old_pa = atomicExch(&tree[leaf], pa);
     float delta = pa - old_pa;
 
     int node = leaf >> 1;
diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
index 454220567..a93d57c87 100644
--- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
+++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
@@ -636,7 +636,8 @@ pub struct GpuDqnTrainer {
     ///   [0]=loss, [1]=mse_loss, [2]=grad_norm_sq          (per-step, from graph_adam)
     ///   [3]=avg_max_q, [4]=q_min, [5]=q_max, [6]=q_mean, [7]=q_var  (every 50 steps)
     ///   [8]=causal_mean_sens                               (every N steps)
-    ///   [9..16]=reserved                                   (nan_flags, future use)
+    ///   [9]=ensemble_diversity_loss                        (epoch boundary)
+    ///   [10..16]=reserved                                  (future use)
     /// Pinned memory enables true async cuMemcpyDtoHAsync without CPU blocking.
     readback_pinned: *mut f32,
     /// Whether there's an in-flight readback to collect