foxhunt/docs/superpowers/specs/2026-04-20-bitonic-sort-integration-design.md

Bitonic Sort Integration — 4 Uses for Existing Sort Infrastructure

Goal

Repurpose the bitonic sort (built for rank normalization, now dormant) for 4 training improvements that all benefit from fast GPU sorting (~14ms for 4M elements).

Use 1: C51 Adaptive Atom Warm-Start from Reward Quantiles

Problem: C51 atoms start uniformly spaced [-50, +50]. After PopArt, rewards land in ~[-3, +3]. The learned step_atom_positions SGD takes many epochs to move atoms into the useful range.

Fix: Each epoch, sort collected rewards via bitonic sort, compute quantile positions, write as warm-start for atom_positions_buf. The existing SGD optimizer refines from there.

sorted_rewards = bitonic_sort(rewards_buf, N)
for branch in 0..4:
    for atom in 0..num_atoms:
        atom_positions[branch * num_atoms + atom] = sorted_rewards[atom * N / num_atoms]

Integration:

• After shape_rewards / experience collection in training_loop.rs
• Call collector.sort_rewards_for_quantiles() → returns sorted buffer
• Call trainer.warm_start_atom_positions(&sorted_rewards, N) → writes to atom_positions_buf
• Existing step_atom_positions() SGD continues refining from the warm start

Cost: ~14ms per epoch (amortized, one sort).

Use 2: Robust PopArt Statistics (Median/IQR)

Problem: PopArt uses running mean/var (Welford). With bimodal rewards (many ±0.1 micro + few ±5.0 trade exits), the mean is pulled toward the dense micro mode. Variance captures both modes but is dominated by the rare large values. This biases normalization.

Fix: Sort rewards, compute median and IQR (interquartile range = Q75 - Q25). Use these instead of mean/var:

median = sorted_rewards[N/2]
q25 = sorted_rewards[N/4]
q75 = sorted_rewards[3*N/4]
iqr = q75 - q25
normalized_r = (r - median) / max(iqr, 1e-6)

Integration:

• Add popart_normalize_robust_kernel that uses median/IQR instead of Welford
• OR: compute quantiles from sorted rewards on GPU, read 3 scalars (median, Q25, Q75) to pinned memory, use in the existing PopArt kernel

Trade-off: Median/IQR is non-differentiable. But PopArt normalization doesn't need gradients — it's a preprocessing step.

Cost: ~14ms per epoch (reuses the same sort from Use 1).

Use 3: Experience Curriculum (Sort by TD-Error)

Problem: PER samples proportional to TD-error |δ|. With dense micro-rewards, most TD-errors are tiny (model quickly learns to predict ±0.1). PER rarely samples these. A deterministic curriculum presents experiences in difficulty order.

Fix: Sort experiences by |TD-error| ascending. Train on the first K (easy) for warmup, then the full sorted buffer:

sorted_indices = bitonic_sort_by_key(|td_errors|, indices, N)
// Epoch warmup (first 10 epochs): train on easy half (low TD-error)
// After warmup: train on full buffer, but sample K% from hard tail

Integration:

• After PER sampling, sort the batch by |td_error|
• Use a curriculum_fraction config (0.0 = all easy, 1.0 = all hard, 0.5 = balanced)
• Gradually increase curriculum_fraction over epochs (start easy, get harder)

Cost: Sort B=16384 batch = ~0.05ms (tiny). Not 4M — just the batch.

Use 4: Top-K Priority Sampling in PER

Problem: PER proportional sampling with alpha=0.3 is nearly uniform. Top-K would guarantee the highest-priority experiences are always in the batch.

Fix: Sort PER priorities, take top K as guaranteed inclusions, fill remaining batch slots with proportional sampling:

sorted_priorities = bitonic_sort(priorities, buffer_size)
top_k_indices = sorted_priorities[-K:]  // highest K priorities
remaining = proportional_sample(buffer_size - K, batch_size - K)
batch = concat(top_k_indices, remaining)

Config: per_topk_fraction = 0.2 → 20% of batch is guaranteed top-K, 80% proportional.

Integration:

• In gpu_replay_buffer.rs::sample_proportional()
• Sort priorities once per epoch (or every N steps)
• Cache the top-K indices
• Mix into each batch

Cost: Sort 500K buffer = ~2ms. Once per epoch.

Implementation Priority

1. Use 1 (atom warm-start) — highest impact, simplest integration, uses existing atom_positions_buf
2. Use 2 (robust PopArt) — high impact, fixes known bimodal distribution issue
3. Use 3 (curriculum) — medium impact, interesting for dense reward regime
4. Use 4 (top-K PER) — medium impact, alternative to alpha tuning

Files Modified

File	Change
gpu_experience_collector.rs	sort_rewards_for_quantiles() method (reuses existing bitonic sort)
gpu_dqn_trainer.rs	warm_start_atom_positions(), robust PopArt kernel call
training_loop.rs	Wire atom warm-start + PopArt quantiles after experience collection
dqn_utility_kernels.cu	popart_normalize_robust kernel variant (median/IQR)
gpu_replay_buffer.rs	Top-K priority cache in sample_proportional()
experience_kernels.cu	Curriculum fraction in env_step or PER sampling
dqn-production.toml	per_topk_fraction, curriculum_warmup_epochs