diff --git a/crates/ml/src/trainers/dqn/smoke_tests/reward_v8.rs b/crates/ml/src/trainers/dqn/smoke_tests/reward_v8.rs index 42c5893f6..49c549672 100644 --- a/crates/ml/src/trainers/dqn/smoke_tests/reward_v8.rs +++ b/crates/ml/src/trainers/dqn/smoke_tests/reward_v8.rs @@ -118,4 +118,56 @@ mod tests { assert!(norm_mean.abs() < 1e-10, "Mean should be 0, got {norm_mean}"); assert!((norm_var - 1.0).abs() < 0.01, "Variance should be 1, got {norm_var}"); } + + // ── 4-Branch Architecture Tests ───────────────────────────────────── + + #[test] + fn test_4branch_action_encoding() { + // All 81 actions encode/decode correctly + for dir in 0..3_usize { + for mag in 0..3_usize { + for order in 0..3_usize { + for urgency in 0..3_usize { + let action = dir * 27 + mag * 9 + order * 3 + urgency; + assert_eq!(action / 27, dir); + assert_eq!((action % 27) / 9, mag); + assert_eq!((action % 9) / 3, order); + assert_eq!(action % 3, urgency); + assert!(action < 81); + } + } + } + } + } + + #[test] + fn test_4branch_position_mapping() { + let max_pos = 2.0_f64; + // direction: 0=Short(-1), 1=Flat(0), 2=Long(+1) + // magnitude: 0=Small(0.25), 1=Medium(0.5), 2=Full(1.0) + // Short × Full = -2.0 + assert!(((- 1.0_f64) * 1.0 * max_pos - (-2.0)).abs() < 1e-10); + // Flat × anything = 0 + assert!((0.0_f64 * 0.25 * max_pos).abs() < 1e-10); + assert!((0.0_f64 * 1.0 * max_pos).abs() < 1e-10); + // Long × Quarter = +0.5 + assert!((1.0_f64 * 0.25 * max_pos - 0.5).abs() < 1e-10); + // Long × Half = +1.0 + assert!((1.0_f64 * 0.5 * max_pos - 1.0).abs() < 1e-10); + // Long × Full = +2.0 + assert!((1.0_f64 * 1.0 * max_pos - 2.0).abs() < 1e-10); + // Short × Small = -0.5 + assert!(((- 1.0_f64) * 0.25 * max_pos - (-0.5)).abs() < 1e-10); + } + + #[test] + fn test_flat_shortcut() { + // Flat direction (idx=1) × any magnitude = 0 position + for mag in 0..3_usize { + let dir_fraction = 0.0_f64; // Flat + let mag_scale = match mag { 0 => 0.25, 1 => 0.5, _ => 1.0 }; + assert!((dir_fraction * mag_scale).abs() < 1e-10, + "Flat × mag={mag} must be zero"); + } + } }