foxhunt/ml/zen_generated.code

<GENERATED-CODE>
Here are the step-by-step instructions to fix the training instability and improve reward scaling.

1.  **Update `ml/src/trainers/dqn.rs`**: Modify the training loop to break the positive feedback loop in the risk-adjusted reward calculation. We will update the PnL history tracker with the raw reward *before* applying the Sharpe ratio adjustment.

2.  **Update `ml/src/dqn/reward.rs`**: Overhaul the reward logic to be more statistically robust. This involves three key changes:
    *   Refactor the P&L reward to be based on percentage portfolio change, which is a more stationary signal.
    *   Introduce a `RewardNormalizer` struct that uses Welford's algorithm to compute running mean and standard deviation.
    *   Integrate the normalizer into the `RewardFunction` to automatically scale all rewards before they are used for training.

<UPDATED_EXISTING_FILE: ml/src/trainers/dqn.rs>
```rust
// ... (code before line 1169)
                    let raw_reward = reward_decimal.to_string().parse::<f32>().unwrap_or(0.0) as f64;

                    // Calculate price return for volatility tracking
                    // Add safety check for division by zero
                    let price_return = if current_close.abs() > 1e-9 {
                        (next_close - current_close) / current_close
                    } else {
                        0.0
                    };

                    // CRITICAL FIX: Update adaptive risk trackers with the raw, unadjusted reward.
                    // This breaks the positive feedback loop where amplified rewards were used
                    // to calculate the next amplification factor (Sharpe ratio).
                    self.update_risk_trackers(raw_reward, price_return);

                    // WAVE 16S: Apply risk-adjusted rewards (Sharpe ratio) using history of raw rewards
                    let risk_adjusted_reward = self.calculate_risk_adjusted_reward(raw_reward);

                    // Convert back to f32 for experience storage
                    let reward = risk_adjusted_reward as f32;

                    // WAVE 16S: Log adaptive features periodically
                    if i % 100 == 0 && i > 0 {
// ... (code from line 1185 to 2677)
    /// Update adaptive risk trackers with new market data
    fn update_risk_trackers(&mut self, reward: f64, price_return: f64) {
        // Update PnL history
        self.pnl_history.push_back(reward);
        if self.pnl_history.len() > 1000 {
// ... (rest of the file is unchanged)
```
</UPDATED_EXISTING_FILE>

