foxhunt/crates/ml-ppo/src/action_masking.rs

//! Action Masking for PPO
//!
//! Prevents the agent from selecting actions that would violate position limits.
//! This is adapted from DQN's action masking but works with PPO's policy network.
//!
//! # Key Differences from DQN:
//! - **DQN**: Masks Q-values after forward pass, selects argmax from unmasked actions
//! - **PPO**: Masks logits BEFORE softmax, samples from masked probability distribution
//!
//! # 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)
//!
//! # 45-Action Factored Space (5 exposure x 3 order x 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

/// Creates an action mask based on current position and position limits.
///
/// # Arguments
/// * `current_position` - Current portfolio position
/// * `max_position` - Maximum allowed position magnitude (typically 2.0)
/// * `num_actions` - Number of actions in the action space (3 or 45)
///
/// # Returns
/// Boolean mask where `true` = action is valid, `false` = action is invalid
pub fn create_action_mask(current_position: f64, max_position: f64, num_actions: usize) -> Vec<bool> {
    let mut mask = vec![true; num_actions];

    if num_actions == 3 {
        if current_position + 1.0 > max_position {
            mask[0] = false;
        }
        if current_position - 1.0 < -max_position {
            mask[1] = false;
        }
    } else if num_actions == 63 {
        for action_idx in 0..num_actions {
            let exposure_idx = action_idx / 9;
            match exposure_idx {
                // Long actions (LongSmall=4, LongHalf=5, LongFull=6)
                4 | 5 | 6 => {
                    if current_position >= max_position {
                        if let Some(m) = mask.get_mut(action_idx) {
                            *m = false;
                        }
                    }
                }
                // Short actions (ShortSmall=0, ShortHalf=1, ShortFull=2)
                0 | 1 | 2 => {
                    if current_position <= -max_position {
                        if let Some(m) = mask.get_mut(action_idx) {
                            *m = false;
                        }
                    }
                }
                _ => {} // Flat=3
            }
        }
    } else {
        // Unsupported action space size, all actions remain valid
    }

    mask
}

/// Applies action mask to logits by setting masked actions to -1e9.
///
/// Invalid actions (mask[i] = false) have their logits set to -1e9 (effectively -inf).
/// This ensures that after softmax, masked actions have probability ~0.
///
/// # Arguments
/// * `logits` - Raw logits from policy network (shape: [`num_actions`])
/// * `mask` - Boolean mask where false = invalid action
///
/// # Returns
/// Masked logits where invalid actions have logit = -1e9
pub fn apply_mask_to_logits_vec(logits: &[f32], mask: &[bool]) -> Vec<f32> {
    logits
        .iter()
        .zip(mask.iter())
        .map(|(&logit, &valid)| if valid { logit } else { -1e9 })
        .collect()
}

#[cfg(test)]
#[allow(
    clippy::bool_assert_comparison,
    clippy::wildcard_enum_match_arm
)]
mod tests {
    use super::*;

    #[test]
    fn test_create_action_mask_flat_position() {
        let mask = create_action_mask(0.0, 2.0, 3);
        assert_eq!(mask, vec![true, true, true]);
    }

    #[test]
    fn test_create_action_mask_at_max_position() {
        let mask = create_action_mask(2.0, 2.0, 3);
        assert_eq!(mask[0], false);
        assert_eq!(mask[1], true);
        assert_eq!(mask[2], true);
    }

    #[test]
    fn test_create_action_mask_at_min_position() {
        let mask = create_action_mask(-2.0, 2.0, 3);
        assert_eq!(mask[0], true);
        assert_eq!(mask[1], false);
        assert_eq!(mask[2], true);
    }

    #[test]
    fn test_apply_mask_to_logits() {
        let logits = vec![0.5_f32, 0.3, 0.2];
        let mask = vec![false, true, true];

        let masked = apply_mask_to_logits_vec(&logits, &mask);

        assert!(masked[0] < -1e8);
        assert_eq!(masked[1], 0.3);
        assert_eq!(masked[2], 0.2);
    }

