# DQN Learning Analysis

## Summary
Analysis of the DQN algorithm and its parent class `RLAlgorithm` to identify potential reasons why the algorithm might not be learning.

## Critical Finding: `set_training_mode()` Does NOT Affect DQN Training

**Important Discovery**: After analyzing the DQN network architecture, **`set_training_mode()` realistically does NOT affect DQN training** because:

1. **DQN uses `nn.LayerNorm`, NOT `nn.BatchNorm`**: 
   - LayerNorm normalizes across features (not batches), so it behaves identically in train and eval modes
   - LayerNorm does NOT use running statistics like BatchNorm does
   - See `agilerl/utils/evolvable_networks.py` line 560-562: `nn.LayerNorm` is used

2. **DQN does NOT use Dropout layers**: 
   - No dropout layers are present in the QNetwork architecture
   - Dropout is the other main layer affected by train/eval mode

3. **`set_training_mode()` doesn't even call `.train()` or `.eval()`**:
   - The method only sets `self.training = training` flag
   - It doesn't actually change the network's training mode

**Conclusion**: The missing `set_training_mode(True)` call and the fact that it doesn't set network mode are **NOT the root cause** of DQN not learning. These are code inconsistencies but don't affect functionality.

## Code Inconsistencies Found (Non-Critical)

### 1. **Missing Training Mode Setting in `train_off_policy.py`**

**Location**: `agilerl/training/train_off_policy.py`, line 238

**Issue**: The training loop does NOT call `agent.set_training_mode(True)` for DQN agents before training begins.

**Comparison**:
- `train_on_policy.py` (line 203): ✅ Calls `agent.set_training_mode(True)`
- `train_multi_agent_off_policy.py` (line 210): ✅ Calls `agent.set_training_mode(True)`
- `train_off_policy.py` (line 238): ❌ **MISSING** `agent.set_training_mode(True)`

**Impact**: **NONE** - This doesn't affect training because DQN networks don't have layers that are affected by train/eval mode.

### 2. **`set_training_mode()` Doesn't Actually Set Network Training Mode**

**Location**: `agilerl/algorithms/core/base.py`, line 648-654

**Current Implementation**:
```python
def set_training_mode(self, training: bool) -> None:
    """Sets the training mode of the algorithm."""
    self.training = training
```

**Issue**: This method only sets a flag `self.training` but does NOT call `.train()` or `.eval()` on the actual neural networks (`self.actor`, `self.actor_target`).

**Impact**: **NONE for DQN** - Since DQN uses LayerNorm (not BatchNorm) and no Dropout, this doesn't affect training.

**Comparison with RainbowDQN**:
- RainbowDQN explicitly calls `self.actor.train()` and `self.actor_target.train()` in `__init__` (lines 221-222)
- RainbowDQN also calls `self.actor.train(mode=training)` in `get_action()` method (line 252)
- DQN does NOT set networks to train mode in `__init__`

**Note**: While this doesn't affect DQN, it's still inconsistent code that could cause issues if BatchNorm or Dropout are added in the future.

## Additional Observations

### 5. **Soft Update Implementation**

**Location**: `agilerl/algorithms/dqn.py`, lines 361-368

The `soft_update()` method looks correct:
```python
def soft_update(self) -> None:
    """Soft updates target network."""
    for eval_param, target_param in zip(
        self.actor.parameters(), self.actor_target.parameters()
    ):
        target_param.data.copy_(
            self.tau * eval_param.data + (1.0 - self.tau) * target_param.data
        )
```

This is called after each learning step (line 358), which is correct.

### 6. **Loss Computation**

**Location**: `agilerl/algorithms/dqn.py`, lines 309-326

The loss computation looks correct:
- Target Q-values are computed with `torch.no_grad()` ✅
- Q-learning update formula is correct: `y_j = rewards + self.gamma * q_target * (1 - dones)` ✅
- Loss is computed and backpropagated correctly ✅

### 7. **Optimizer Setup**

**Location**: `agilerl/algorithms/dqn.py`, lines 151-156

The optimizer is set up correctly with `OptimizerWrapper` and includes the actor network.

## Recommended Fixes

### Fix 1: Add `set_training_mode(True)` to Training Loop

**File**: `agilerl/training/train_off_policy.py`

**Change**: Add after line 238:
```python
for agent_idx, agent in enumerate(pop):  # Loop through population
    agent.set_training_mode(True)  # ADD THIS LINE
    state, info = env.reset()  # Reset environment at start of episode
```

### Fix 2: Override `set_training_mode()` in DQN Class

**File**: `agilerl/algorithms/dqn.py`

**Change**: Add method after `init_hook()`:
```python
def set_training_mode(self, training: bool) -> None:
    """Sets the training mode of the algorithm and networks."""
    super().set_training_mode(training)
    if training:
        self.actor.train()
        self.actor_target.train()
    else:
        self.actor.eval()
        self.actor_target.eval()
```

### Fix 3: Initialize Networks to Train Mode in `__init__`

**File**: `agilerl/algorithms/dqn.py`

**Change**: Add after line 159 (after `self.criterion = nn.MSELoss()`):
```python
# Put the nets into training mode
self.actor.train()
self.actor_target.train()
```

## What Could Actually Be Causing Learning Issues?

Since `set_training_mode()` doesn't affect DQN training, the real issues might be:

1. **Learning rate too high/low**: Check if `lr=0.001` in config is appropriate
2. **Target network update frequency**: `tau=0.001` means very slow updates - is this intentional?
3. **Replay buffer issues**: Is the buffer being populated correctly? Is `learning_delay=0` appropriate?
4. **Epsilon decay**: Check if exploration is decaying too fast/slow
5. **Network initialization**: Are weights initialized properly?
6. **Gradient clipping**: Is there any gradient explosion/vanishing?
7. **Loss computation**: Verify the loss is actually being computed and backpropagated correctly
8. **Experience sampling**: Are experiences being sampled correctly from the replay buffer?

## Priority

1. **LOW**: Fix 1 (add `set_training_mode(True)` to training loop) - For code consistency only
2. **LOW**: Fix 2 (override `set_training_mode()` to actually set network mode) - For future-proofing only
3. **LOW**: Fix 3 (initialize networks to train mode) - For code consistency only

**These fixes are NOT urgent** - they won't fix the learning issue, but they're good practice for code consistency and future-proofing.

## Testing Recommendations

To debug why DQN isn't learning, check:
1. **Loss values**: Are they decreasing? Are they NaN or exploding?
2. **Q-values**: Are they reasonable? Are they converging?
3. **Gradients**: Are gradients flowing? Check `agent.actor.parameters()` gradients
4. **Replay buffer**: Is it filling up? Are experiences diverse?
5. **Target network**: Is it updating? Check if `soft_update()` is being called
6. **Exploration**: Is epsilon decaying appropriately? Are actions diverse?
7. **Rewards**: Are rewards being collected? Are they normalized/scaled appropriately?
