docs: DQN action collapse fix design — 5 fixes for 45-action diversity
Root causes: entropy regularizer dead code, noisy nets disable epsilon, diversity penalty calibrated for 3 actions not 45. Fixes: wire entropy into loss, fix normalization, epsilon floor for noisy nets, sigma scheduling, UCB count-based bonus. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
389
docs/plans/2026-03-03-dqn-action-collapse-fix.md
Normal file
389
docs/plans/2026-03-03-dqn-action-collapse-fix.md
Normal file
@@ -0,0 +1,389 @@
|
||||
# DQN Action Collapse Fix — Design & Implementation Plan
|
||||
|
||||
> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task.
|
||||
|
||||
**Goal:** Fix 45-action DQN collapse where the model converges to a single action (A38 = Long100/Market/Aggressive) across all hyperopt trials, despite existing anti-collapse infrastructure.
|
||||
|
||||
**Root Causes:**
|
||||
1. Entropy regularizer is dead code (initialized, never called)
|
||||
2. Noisy nets disable epsilon-greedy but provide insufficient exploration for 45 actions
|
||||
3. Diversity penalty entropy normalized for 3 actions, not 45
|
||||
|
||||
**Approach:** Fix all three existing mechanisms + add two new ones (epsilon floor for noisy nets, count-based exploration bonus).
|
||||
|
||||
---
|
||||
|
||||
## Task 1: Wire entropy regularizer into TD loss
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/dqn/dqn.rs` (add entropy bonus to `compute_loss` or loss path)
|
||||
- Modify: `crates/ml/src/trainers/dqn/trainer.rs` (pass entropy regularizer into training step)
|
||||
|
||||
**Step 1: Add entropy bonus computation to the Q-learning loss**
|
||||
|
||||
The entropy regularizer at `crates/ml/src/dqn/entropy_regularization.rs` is fully implemented but never called. Wire it into the training loop.
|
||||
|
||||
In the training step where loss is computed (trainer.rs, the `train_step` / gradient update section), after computing Q-values for the current batch:
|
||||
|
||||
```rust
|
||||
// Compute softmax policy from Q-values: π(a|s) = softmax(Q(s,a) / τ)
|
||||
// Then entropy: H(π) = -Σ π(a|s) * log(π(a|s))
|
||||
// Add to loss: loss = td_loss - entropy_coefficient * mean_entropy
|
||||
```
|
||||
|
||||
Specifically:
|
||||
1. After the forward pass that produces `q_values` tensor (shape [batch, 45])
|
||||
2. Compute `log_softmax(q_values, dim=1)` and `softmax(q_values, dim=1)`
|
||||
3. Entropy per sample: `H = -(softmax * log_softmax).sum(dim=1)` → shape [batch]
|
||||
4. Mean entropy: `mean_H = H.mean()`
|
||||
5. Final loss: `loss = td_loss - entropy_coefficient * mean_H`
|
||||
|
||||
The `entropy_coefficient` field already exists in `DQNHyperparameters` (default 0.01, hyperopt range [0.0, 0.1]).
|
||||
|
||||
**Step 2: Remove unused `entropy_regularizer` field from trainer struct**
|
||||
|
||||
The `Arc<EntropyRegularizer>` stored at trainer.rs:754 is no longer needed — entropy is computed directly on Q-value tensors. Remove the field and the initialization at lines 531-535.
|
||||
|
||||
**Step 3: Add unit test**
|
||||
|
||||
Test that with `entropy_coefficient > 0`, training on a batch where all Q-values are identical (uniform policy) produces lower loss than a batch with one dominant Q-value (collapsed policy).
|
||||
|
||||
**Step 4: Build and test**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo test -p ml --lib -- trainers::dqn
|
||||
```
|
||||
|
||||
**Step 5: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/src/dqn/ crates/ml/src/trainers/dqn/
|
||||
git commit -m "fix(dqn): wire entropy regularization into TD loss (was dead code)"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Task 2: Fix diversity penalty entropy normalization
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/dqn/reward.rs` (fix `calculate_entropy()` normalization)
|
||||
|
||||
**Step 1: Fix entropy normalization from log(3) to log(45)**
|
||||
|
||||
In `reward.rs`, find `calculate_entropy()` (around line 364-395). The function computes Shannon entropy of recent actions. Currently it normalizes by `log(3)` (3 legacy action categories). Change to `log(45)` to match the 45 factored actions.
|
||||
|
||||
Also change the entropy threshold from 0.5 to 0.3. With 45 actions, entropy of 0.3 (normalized) means ~4 distinct actions — a reasonable collapse threshold.
|
||||
|
||||
```rust
|
||||
// Before:
|
||||
let max_entropy = (3.0_f64).ln(); // log(3) = 1.099
|
||||
let entropy_threshold = 0.5;
|
||||
|
||||
// After:
|
||||
let max_entropy = (45.0_f64).ln(); // log(45) = 3.807
|
||||
let entropy_threshold = 0.3;
|
||||
```
|
||||
|
||||
**Step 2: Update the action counting to use 45 action indices**
|
||||
|
||||
If the entropy calculation currently maps actions to 3 categories (Buy/Sell/Hold), change it to use the full `to_index()` (0-44) for the frequency distribution.
|
||||
|
||||
**Step 3: Add test**
|
||||
|
||||
Test that:
|
||||
- Uniform distribution over 45 actions → entropy ≈ 1.0 (normalized) → no penalty
|
||||
- Single-action collapse (all A38) → entropy ≈ 0.0 → penalty fires (-0.1)
|
||||
- 5 distinct actions → entropy ≈ 0.42 → penalty fires (< 0.3 threshold? depends on distribution)
|
||||
|
||||
**Step 4: Build and test**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo test -p ml --lib -- dqn::reward
|
||||
```
|
||||
|
||||
**Step 5: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/src/dqn/reward.rs
|
||||
git commit -m "fix(dqn): normalize diversity entropy for 45 actions (was calibrated for 3)"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Task 3: Add epsilon floor for noisy nets
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/trainers/dqn/config.rs` (add `noisy_epsilon_floor` field)
|
||||
- Modify: `crates/ml/src/dqn/dqn.rs` (use floor in `select_action()`)
|
||||
- Modify: `crates/ml/src/hyperopt/adapters/dqn.rs` (add to search space)
|
||||
|
||||
**Step 1: Add config field**
|
||||
|
||||
In `config.rs`, add to `DQNHyperparameters`:
|
||||
|
||||
```rust
|
||||
/// Minimum epsilon when noisy nets are enabled (guarantees some random exploration)
|
||||
pub noisy_epsilon_floor: f64, // default: 0.05
|
||||
```
|
||||
|
||||
Add default value in the `Default` impl: `noisy_epsilon_floor: 0.05`.
|
||||
|
||||
**Step 2: Use floor in select_action()**
|
||||
|
||||
In `dqn.rs`, change the noisy nets epsilon logic:
|
||||
|
||||
```rust
|
||||
// Before:
|
||||
let effective_epsilon = if self.config.use_noisy_nets {
|
||||
0.0 // Learned exploration via parameter noise
|
||||
} else {
|
||||
self.epsilon
|
||||
};
|
||||
|
||||
// After:
|
||||
let effective_epsilon = if self.config.use_noisy_nets {
|
||||
self.config.noisy_epsilon_floor // Minimum random exploration even with noisy nets
|
||||
} else {
|
||||
self.epsilon
|
||||
};
|
||||
```
|
||||
|
||||
**Step 3: Add to hyperopt search space**
|
||||
|
||||
In `dqn.rs` hyperopt adapter, add `noisy_epsilon_floor` to:
|
||||
- `DQNParams` struct
|
||||
- `continuous_bounds()` — range [0.02, 0.10]
|
||||
- `from_continuous()` / `to_continuous()`
|
||||
- `to_hyperparams()` mapping
|
||||
|
||||
**Step 4: Skip epsilon decay when noisy nets enabled**
|
||||
|
||||
In `trainer.rs`, guard the `agent.update_epsilon()` call:
|
||||
|
||||
```rust
|
||||
// Before:
|
||||
agent.update_epsilon();
|
||||
|
||||
// After:
|
||||
if !self.hyperparams.use_noisy_nets {
|
||||
agent.update_epsilon();
|
||||
}
|
||||
```
|
||||
|
||||
**Step 5: Build and test**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo check -p ml
|
||||
SQLX_OFFLINE=true cargo test -p ml --lib -- dqn::dqn::tests
|
||||
```
|
||||
|
||||
**Step 6: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/src/trainers/dqn/config.rs crates/ml/src/dqn/dqn.rs crates/ml/src/hyperopt/adapters/dqn.rs crates/ml/src/trainers/dqn/trainer.rs
|
||||
git commit -m "feat(dqn): add epsilon floor for noisy nets (0.05 default, hyperopt-tunable)"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Task 4: Enable noisy sigma scheduling with higher initial sigma
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/trainers/dqn/config.rs` (change defaults)
|
||||
- Modify: `crates/ml/src/dqn/dqn.rs` (add sigma floor)
|
||||
|
||||
**Step 1: Change defaults for sigma scheduling**
|
||||
|
||||
In `config.rs`, update the defaults:
|
||||
|
||||
```rust
|
||||
// Before:
|
||||
enable_noisy_sigma_scheduler: false,
|
||||
noisy_sigma_initial: 0.6,
|
||||
noisy_sigma_final: 0.4,
|
||||
|
||||
// After:
|
||||
enable_noisy_sigma_scheduler: true, // Enable by default
|
||||
noisy_sigma_initial: 0.8, // Wider initial exploration
|
||||
noisy_sigma_final: 0.4,
|
||||
```
|
||||
|
||||
**Step 2: Add sigma floor**
|
||||
|
||||
In the sigma scheduling logic (wherever sigma is computed/annealed), add a minimum:
|
||||
|
||||
```rust
|
||||
let sigma = computed_sigma.max(0.3); // Never let sigma drop below 0.3
|
||||
```
|
||||
|
||||
**Step 3: Update hyperopt adapter defaults**
|
||||
|
||||
Make sure the hyperopt adapter uses the new defaults when constructing `DQNHyperparameters`.
|
||||
|
||||
**Step 4: Build and test**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo check -p ml
|
||||
```
|
||||
|
||||
**Step 5: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/src/trainers/dqn/config.rs crates/ml/src/dqn/dqn.rs
|
||||
git commit -m "feat(dqn): enable noisy sigma scheduling (0.8→0.4, floor=0.3)"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Task 5: Add count-based exploration bonus (UCB)
|
||||
|
||||
**Files:**
|
||||
- Create: `crates/ml/src/dqn/count_bonus.rs`
|
||||
- Modify: `crates/ml/src/dqn/mod.rs` (add module)
|
||||
- Modify: `crates/ml/src/trainers/dqn/config.rs` (add config fields)
|
||||
- Modify: `crates/ml/src/dqn/dqn.rs` (apply bonus in action selection)
|
||||
- Modify: `crates/ml/src/hyperopt/adapters/dqn.rs` (add to search space)
|
||||
|
||||
**Step 1: Create count_bonus.rs**
|
||||
|
||||
```rust
|
||||
/// UCB-style count-based exploration bonus for action selection.
|
||||
///
|
||||
/// Adds bonus = β * sqrt(log(N) / (1 + n_a)) to Q-values during action selection,
|
||||
/// where N = total actions taken this epoch, n_a = count of action a.
|
||||
/// This gives direct exploration incentive to under-visited actions.
|
||||
pub struct CountBonus {
|
||||
action_counts: [u64; 45],
|
||||
total_count: u64,
|
||||
coefficient: f64,
|
||||
}
|
||||
|
||||
impl CountBonus {
|
||||
pub fn new(coefficient: f64) -> Self { ... }
|
||||
|
||||
/// Record that action was taken
|
||||
pub fn record_action(&mut self, action_idx: usize) { ... }
|
||||
|
||||
/// Compute bonus for each of 45 actions, returns [f64; 45]
|
||||
pub fn bonuses(&self) -> [f64; 45] {
|
||||
if self.total_count == 0 {
|
||||
return [self.coefficient; 45]; // Max bonus when no data
|
||||
}
|
||||
let log_n = (self.total_count as f64).ln();
|
||||
let mut result = [0.0; 45];
|
||||
for a in 0..45 {
|
||||
result[a] = self.coefficient * (log_n / (1.0 + self.action_counts[a] as f64)).sqrt();
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
/// Reset counts (call at epoch boundary)
|
||||
pub fn reset(&mut self) { ... }
|
||||
}
|
||||
```
|
||||
|
||||
**Step 2: Add config fields**
|
||||
|
||||
In `config.rs`:
|
||||
|
||||
```rust
|
||||
pub use_count_bonus: bool, // default: true
|
||||
pub count_bonus_coefficient: f64, // default: 0.1
|
||||
```
|
||||
|
||||
**Step 3: Wire into action selection**
|
||||
|
||||
In `dqn.rs` `select_action()`, when choosing the greedy action:
|
||||
|
||||
```rust
|
||||
// After computing q_values from forward pass:
|
||||
if self.config.use_count_bonus {
|
||||
let bonuses = self.count_bonus.bonuses();
|
||||
let bonus_tensor = Tensor::from_vec(bonuses.to_vec(), 45, device)?;
|
||||
let q_with_bonus = q_values.broadcast_add(&bonus_tensor)?;
|
||||
// Use q_with_bonus for argmax instead of q_values
|
||||
}
|
||||
```
|
||||
|
||||
Record action after selection:
|
||||
|
||||
```rust
|
||||
if self.config.use_count_bonus {
|
||||
self.count_bonus.record_action(action.to_index());
|
||||
}
|
||||
```
|
||||
|
||||
**Step 4: Reset at epoch boundary**
|
||||
|
||||
In `trainer.rs`, at the start of each epoch:
|
||||
|
||||
```rust
|
||||
agent.reset_count_bonus(); // Resets per-epoch exploration pressure
|
||||
```
|
||||
|
||||
**Step 5: Add to hyperopt search space**
|
||||
|
||||
- `count_bonus_coefficient` — range [0.01, 0.5]
|
||||
- `use_count_bonus` — always true (not tuned, it's a safety mechanism)
|
||||
|
||||
**Step 6: Add tests**
|
||||
|
||||
- Test that bonuses are high for unvisited actions and low for frequent actions
|
||||
- Test that after recording 100 of action 38, action 38's bonus is near zero while others are high
|
||||
- Test reset clears all counts
|
||||
|
||||
**Step 7: Build and test**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo test -p ml --lib -- dqn::count_bonus
|
||||
SQLX_OFFLINE=true cargo check -p ml
|
||||
```
|
||||
|
||||
**Step 8: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/src/dqn/count_bonus.rs crates/ml/src/dqn/mod.rs crates/ml/src/trainers/dqn/config.rs crates/ml/src/dqn/dqn.rs crates/ml/src/hyperopt/adapters/dqn.rs
|
||||
git commit -m "feat(dqn): add UCB count-based exploration bonus for 45-action diversity"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Task 6: Update monitoring to report exploration diagnostics
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/trainers/dqn/trainer.rs` (log exploration source stats)
|
||||
|
||||
**Step 1: Track exploration source in action diversity log**
|
||||
|
||||
In the per-epoch action diversity logging, add:
|
||||
- Count of random (epsilon) vs greedy vs noisy-net actions
|
||||
- Effective entropy over 45 actions (normalized by log(45))
|
||||
- Count bonus range (min/max/mean)
|
||||
|
||||
**Step 2: Add Prometheus metrics**
|
||||
|
||||
In `crates/common/src/metrics/training_metrics.rs`, add:
|
||||
|
||||
```rust
|
||||
foxhunt_training_action_entropy // Normalized entropy (0-1) of epoch action distribution
|
||||
foxhunt_training_epsilon_actions // Count of epsilon-random actions this epoch
|
||||
foxhunt_training_greedy_actions // Count of greedy actions this epoch
|
||||
```
|
||||
|
||||
**Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/src/trainers/dqn/trainer.rs crates/common/src/metrics/training_metrics.rs
|
||||
git commit -m "feat(monitoring): add exploration diagnostics (entropy, epsilon/greedy split)"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Task 7: Full build verification
|
||||
|
||||
**Step 1:** `SQLX_OFFLINE=true cargo check --workspace`
|
||||
**Step 2:** `SQLX_OFFLINE=true cargo test -p ml --lib`
|
||||
**Step 3:** `SQLX_OFFLINE=true cargo test -p common --lib`
|
||||
**Step 4:** `SQLX_OFFLINE=true cargo clippy --workspace -- -D warnings`
|
||||
|
||||
Expected: 0 errors, 0 warnings, all tests pass.
|
||||
Reference in New Issue
Block a user