================================================================================ PRIORITIZED EXPERIENCE REPLAY - DATA FLOW DIAGRAM ================================================================================ ┌─────────────────────────────────────────────────────────────────────────────┐ │ DQN TRAINING LOOP │ └─────────────────────────────────────────────────────────────────────────────┘ ┌────────────────────┐ │ 1. Store │ │ Experience │ │ (state, action, │ │ reward, next, │ │ done) │ └──────┬─────────────┘ │ v ┌────────────────────────────────────────────────────────┐ │ PrioritizedReplayBuffer::push() │ │ • Initial priority = max_priority (ensures sampling) │ │ • Segment tree update: O(log n) │ └──────┬─────────────────────────────────────────────────┘ │ v ┌────────────────────┐ │ 2. Sample Batch │ │ (batch_size=32) │ └──────┬─────────────┘ │ v ┌──────────────────────────────────────────────────────────────┐ │ PrioritizedReplayBuffer::sample() │ │ ┌────────────────────────────────────────────────────────┐ │ │ │ For each sample i: │ │ │ │ │ │ │ │ 1. Get priority: p_i from segment tree │ │ │ │ │ │ │ │ 2. Calculate probability: │ │ │ │ P(i) = p_i / Σ(p_j) │ │ │ │ │ │ │ │ 3. Calculate current beta (with annealing): │ │ │ │ progress = step / beta_annealing_steps │ │ │ │ β = β_start + (β_max - β_start) × progress │ │ │ │ β ∈ [0.4, 1.0] │ │ │ │ │ │ │ │ 4. Calculate IS weight: │ │ │ │ w_raw = (N × P(i))^(-β) │ │ │ │ w_i = w_raw / max(w_j) ← normalization │ │ │ └────────────────────────────────────────────────────────┘ │ │ │ │ Returns: (experiences, weights, indices) │ └──────┬───────────────────────────────────────────────────────┘ │ v ┌────────────────────┐ │ 3. Forward Pass │ │ • Main network │ │ • Target network │ └──────┬─────────────┘ │ v ┌──────────────────────────────────────────────────────────────┐ │ WorkingDQN::train_step() │ │ ┌────────────────────────────────────────────────────────┐ │ │ │ 1. Compute TD errors: │ │ │ │ δ_i = Q(s,a) - (r + γ max Q'(s',a')) │ │ │ │ │ │ │ │ 2. Apply IS weights to loss: │ │ │ │ │ │ │ │ Standard Loss (MSE/Huber): │ │ │ │ weighted_diff = (Q - target) × w_i │ │ │ │ loss = mean(weighted_diff²) │ │ │ │ │ │ │ │ Distributional Loss (C51): │ │ │ │ per_sample_loss = -Σ target_i × log(pred_i) │ │ │ │ weighted_loss = per_sample_loss × w_i │ │ │ │ loss = mean(weighted_loss) │ │ │ └────────────────────────────────────────────────────────┘ │ └──────┬───────────────────────────────────────────────────────┘ │ v ┌────────────────────┐ │ 4. Backward Pass │ │ • Gradients │ │ • Optimizer step │ └──────┬─────────────┘ │ v ┌──────────────────────────────────────────────────────────────┐ │ 5. Update Priorities │ │ ┌────────────────────────────────────────────────────────┐ │ │ │ For each sampled index i: │ │ │ │ │ │ │ │ 1. Get TD error: δ_i │ │ │ │ │ │ │ │ 2. Calculate new priority: │ │ │ │ p_i = |δ_i| + ε (ε = 1e-6) │ │ │ │ │ │ │ │ 3. Update segment tree: O(log n) │ │ │ │ memory.update_priorities(indices, priorities) │ │ │ └────────────────────────────────────────────────────────┘ │ └──────┬───────────────────────────────────────────────────────┘ │ v ┌────────────────────┐ │ 6. Step Beta │ │ memory.step() │ │ training_step++ │ └────────────────────┘ ================================================================================ SEGMENT TREE STRUCTURE ================================================================================ Example for capacity = 8: Tree array indices: [0] [1] [2] [3] [4] [5] [6] [7] [8] [9] [10][11][12][13][14][15] x ROOT L0 R0 L1 R1 L2 R2 EXP EXP EXP EXP EXP EXP EXP EXP ↑ ↑ ↑ ↑ ↑ ↑ ↑ ↑ ↑ │ └───┴───┴───┴───┴───┴───┴───┘ │ Leaf nodes (experiences 0-7) │ Sum of all priorities Internal nodes store cumulative sums: - tree[1] = total priority sum - tree[2] = sum of left subtree (experiences 0-3) - tree[3] = sum of right subtree (experiences 4-7) - tree[i] = tree[2i] + tree[2i+1] Sampling algorithm (O(log n)): 1. Generate random value v ∈ [0, tree[1]] 2. Start at root (idx=1) 3. While not at leaf: - If v <= tree[left_child]: go left - Else: subtract tree[left_child] from v, go right 4. Return (idx - capacity) as experience index Update algorithm (O(log n)): 1. Set tree[idx + capacity] = new_priority 2. Propagate up: tree[parent] = tree[left] + tree[right] 3. Repeat until root ================================================================================ BETA ANNEALING SCHEDULE ================================================================================ Training Step │ Progress │ Beta │ IS Correction Strength ──────────────┼──────────┼──────────┼──────────────────────── 0 │ 0.0% │ 0.40 │ Minimal (exploration) 50,000 │ 10.0% │ 0.46 │ Growing 100,000 │ 20.0% │ 0.52 │ ↓ 150,000 │ 30.0% │ 0.58 │ ↓ 200,000 │ 40.0% │ 0.64 │ ↓ 250,000 │ 50.0% │ 0.70 │ ↓ 300,000 │ 60.0% │ 0.76 │ ↓ 350,000 │ 70.0% │ 0.82 │ ↓ 400,000 │ 80.0% │ 0.88 │ ↓ 450,000 │ 90.0% │ 0.94 │ ↓ 500,000 │ 100.0% │ 1.00 │ Full correction (convergence) Formula: β = β_start + (β_max - β_start) × min(1.0, step / annealing_steps) ================================================================================ IMPORTANCE SAMPLING WEIGHT CALCULATION ================================================================================ Given: - N = buffer size (e.g., 100,000) - P(i) = sampling probability for experience i - β = current beta value (0.4 → 1.0) Step-by-step calculation: 1. Raw weight: w_raw(i) = (N × P(i))^(-β) 2. Find maximum weight: P_min = min(P(j)) for all j w_max = (N × P_min)^(-β) 3. Normalize weight: w_i = w_raw(i) / w_max 4. Clamp to reasonable range: w_i = min(w_i, 10.0) ← prevents extreme values Properties: - w_i ∈ [0, 1] after normalization - Higher priority → higher P(i) → lower w_i (compensates for bias) - Lower priority → lower P(i) → higher w_i (upweights rare samples) - β=0 → w_i=1 (no correction, pure prioritization) - β=1 → full correction (unbiased gradient estimates) ================================================================================ PRIORITY UPDATE EXAMPLES ================================================================================ Example 1: High TD error (important transition) TD error: δ = 5.0 Priority: p = |5.0| + 1e-6 = 5.000001 Result: High sampling probability in next batch Example 2: Low TD error (well-learned transition) TD error: δ = 0.01 Priority: p = |0.01| + 1e-6 = 0.010001 Result: Low sampling probability Example 3: New experience (no TD error yet) TD error: N/A Priority: p = max_priority (e.g., 10.0) Result: Guaranteed to be sampled at least once ================================================================================ MEMORY EFFICIENCY ================================================================================ For 1M capacity buffer: Component Memory Usage ───────────────────────────────────────── Segment tree 8 MB (2 × 1M × 4 bytes) Experience buffer Varies (depends on state size) Atomic counters 32 bytes RNG state ~100 bytes ───────────────────────────────────────── Total PER overhead ~8 MB + experiences Comparison to uniform buffer: - Uniform: experiences only - PER: experiences + 8 MB overhead - Overhead: ~0.8% for typical state sizes ================================================================================