foxhunt/WAVE1_A5_IMPLEMENTATION_PLAN.md

Wave 1 Agent A5: Factored Actions Training Integration

Implementation Plan

Phase 1: Trainer Modifications (ml/src/trainers/dqn.rs)

Changes Required:

1. Import Factored Types (lines 20-30)

// Add conditional imports
#[cfg(feature = "factored-actions")]
use crate::dqn::{FactoredAction, FactoredQNetwork, FactoredQNetworkConfig};

// Keep existing TradingAction for non-factored builds
#[cfg(not(feature = "factored-actions"))]
use crate::dqn::{Experience, TradingAction, TradingState};

#[cfg(feature = "factored-actions")]
use crate::dqn::{Experience, TradingState};

2. Struct Fields (around line 403)

pub struct DQNTrainer {
    #[cfg(feature = "factored-actions")]
    /// Flag to indicate if factored actions are enabled (runtime toggle)
    use_factored_actions: bool,

    #[cfg(not(feature = "factored-actions"))]
    _use_factored_actions: bool,  // Placeholder

    // ... existing fields ...

    /// Recent actions (type changes with feature)
    #[cfg(not(feature = "factored-actions"))]
    recent_actions: VecDeque<TradingAction>,

    #[cfg(feature = "factored-actions")]
    recent_actions: VecDeque<u8>,  // Store action indices (0-44)
}

3. Constructor (around line 458)

pub fn new_with_reward_system(
    hyperparams: DQNHyperparameters,
    reward_system: RewardSystem,
    #[cfg(feature = "factored-actions")]
    use_factored_actions: bool,
) -> Result<Self> {
    // ... existing validation ...

    let config = WorkingDQNConfig {
        state_dim: 128,

        #[cfg(feature = "factored-actions")]
        num_actions: if use_factored_actions { 45 } else { 3 },

        #[cfg(not(feature = "factored-actions"))]
        num_actions: 3,

        // ... rest of config ...
    };

    // ... after agent creation ...

    Ok(Self {
        #[cfg(feature = "factored-actions")]
        use_factored_actions,

        #[cfg(not(feature = "factored-actions"))]
        _use_factored_actions: false,

        // ... existing fields ...

        recent_actions: VecDeque::with_capacity(100),
    })
}

4. Action Selection (around line 2158)

async fn select_action(&self, state: &TradingState) -> Result<TradingAction> {
    let agent = self.agent.read().await;

    #[cfg(feature = "factored-actions")]
    if self.use_factored_actions {
        // Get action index (0-44)
        let action_idx = agent.select_action(&state.feature_vector)?;

        // Convert to FactoredAction
        let factored = FactoredAction::from_index(action_idx as usize)
            .map_err(|e| anyhow::anyhow!("Invalid factored action index: {}", e))?;

        // Store index in recent_actions
        self.recent_actions.push_back(action_idx);
        if self.recent_actions.len() > 100 {
            self.recent_actions.pop_front();
        }

        // Map to TradingAction based on exposure
        return Ok(match factored.exposure {
            ExposureLevel::Short100 | ExposureLevel::Short50 => TradingAction::Sell,
            ExposureLevel::Flat => TradingAction::Hold,
            ExposureLevel::Long50 | ExposureLevel::Long100 => TradingAction::Buy,
        });
    }

    // Original 3-action selection
    let action_idx = agent.select_action(&state.feature_vector)?;
    Ok(TradingAction::from_int(action_idx as u8)
        .ok_or_else(|| anyhow::anyhow!("Invalid action index: {}", action_idx))?)
}

5. Experience Storage (around line 2320)

async fn store_experience(&self, experience: Experience) -> Result<()> {
    let agent = self.agent.read().await;

    #[cfg(feature = "factored-actions")]
    if self.use_factored_actions {
        // Convert TradingAction back to factored action index
        // This is approximate - we lose order type and urgency info
        let action_idx = match experience.action {
            TradingAction::Buy => 27,    // Long50, Market, Normal (index 27)
            TradingAction::Sell => 9,    // Short50, Market, Normal (index 9)
            TradingAction::Hold => 19,   // Flat, Market, Normal (index 19)
        };

        let mut factored_exp = experience.clone();
        // Store the factored index (loss of granularity acceptable for now)
        agent.memory.lock().unwrap().push(factored_exp);
        return Ok(());
    }

    // Original storage
    agent.memory.lock().unwrap().push(experience);
    Ok(())
}

