# 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**:
```rust
let best_action_idx = q_values
    .argmax(1)?
    .to_scalar::<u32>()
```

**After**:
```rust
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**:
```rust
let action = q_values
    .argmax(1)
    .map_err(|e| MLError::ModelError(format!("Failed to select action: {}", e)))?
    .to_scalar::<i64>()
```

**After**:
```rust
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**:
```rust
action_values.argmax(1)?
    .to_scalar::<i64>()
```

**After**:
```rust
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**:
```rust
action_values.argmax(1)?
    .to_scalar::<i64>()
    .map_err(|e| MLError::TrainingError(format!("Action extraction failed: {}", e)))? as usize
```

**After**:
```rust
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
```bash
cargo build -p ml
```

### 2. Unit Tests
```bash
cargo test -p ml dqn::dqn::tests
cargo test -p ml dqn::trainable_adapter
```

### 3. Integration Tests
```bash
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.