    #[test]
    fn test_apply_mask_all_valid() {
        let logits = vec![0.5_f32, 0.3, 0.2];
        let mask = vec![true, true, true];

        let masked = apply_mask_to_logits_vec(&logits, &mask);

        assert_eq!(masked[0], 0.5);
        assert_eq!(masked[1], 0.3);
        assert_eq!(masked[2], 0.2);
    }

    #[test]
    fn test_create_action_mask_63_flat_position() {
        let mask = create_action_mask(0.0, 2.0, 63);
        assert_eq!(mask.len(), 63);
        assert!(mask.iter().all(|&v| v));
    }

    #[test]
    fn test_create_action_mask_63_at_max_position() {
        let mask = create_action_mask(2.0, 2.0, 63);
        assert_eq!(mask.len(), 63);
        // Short actions (exposure 0,1,2) should be valid
        for i in 0..27 {
            assert_eq!(mask.get(i).copied().unwrap_or(false), true);
        }
        // Flat (exposure 3) should be valid
        for i in 27..36 {
            assert_eq!(mask.get(i).copied().unwrap_or(false), true);
        }
        // Long actions (exposure 4,5,6) should be masked
        for i in 36..63 {
            assert_eq!(mask.get(i).copied().unwrap_or(true), false);
        }
    }

    #[test]
    fn test_create_action_mask_63_at_min_position() {
        let mask = create_action_mask(-2.0, 2.0, 63);
        // Short actions (exposure 0,1,2) should be masked
        for i in 0..27 {
            assert_eq!(mask.get(i).copied().unwrap_or(true), false);
        }
        // Flat (exposure 3) should be valid
        for i in 27..36 {
            assert_eq!(mask.get(i).copied().unwrap_or(false), true);
        }
        // Long actions (exposure 4,5,6) should be valid
        for i in 36..63 {
            assert_eq!(mask.get(i).copied().unwrap_or(false), true);
        }
    }

    #[test]
    fn test_create_action_mask_63_partial_position() {
        let mask = create_action_mask(1.0, 2.0, 63);
        assert!(mask.iter().all(|&v| v));
    }

    #[test]
    fn test_apply_mask_all_invalid() {
        let logits = vec![0.5_f32, 0.3, 0.2];
        let mask = vec![false, false, false];

        let masked = apply_mask_to_logits_vec(&logits, &mask);

        assert!(masked[0] < -1e8);
        assert!(masked[1] < -1e8);
        assert!(masked[2] < -1e8);
    }

    #[test]
    fn test_63_action_masking_canonical_layout() {
        // 63-action layout: 7 exposure levels × 9 sub-actions each
        // exposure_idx = action_idx / 9:
        //   0=ShortSmall, 1=ShortHalf, 2=ShortFull, 3=Flat, 4=LongSmall, 5=LongHalf, 6=LongFull
        let max_position = 5.0;

        // At max position: Long actions (idx 4,5,6) should be masked
        let mask = create_action_mask(max_position, max_position, 63);
        for idx in 0..63 {
            let exposure_idx = idx / 9;
            match exposure_idx {
                4 | 5 | 6 => {
                    assert!(
                        !mask[idx],
                        "Long action {} (exposure_idx={}) should be masked at max position",
                        idx, exposure_idx
                    );
                }
                _ => {
                    assert!(
                        mask[idx],
                        "Non-long action {} (exposure_idx={}) should NOT be masked at max position",
                        idx, exposure_idx
                    );
                }
            }
        }

        // At min position: Short actions (idx 0,1,2) should be masked
        let mask = create_action_mask(-max_position, max_position, 63);
        for idx in 0..63 {
            let exposure_idx = idx / 9;
            match exposure_idx {
                0 | 1 | 2 => {
                    assert!(
                        !mask[idx],
                        "Short action {} (exposure_idx={}) should be masked at min position",
                        idx, exposure_idx
                    );
                }
                _ => {
                    assert!(
                        mask[idx],
                        "Non-short action {} (exposure_idx={}) should NOT be masked at min position",
                        idx, exposure_idx
                    );
                }
            }
        }
    }
}