# 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)
```rust
// 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)
```rust
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)
```rust
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)
```rust
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)
```rust
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)
```rust
#[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)
```rust
#[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:

```rust
#[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

```bash
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