foxhunt/docs/archive/waves/WAVE_7_1_QUICK_FIX_GUIDE.md

Wave 7.1: DQN Tensor Rank Quick Fix Guide

Fix Type: Add .squeeze(0) after argmax(1) before to_scalar()

Time to Fix: 5 minutes (3 files, 1 line each)

---

Fix Locations

1. WorkingDQN (PRIMARY)

File: /home/jgrusewski/Work/foxhunt/ml/src/dqn/dqn.rs

Line: 357

Before:

let best_action_idx = q_values
    .argmax(1)?
    .to_scalar::<u32>()

After:

let best_action_idx = q_values
    .argmax(1)?
    .squeeze(0)?        // ✅ ADD THIS LINE
    .to_scalar::<u32>()

---

2. RainbowAgentImpl

File: /home/jgrusewski/Work/foxhunt/ml/src/dqn/rainbow_agent_impl.rs

Line: 151

Before:

let action = q_values
    .argmax(1)
    .map_err(|e| MLError::ModelError(format!("Failed to select action: {}", e)))?
    .to_scalar::<i64>()

After:

let action = q_values
    .argmax(1)
    .map_err(|e| MLError::ModelError(format!("Failed to select action: {}", e)))?
    .squeeze(0)?        // ✅ ADD THIS LINE
    .to_scalar::<i64>()

---

3. RainbowAgent (First Instance)

File: /home/jgrusewski/Work/foxhunt/ml/src/dqn/rainbow_types.rs

Line: 395

Before:

action_values.argmax(1)?
    .to_scalar::<i64>()

After:

action_values.argmax(1)?
    .squeeze(0)?        // ✅ ADD THIS LINE
    .to_scalar::<i64>()

---

4. RainbowAgent (Second Instance)

File: /home/jgrusewski/Work/foxhunt/ml/src/dqn/rainbow_types.rs

Line: 407

Before:

action_values.argmax(1)?
    .to_scalar::<i64>()
    .map_err(|e| MLError::TrainingError(format!("Action extraction failed: {}", e)))? as usize

After:

action_values.argmax(1)?
    .squeeze(0)?        // ✅ ADD THIS LINE
    .to_scalar::<i64>()
    .map_err(|e| MLError::TrainingError(format!("Action extraction failed: {}", e)))? as usize

---

Verification Commands

1. Compile Check

cargo build -p ml

2. Unit Tests

cargo test -p ml dqn::dqn::tests
cargo test -p ml dqn::trainable_adapter

3. Integration Tests

cargo test -p ml dqn_checkpoint_validation
cargo test -p ml dqn_edge_cases

---

Expected Outcomes

✅ Compilation: No more tensor rank errors ✅ Action Selection: Works with batch_size=1 input ✅ Test Pass Rate: 100% for DQN unit tests

---

Why This Fix Works

Problem: argmax(1) on [1, num_actions] returns [1] (rank-1 tensor)

Solution: squeeze(0) reduces [1] to [] (rank-0 scalar)

Result: to_scalar() works on rank-0 tensor

Tensor Shape Flow:

[1, 3] --argmax(1)--> [1] --squeeze(0)--> [] --to_scalar()--> u32

---

Related Patterns in Codebase

This pattern already exists in other parts of DQN:

1. train_step() (dqn.rs:469): .squeeze(1)? after gather
2. network.rs (line 183): .squeeze(0)? before to_vec1()
3. agent.rs (line 397): .squeeze(1)? after gather

Rule: Always squeeze before scalar/vector extraction if batch dimension exists.