fix(ml): align PPO action masking to canonical exposure*9+order*3+urgency layout

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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-02 14:53:13 +01:00
parent 512535076f
commit 63b2c6e72d

View File

@@ -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
);
}
}
}
}
}
}