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:
@@ -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
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user