6. Transaction Cost Integration (in reward calculation, around line 827)

#[cfg(feature = "factored-actions")]
if self.use_factored_actions {
    // Extract transaction cost from factored action
    let action_idx = /* get from experience */;
    let factored = FactoredAction::from_index(action_idx as usize)?;
    let tx_cost = factored.transaction_cost();

    // Adjust P&L by transaction cost
    let adjusted_pnl = pnl * (1.0 - tx_cost);
    // Use adjusted_pnl in reward calculation
}

7. Position Masking (around line 2276, in epsilon_greedy_action)

#[cfg(feature = "factored-actions")]
if self.use_factored_actions {
    let current_position = self.portfolio_tracker.get_position_pct();

    // Get Q-values for all 45 actions
    let q_values = agent.forward(&state)?;

    // Apply position masking (prevent exceeding ±100%)
    let masked_q = apply_position_mask(&q_values, current_position)?;

    // Select action from masked Q-values
    let action_idx = if epsilon_greedy {
        sample_random_action()
    } else {
        masked_q.argmax(1)?
    };

    return Ok(action_idx);
}

Phase 2: CLI Flags (ml/examples/train_dqn.rs)

Add new flags:

#[derive(Debug, Parser)]
struct Opts {
    // ... existing fields ...

    /// Enable factored action space (45 actions: 5 exposure × 3 order × 3 urgency)
    /// Requires feature flag: --features factored-actions
    #[cfg(feature = "factored-actions")]
    #[arg(long)]
    use_factored_actions: bool,
}

// In main():
info!("Action space: {} actions",
    if opts.use_factored_actions { 45 } else { 3 });

#[cfg(feature = "factored-actions")]
if opts.use_factored_actions {
    info!("  • Exposure levels: 5 (Short100, Short50, Flat, Long50, Long100)");
    info!("  • Order types: 3 (Market 0.20%, LimitMaker 0.10%, IoC 0.15%)");
    info!("  • Urgency: 3 (Patient 0.5x, Normal 1.0x, Aggressive 1.5x)");
}

// Create trainer with factored actions
#[cfg(feature = "factored-actions")]
let mut trainer = DQNTrainer::new_with_reward_system(
    hyperparams,
    reward_system,
    opts.use_factored_actions,
)?;

#[cfg(not(feature = "factored-actions"))]
if opts.use_factored_actions {
    return Err(anyhow::anyhow!(
        "--use-factored-actions requires compiling with --features factored-actions"
    ));
}

Phase 3: Smoke Tests (ml/tests/dqn_factored_smoke_tests.rs)

Create 5 smoke tests:

1. test_factored_training_5_epochs - Full 5-epoch training
2. test_factored_checkpoint_save_load - Checkpoint persistence
3. test_factored_action_diversity - All 45 actions selectable
4. test_transaction_cost_application - OrderType costs reduce P&L
5. test_position_limit_enforcement - ±100% limits respected

Testing Command

cargo run -p ml --example train_dqn --release --features cuda,factored-actions -- \
  --parquet-file test_data/ES_FUT_180d.parquet \
  --epochs 5 \
  --use-factored-actions \
  --output-dir /tmp/wave1_factored_smoke_test

Key Design Decisions

1. Backward Compatibility: 3-action code path remains untouched when feature flag disabled
2. Runtime Toggle: use_factored_actions flag allows same binary to support both modes
3. Approximate Mapping: TradingAction → factored index uses default order type (Market) and urgency (Normal)
4. Transaction Costs: Integrated directly into reward calculation
5. Position Masking: Applied during action selection to enforce ±100% limits

Implementation Order

1. ✅ Add imports and struct fields
2. ✅ Modify constructor
3. ✅ Update action selection
4. ✅ Add transaction cost logic
5. ✅ Add position masking
6. ✅ Create CLI flags
7. ✅ Write smoke tests
8. ✅ Run validation