<UPDATED_EXISTING_FILE: ml/src/dqn/reward.rs>
```rust
//! Trading-specific reward functions for DQN

// CANONICAL TYPE IMPORTS - Use common::Decimal
use serde::{Deserialize, Serialize};
// For Decimal::from_f64
use common::types::Price;
use rust_decimal::Decimal;

use super::action_space::FactoredAction;
use super::agent::{TradingAction, TradingState};
use crate::MLError;

/// Configuration for reward function
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RewardConfig {
    /// Weight for P&L component
    pub pnl_weight: Decimal,
    /// Weight for risk penalty
    pub risk_weight: Decimal,
    /// Weight for transaction cost penalty
    pub cost_weight: Decimal,
    /// Weight for hold reward (to reduce over-trading)
    pub hold_reward: Decimal,
    /// Price movement threshold for dynamic HOLD reward (as fraction, e.g., 0.02 = 2%)
    pub movement_threshold: Decimal,
    /// Weight for HOLD action penalty during high volatility (negative value applied)
    pub hold_penalty_weight: Decimal,
    /// Weight for diversity penalty (negative value to penalize low entropy, -0.1 default)
    pub diversity_weight: Decimal,
    /// Enable running normalization of rewards
    pub enable_reward_normalization: bool,
}

impl Default for RewardConfig {
    fn default() -> Self {
        Self {
            pnl_weight: Decimal::ONE,
            risk_weight: Decimal::try_from(0.1).unwrap_or(Decimal::ZERO),
            cost_weight: Decimal::try_from(0.05).unwrap_or(Decimal::ZERO),
            hold_reward: Decimal::try_from(0.001).unwrap_or(Decimal::ZERO),
            movement_threshold: Decimal::try_from(0.01).unwrap_or(Decimal::ZERO), // 1% matches data distribution
            hold_penalty_weight: Decimal::try_from(0.01).unwrap_or(Decimal::ZERO), // 1% default penalty
            diversity_weight: Decimal::try_from(-0.1).unwrap_or(Decimal::ZERO), // -0.1 default (100x stronger than hold_reward)
            enable_reward_normalization: true, // Default to enabled for robustness
        }
    }
}

/// A running normalizer for rewards using Welford's online algorithm.
/// This adaptively scales rewards to have a mean of ~0 and stddev of ~1,
/// which stabilizes DQN training with non-stationary reward distributions.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RewardNormalizer {
    count: u64,
    mean: f64,
    m2: f64, // Sum of squares of differences from the current mean
    enabled: bool,
}

impl RewardNormalizer {
    /// Creates a new `RewardNormalizer`.
    pub fn new(enabled: bool) -> Self {
        Self {
            count: 0,
            mean: 0.0,
            m2: 0.0,
            enabled,
        }
    }

    /// Updates the running statistics with a new reward value using Welford's algorithm.
    pub fn update(&mut self, reward: Decimal) {
        if !self.enabled {
            return;
        }
        if let Ok(reward_f64) = reward.try_into::<f64>() {
            self.count += 1;
            let delta = reward_f64 - self.mean;
            self.mean += delta / self.count as f64;
            let delta2 = reward_f64 - self.mean;
            self.m2 += delta * delta2;
        }
    }

    /// Normalizes a reward value based on the current running statistics.
    pub fn normalize(&self, reward: Decimal) -> Decimal {
        // Do not normalize until we have seen at least a few samples
        if !self.enabled || self.count < 10 {
            return reward;
        }

        let variance = self.m2 / self.count as f64;
        let std_dev = variance.sqrt();

        if let Ok(reward_f64) = reward.try_into::<f64>() {
            // Clip std_dev to avoid division by zero or near-zero values, which can cause explosions.
            let safe_std_dev = std_dev.max(1e-8);
            let normalized_reward = (reward_f64 - self.mean) / safe_std_dev;
            Decimal::try_from(normalized_reward).unwrap_or(reward)
        } else {
            reward
        }
    }
}

/// Risk metrics for trading state
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct RiskMetrics {
    /// Value at Risk (95%)
    pub var_95: Decimal,
    /// Maximum drawdown
    pub max_drawdown: Decimal,
    /// Sharpe ratio
    pub sharpe_ratio: Decimal,
    /// Volatility
    pub volatility: Decimal,
}

/// Market data snapshot
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct MarketData {
    /// Current bid price
    pub bid: Price,
    /// Current ask price
    pub ask: Price,
    /// Bid-ask spread
    pub spread: Price,
    /// Volume
    pub volume: Decimal,
}

/// Calculate Shannon entropy for action distribution
///
/// # Arguments
/// * `recent_actions` - Sliding window of recent actions (last 100 actions)
///
/// # Returns
/// Shannon entropy H = -Σ(p_i * log2(p_i)) where p_i is frequency of action i
/// Range: [0.0, 1.585] (0 = all same action, 1.585 = perfectly balanced)
///
/// # Note
/// For FactoredAction, we convert to legacy TradingAction (Buy/Sell/Hold) for entropy calculation
/// to maintain backward compatibility with the original 3-action entropy range.
fn calculate_entropy(recent_actions: &[FactoredAction]) -> Decimal {
    if recent_actions.is_empty() {
        return Decimal::ONE; // Default to high entropy if no history
    }

    // Convert to legacy actions and count
    let mut counts = [0, 0, 0]; // BUY, SELL, HOLD
    for action in recent_actions {
        let legacy_action = action.to_legacy_action();
        counts[legacy_action as usize] += 1;
    }

    let total = recent_actions.len() as f64;
    let mut entropy = Decimal::ZERO;

    for &count in &counts {
        if count > 0 {
            let p = Decimal::try_from(count as f64 / total).unwrap_or(Decimal::ZERO);
            // Shannon entropy: -p * log2(p)
            // Use natural log and convert: log2(x) = ln(x) / ln(2)
            if let Ok(p_f64) = TryInto::<f64>::try_into(p) {
                if p_f64 > 0.0 {
                    let log2_p = p_f64.ln() / 2.0_f64.ln();
                    let term = p * Decimal::try_from(log2_p).unwrap_or(Decimal::ZERO);
                    entropy -= term;
                }
            }
        }
    }

    entropy
}

/// Reward function for `DQN` training
#[derive(Debug)]
pub struct RewardFunction {
    /// Configuration
    config: RewardConfig,
    /// Previous rewards for tracking
    reward_history: Vec<Decimal>,
    /// Running normalizer for rewards
    normalizer: RewardNormalizer,
}

impl RewardFunction {
    /// Create a new reward function
    pub fn new(config: RewardConfig) -> Self {
        Self {
            normalizer: RewardNormalizer::new(config.enable_reward_normalization),
            config,
            reward_history: Vec::new(),
        }
    }

    /// Validate that a TradingState has the required portfolio features
    ///
    /// # Required Structure
    /// Portfolio features must have length >= 3:
    /// - [0]: Portfolio value (cash + unrealized P&L)
    /// - [1]: Position size (signed: +Long, -Short, 0=flat)
    /// - [2]: Bid-ask spread
    ///
    /// # Returns
    /// Ok(()) if valid, Err(MLError) with descriptive message if invalid
    fn validate_portfolio_features(state: &TradingState) -> Result<(), MLError> {
        if state.portfolio_features.len() < 3 {
            return Err(MLError::InvalidInput(
                format!(
                    "TradingState portfolio_features must have length >= 3 (got {}). Expected: [portfolio_value, position_size, spread]",
                    state.portfolio_features.len()
                )
            ));
        }
        Ok(())
    }

    /// Calculate reward for a state transition
    ///
    /// # Arguments
    /// * `action` - Trading action taken (FactoredAction with exposure, order type, urgency)
    /// * `current_state` - Current trading state (128-dim: 4 price + 121 technical + 3 portfolio)
    /// * `next_state` - Next trading state after action (128-dim structure)
    /// * `recent_actions` - Sliding window of last 100 actions for diversity penalty
    ///
    /// # Validation
    /// Validates that both states have proper portfolio_features (length >= 3).
    /// Logs warnings if features are missing but continues with default values.
    ///
    /// # Note
    /// Converts FactoredAction to legacy TradingAction for reward calculation logic.
    pub fn calculate_reward(
        &mut self,
        action: FactoredAction,
        current_state: &TradingState,
        next_state: &TradingState,
        recent_actions: &[FactoredAction],
    ) -> Result<Decimal, MLError> {
        // Validate portfolio features (non-fatal, logs warnings)
        if let Err(e) = Self::validate_portfolio_features(current_state) {
            tracing::warn!("Current state validation: {}", e);
        }
        if let Err(e) = Self::validate_portfolio_features(next_state) {
            tracing::warn!("Next state validation: {}", e);
        }

        // Convert to legacy action for reward logic
        let legacy_action = action.to_legacy_action();

        let base_reward = match legacy_action {
            TradingAction::Buy | TradingAction::Sell => {
                // Calculate P&L-based reward
                let pnl_reward = self.calculate_pnl_reward(current_state, next_state)?;

                // Calculate risk penalty
                let risk_penalty = self.calculate_risk_penalty(next_state);

                // Calculate transaction cost penalty
                let cost_penalty = self.calculate_cost_penalty(current_state, next_state);

                self.config.pnl_weight * pnl_reward
                    - self.config.risk_weight * risk_penalty
                    - self.config.cost_weight * cost_penalty
            },
            TradingAction::Hold => {
                // Dynamic HOLD reward based on price movement
                self.calculate_hold_reward(current_state, next_state)?
            },
        };

        // Calculate diversity bonus (in-training regularization)
        // Entropy threshold: 0.5 (50% of max entropy for 3 actions = 1.585)
        // Low entropy → action bias → apply penalty (-0.1)
        let entropy = calculate_entropy(recent_actions);
        let entropy_threshold = Decimal::try_from(0.5).unwrap_or(Decimal::ZERO);
        let diversity_bonus = if entropy < entropy_threshold {
            self.config.diversity_weight // -0.1 (penalty for low diversity)
        } else {
            Decimal::ZERO // No penalty for balanced actions
        };

        let final_reward = base_reward + diversity_bonus;

        // Update normalizer with the raw, unclamped reward
        self.normalizer.update(final_reward);

        // Normalize the reward before clamping
        let normalized_reward = self.normalizer.normalize(final_reward);

        // Clamp reward to prevent cumulative explosions. This is a final safeguard.
        let clamped_reward = normalized_reward.clamp(Decimal::from(-1), Decimal::ONE);

        // Store the original, unclamped reward for statistical analysis
        self.reward_history.push(final_reward);
        if self.reward_history.len() > 1000 {
            self.reward_history.remove(0);
        }

        Ok(clamped_reward)
    }

    /// Calculate P&L-based reward component as a percentage of portfolio value.
    ///
    /// This approach creates a more stationary reward signal, as the reward is
    /// proportional to the percentage gain/loss rather than the absolute dollar amount.
    /// A 0.5% gain yields a similar reward regardless of whether the portfolio is $100K or $1M.
    fn calculate_pnl_reward(
        &self,
        current_state: &TradingState,
        next_state: &TradingState,
    ) -> Result<Decimal, MLError> {
        // Validate portfolio_features length (defensive check)
        if current_state.portfolio_features.is_empty() {
            tracing::warn!("Current state portfolio_features is empty, using 0.0 for P&L reward");
            return Ok(Decimal::ZERO);
        }
        if next_state.portfolio_features.is_empty() {
            tracing::warn!("Next state portfolio_features is empty, using 0.0 for P&L reward");
            return Ok(Decimal::ZERO);
        }

        let current_value =
            Decimal::try_from(*current_state.portfolio_features.get(0).unwrap_or(&100000.0) as f64)
                .unwrap_or_else(|_| Decimal::from(100000));
        let next_value =
            Decimal::try_from(*next_state.portfolio_features.get(0).unwrap_or(&100000.0) as f64)
                .unwrap_or(current_value);

        // Avoid division by zero if portfolio value is somehow zero.
        if current_value.is_zero() {
            return Ok(Decimal::ZERO);
        }

        let pnl_change = next_value - current_value;

        // Calculate P&L as a percentage of portfolio value to create a stationary reward signal.
        let pnl_percentage = pnl_change / current_value;

        // Scale percentage to a more intuitive reward magnitude.
        // A 1% portfolio gain (pnl_percentage = 0.01) becomes a reward of 1.0.
        // A 0.1% gain (0.001) becomes a reward of 0.1.
        let scaled_pnl = pnl_percentage * Decimal::from(100);

        Ok(scaled_pnl)
    }

    /// Calculate risk penalty
    ///
// ... (rest of the file is unchanged)
```
</UPDATED_EXISTING_FILE>
</GENERATED-CODE>