foxhunt/crates/ml-dqn/src/nstep_gpu_tests.rs

//! CPU-vs-GPU parity tests for n-step return accumulator.
//!
//! These tests verify the GPU `nstep_push_and_sum()` ring buffer logic
//! by running an equivalent pure-Rust implementation and checking
//! discounted return values.

#[cfg(test)]
#[allow(clippy::doc_markdown)]
mod tests {
    /// Pure-Rust port of the GPU nstep_push_and_sum() for parity testing.
    fn nstep_push_and_sum_cpu(
        ring: &mut [f32],
        ring_idx: &mut usize,
        ring_len: &mut usize,
        reward: f32,
        gamma: f32,
        n_steps: usize,
    ) -> f32 {
        ring[*ring_idx] = reward;
        *ring_idx = (*ring_idx + 1) % n_steps;
        if *ring_len < n_steps {
            *ring_len += 1;
        }

        let mut g = 0.0_f32;
        let mut discount = 1.0_f32;
        let len = *ring_len;
        let start = (*ring_idx + n_steps - len) % n_steps;
        for k in 0..len {
            let i = (start + k) % n_steps;
            g += discount * ring[i];
            discount *= gamma;
        }
        g
    }

    fn nstep_reset_cpu(ring: &mut [f32], ring_idx: &mut usize, ring_len: &mut usize) {
        *ring_idx = 0;
        *ring_len = 0;
        for v in ring.iter_mut() {
            *v = 0.0;
        }
    }

    #[test]
    fn test_nstep_1_is_identity() {
        // n_steps=1 should return the reward itself (standard TD)
        let mut ring = [0.0_f32; 8];
        let mut idx = 0_usize;
        let mut len = 0_usize;
        let gamma = 0.99_f32;

        let rewards = [0.5_f32, -0.3, 1.0, 0.0, -0.1];
        for &r in &rewards {
            let g = nstep_push_and_sum_cpu(&mut ring, &mut idx, &mut len, r, gamma, 1);
            assert!(
                (g - r).abs() < 1e-6,
                "n_step=1 should return reward: got {} expected {}",
                g,
                r
            );
        }
    }

    #[test]
    fn test_nstep_3_discounted_return() {
        let mut ring = [0.0_f32; 8];
        let mut idx = 0_usize;
        let mut len = 0_usize;
        let gamma = 0.99_f32;
        let n = 3_usize;

        // Push 3 rewards: [1.0, 2.0, 3.0]
        let g1 = nstep_push_and_sum_cpu(&mut ring, &mut idx, &mut len, 1.0, gamma, n);
        assert!((g1 - 1.0).abs() < 1e-6, "First push: {}", g1);

        let g2 = nstep_push_and_sum_cpu(&mut ring, &mut idx, &mut len, 2.0, gamma, n);
        // G = 1.0 + 0.99*2.0 = 2.98
        let expected2 = 1.0 + gamma * 2.0;
        assert!((g2 - expected2).abs() < 1e-5, "Second push: {} vs {}", g2, expected2);

        let g3 = nstep_push_and_sum_cpu(&mut ring, &mut idx, &mut len, 3.0, gamma, n);
        // G = 1.0 + 0.99*2.0 + 0.99^2*3.0 = 1.0 + 1.98 + 2.9403 = 5.9203
        let expected3 = 1.0 + gamma * 2.0 + gamma * gamma * 3.0;
        assert!((g3 - expected3).abs() < 1e-4, "Third push: {} vs {}", g3, expected3);

        // Push 4th: ring wraps, oldest (1.0) dropped
        let g4 = nstep_push_and_sum_cpu(&mut ring, &mut idx, &mut len, 4.0, gamma, n);
        // G = 2.0 + 0.99*3.0 + 0.99^2*4.0
        let expected4 = 2.0 + gamma * 3.0 + gamma * gamma * 4.0;
        assert!((g4 - expected4).abs() < 1e-4, "Fourth push: {} vs {}", g4, expected4);
    }

    #[test]
    fn test_nstep_reset_clears_state() {
        let mut ring = [0.0_f32; 8];
        let mut idx = 0_usize;
        let mut len = 0_usize;
        let gamma = 0.99_f32;

        // Fill some values
        nstep_push_and_sum_cpu(&mut ring, &mut idx, &mut len, 1.0, gamma, 3);
        nstep_push_and_sum_cpu(&mut ring, &mut idx, &mut len, 2.0, gamma, 3);

        // Reset
        nstep_reset_cpu(&mut ring, &mut idx, &mut len);

        // After reset, first push should return just the reward
        let g = nstep_push_and_sum_cpu(&mut ring, &mut idx, &mut len, 5.0, gamma, 3);
        assert!((g - 5.0).abs() < 1e-6, "After reset: {} vs 5.0", g);
    }

    #[test]
    fn test_nstep_5_accumulation() {
        let mut ring = [0.0_f32; 8];
        let mut idx = 0_usize;
        let mut len = 0_usize;
        let gamma = 0.95_f32;
        let n = 5_usize;

        // Push 5 rewards: all 1.0
        let mut last_g = 0.0;
        for i in 0..5 {
            last_g = nstep_push_and_sum_cpu(&mut ring, &mut idx, &mut len, 1.0, gamma, n);
            // Each step adds more terms to the sum
            assert!(last_g > 0.0, "Step {} should be positive", i);
        }

        // After 5 steps of reward=1.0: G = 1 + 0.95 + 0.95^2 + 0.95^3 + 0.95^4
        let expected = 1.0 + 0.95 + 0.95_f32.powi(2) + 0.95_f32.powi(3) + 0.95_f32.powi(4);
        assert!(
            (last_g - expected).abs() < 1e-4,
            "5-step constant reward: {} vs {}",
            last_g,
            expected
        );
    }
}