From 63b2c6e72d7bca5ac06e2ea9de426bfec3d9994e Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 2 Mar 2026 14:53:13 +0100 Subject: [PATCH] fix(ml): align PPO action masking to canonical exposure*9+order*3+urgency layout MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Was using direction=idx/15 (3 groups of 15: Buy/Sell/Hold) — incompatible with DQN's FactoredAction encoding. Now uses exposure_idx=idx/9 (5 groups of 9: Short100/Short50/Flat/Long50/Long100) matching FactoredAction. Co-Authored-By: Claude Opus 4.6 --- crates/ml/src/ppo/action_masking.rs | 206 +++++++++++++++++++++------- 1 file changed, 160 insertions(+), 46 deletions(-) diff --git a/crates/ml/src/ppo/action_masking.rs b/crates/ml/src/ppo/action_masking.rs index 30153a3c0..e390cde3d 100644 --- a/crates/ml/src/ppo/action_masking.rs +++ b/crates/ml/src/ppo/action_masking.rs @@ -7,16 +7,20 @@ //! - **DQN**: Masks Q-values after forward pass, selects argmax from unmasked actions //! - **PPO**: Masks logits BEFORE softmax, samples from masked probability distribution //! -//! # Current Implementation (3 Actions): +//! # 3-Action Space: //! - Action 0: BUY (increase position by +1.0) //! - Action 1: SELL (decrease position by -1.0) //! - Action 2: HOLD (keep position unchanged) //! -//! # Future Expansion (Phase 3 - 45 Actions): -//! When expanding to DQN's 45-action factored space (5 exposure × 3 order × 3 urgency): -//! - Exposure levels: Short100 (-1.0), Short50 (-0.5), Flat (0.0), Long50 (+0.5), Long100 (+1.0) -//! - Masking logic remains the same, just more granular exposure levels -//! - Will need to adapt masking to handle factored action space structure +//! # 45-Action Factored Space (5 exposure × 3 order × 3 urgency): +//! Index = exposure * 9 + order * 3 + urgency +//! - exposure_idx = idx / 9 +//! - 0 = Short100 (indices 0..9) +//! - 1 = Short50 (indices 9..18) +//! - 2 = Flat (indices 18..27) +//! - 3 = Long50 (indices 27..36) +//! - 4 = Long100 (indices 36..45) +//! - Masking: Long50/Long100 masked at max position, Short100/Short50 masked at min position use candle_core::{Result, Tensor}; @@ -42,12 +46,12 @@ use candle_core::{Result, Tensor}; /// - Mask SELL if current_position - 1.0 < -max_position /// - HOLD is always valid /// -/// # Future Expansion Notes (45 Actions): -/// When expanding to factored action space: -/// - Action index → (exposure, order, urgency) -/// - Exposure determines position change: Short100 (-1.0), Short50 (-0.5), Flat (0.0), Long50 (+0.5), Long100 (+1.0) -/// - Masking based on target_exposure = action.exposure.target_exposure() -/// - Same logic: mask if |new_position| > max_position +/// # 45-Action Factored Space: +/// Index = exposure * 9 + order * 3 + urgency (matches `FactoredAction::from_index`) +/// - exposure_idx = idx / 9: 0=Short100, 1=Short50, 2=Flat, 3=Long50, 4=Long100 +/// - Long50/Long100 (indices 27..45) masked when current_position >= max_position +/// - Short100/Short50 (indices 0..18) masked when current_position <= -max_position +/// - Flat (indices 18..27) always valid /// /// # Example /// ```rust @@ -77,28 +81,29 @@ pub fn create_action_mask(current_position: f64, max_position: f64, num_actions: // HOLD (action 2) is always valid } else if num_actions == 45 { - // Factored 45-action space: (direction: 3) x (size: 5) x (urgency: 3) = 45 - // direction = action_idx / 15: 0=Buy, 1=Sell, 2=Hold + // Factored 45-action space: exposure(5) × order(3) × urgency(3) + // Index = exposure * 9 + order * 3 + urgency + // exposure: 0=Short100, 1=Short50, 2=Flat, 3=Long50, 4=Long100 for action_idx in 0..num_actions { - let direction = action_idx / 15; - match direction { - 0 => { - // Buy: mask if at or above max position + let exposure_idx = action_idx / 9; + match exposure_idx { + 3 | 4 => { + // Long50, Long100: mask if at or above max position if current_position >= max_position { if let Some(m) = mask.get_mut(action_idx) { *m = false; } } } - 1 => { - // Sell: mask if at or below negative max position + 0 | 1 => { + // Short100, Short50: mask if at or below negative max position if current_position <= -max_position { if let Some(m) = mask.get_mut(action_idx) { *m = false; } } } - // 2 = Hold: always valid, no masking needed + // 2 = Flat: always valid _ => {} } } @@ -222,33 +227,52 @@ mod tests { #[test] fn test_create_action_mask_45_at_max_position() { - // At max long: Buy actions (0..15) masked, Sell (15..30) and Hold (30..45) valid + // At max long: Long50 (27..36) and Long100 (36..45) masked + // Short100 (0..9), Short50 (9..18), Flat (18..27) valid let mask = create_action_mask(2.0, 2.0, 45); assert_eq!(mask.len(), 45); - // Buy actions (direction=0, indices 0..15) should be masked - for i in 0..15 { + // Short100 (indices 0..9) should be valid + for i in 0..9 { + assert_eq!( + mask.get(i).copied().unwrap_or(false), + true, + "Short100 action {} should be valid at max position", + i + ); + } + // Short50 (indices 9..18) should be valid + for i in 9..18 { + assert_eq!( + mask.get(i).copied().unwrap_or(false), + true, + "Short50 action {} should be valid at max position", + i + ); + } + // Flat (indices 18..27) should be valid + for i in 18..27 { + assert_eq!( + mask.get(i).copied().unwrap_or(false), + true, + "Flat action {} should be valid at max position", + i + ); + } + // Long50 (indices 27..36) should be masked + for i in 27..36 { assert_eq!( mask.get(i).copied().unwrap_or(true), false, - "Buy action {} should be masked at max position", + "Long50 action {} should be masked at max position", i ); } - // Sell actions (direction=1, indices 15..30) should be valid - for i in 15..30 { + // Long100 (indices 36..45) should be masked + for i in 36..45 { assert_eq!( - mask.get(i).copied().unwrap_or(false), - true, - "Sell action {} should be valid at max position", - i - ); - } - // Hold actions (direction=2, indices 30..45) should be valid - for i in 30..45 { - assert_eq!( - mask.get(i).copied().unwrap_or(false), - true, - "Hold action {} should be valid at max position", + mask.get(i).copied().unwrap_or(true), + false, + "Long100 action {} should be masked at max position", i ); } @@ -256,16 +280,53 @@ mod tests { #[test] fn test_create_action_mask_45_at_min_position() { - // At max short: Sell actions (15..30) masked, Buy (0..15) and Hold (30..45) valid + // At max short: Short100 (0..9) and Short50 (9..18) masked + // Flat (18..27), Long50 (27..36), Long100 (36..45) valid let mask = create_action_mask(-2.0, 2.0, 45); - for i in 0..15 { - assert_eq!(mask.get(i).copied().unwrap_or(false), true); + // Short100 (0..9) masked + for i in 0..9 { + assert_eq!( + mask.get(i).copied().unwrap_or(true), + false, + "Short100 action {} should be masked at min position", + i + ); } - for i in 15..30 { - assert_eq!(mask.get(i).copied().unwrap_or(true), false); + // Short50 (9..18) masked + for i in 9..18 { + assert_eq!( + mask.get(i).copied().unwrap_or(true), + false, + "Short50 action {} should be masked at min position", + i + ); } - for i in 30..45 { - assert_eq!(mask.get(i).copied().unwrap_or(false), true); + // Flat (18..27) valid + for i in 18..27 { + assert_eq!( + mask.get(i).copied().unwrap_or(false), + true, + "Flat action {} should be valid at min position", + i + ); + } + // Long50 (27..36) valid + for i in 27..36 { + assert_eq!( + mask.get(i).copied().unwrap_or(false), + true, + "Long50 action {} should be valid at min position", + i + ); + } + // Long100 (36..45) valid + for i in 36..45 { + assert_eq!( + mask.get(i).copied().unwrap_or(false), + true, + "Long100 action {} should be valid at min position", + i + ); } } @@ -289,4 +350,57 @@ mod tests { assert!(values[1] < -1e8); assert!(values[2] < -1e8); } + + #[test] + fn test_45_action_masking_canonical_layout() { + use crate::common::action::{ExposureLevel, FactoredAction}; + + let max_position = 5.0; + + // At max long position, all long actions should be masked + let mask = create_action_mask(max_position, max_position, 45); + for idx in 0..45 { + if let Ok(action) = FactoredAction::from_index(idx) { + match action.exposure { + ExposureLevel::Long50 | ExposureLevel::Long100 => { + assert!( + !mask.get(idx).copied().unwrap_or(true), + "Long action {} should be masked at max position", + idx + ); + } + _ => { + assert!( + mask.get(idx).copied().unwrap_or(false), + "Non-long action {} should NOT be masked at max position", + idx + ); + } + } + } + } + + // At max short position, all short actions should be masked + let mask = create_action_mask(-max_position, max_position, 45); + for idx in 0..45 { + if let Ok(action) = FactoredAction::from_index(idx) { + match action.exposure { + ExposureLevel::Short50 | ExposureLevel::Short100 => { + assert!( + !mask.get(idx).copied().unwrap_or(true), + "Short action {} should be masked at min position", + idx + ); + } + _ => { + assert!( + mask.get(idx).copied().unwrap_or(false), + "Non-short action {} should NOT be masked at min position", + idx + ); + } + } + } + } + } }