Markov Fleet C51 Distributional DQN (v0.8)

C51 Distributional Reinforcement Learning implementation for 100-bridge fleet maintenance optimization.

Based on: "A Distributional Perspective on Reinforcement Learning" (Bellemare et al., PMLR 2017)

---

🎯 Key Features (v0.8)

Core Algorithm: C51 Distributional RL

• Categorical distribution over return values (51 atoms)
• Distributional Bellman update with projection step
• Cross-entropy loss instead of MSE loss
• Support range: [V_min, V_max] = [-100, +100]
• Delta_z: 4.0 (uniform spacing)

Previous Optimizations (v0.7)

• ✅ Noisy Networks for exploration (no ε-greedy needed)
• ✅ Dueling DQN architecture
• ✅ Double DQN for target calculation
• ✅ Prioritized Experience Replay (PER)
• ✅ N-step Learning (n=3)
• ✅ AsyncVectorEnv for parallel training
• ✅ Mixed Precision Training (AMP)

---

📊 What's New in v0.8

1. C51 Network Architecture (FleetC51)

# Output: Categorical distributions instead of Q-values
q_values, distributions = agent(state)
# q_values: [batch, n_bridges, n_actions]          # Expected values
# distributions: [batch, n_bridges, n_actions, 51]  # Probability distributions

Architecture:

• Shared network: [100] -> [512, 256]
• Value stream: [256] -> [128] -> [51] (NoisyLinear)
• Advantage stream: [256] -> [128] -> [30600] (NoisyLinear, 100×6×51)

2. Distributional Bellman Update

# Project T_z = r + γ * z onto fixed support
# Cross-entropy loss: -sum(m * log(p))
loss = c51_distributional_loss(agent, target_net, s, a, r, s', done)

Key Steps:

1. Compute target distribution with projection
2. Calculate cross-entropy between current and target distributions
3. Update priorities using KL divergence

3. Training Improvements

• Loss Function: Cross-entropy loss with distributional projection
• TD Errors: KL divergence for PER priority updates
• Q-Value Estimation: Expected value from distribution

---

🚀 Quick Start

1. Install Dependencies

pip install -r requirements.txt

2. Test C51 Implementation

python test_noisy_net.py

Expected output:

C51 DISTRIBUTIONAL DQN VERIFICATION (v0.8)
✓ NoisyLinear test PASSED!
✓ FleetC51 test PASSED!
✓ Training loop test PASSED!
ALL TESTS PASSED! ✓

3. Train C51 Agent

# Full training (1000 episodes)
python train_markov_fleet.py --episodes 1000 --device cuda

# Quick test (100 episodes)
python train_markov_fleet.py --episodes 100 --device cuda

Training output:

MARKOV FLEET C51 DISTRIBUTIONAL DQN TRAINING (v0.8)
================================================================================
C51 Distribution (Bellemare et al., PMLR 2017):
  N_atoms: 51
  V_min: -100.0, V_max: 100.0
  Delta_z: 4.00

Optimizations:
  ✓ C51 Distributional RL (categorical distribution)
  ✓ Cross-entropy loss with projection
  ✓ Noisy Networks (no ε-greedy needed!)
  ...

4. Visualize Results

# Training curves
python visualize_markov_v08.py outputs_markov/models/markov_fleet_c51_final_1000ep.pt

# Training curves + C51 distributions
python visualize_markov_v08.py outputs_markov/models/markov_fleet_c51_final_1000ep.pt --plot-dist

Training Results (200-Bridge Fleet, 25k Episodes)

Training Curves:

C51 Return Distributions:

---

📁 Project Structure

markov-dqn-v08-c51/
├── train_markov_fleet.py        # Main training script (C51)
│   ├── FleetC51                  # C51 Distributional DQN
│   ├── c51_distributional_loss() # Cross-entropy loss
│   └── train_markov_fleet()      # Training loop
│
├── test_noisy_net.py             # C51 verification tests
├── visualize_markov_v08.py       # Training visualization + C51 distributions
├── config.yaml                   # Hyperparameters (C51 params)
│
├── src/
│   ├── markov_fleet_environment.py  # Fleet environment
│   └── fleet_environment_gym.py     # Gym interface
│
└── outputs_markov/
    ├── models/                   # Trained models
    ├── plots/                    # Visualizations
    └── logs/                     # Training logs

---

🔧 Configuration (config.yaml)

C51 Distribution Parameters

network:
  n_atoms: 51              # Number of atoms in distribution
  v_min: -100.0            # Minimum return value
  v_max: 100.0             # Maximum return value
  # Support: z_i = v_min + i * delta_z, i = 0, ..., 50
  # delta_z = (v_max - v_min) / (n_atoms - 1) = 4.0

Training Parameters

training:
  num_episodes: 1000
  learning_rate: 0.0015
  batch_size: 64
  buffer_capacity: 10000
  target_sync_steps: 500
  n_steps: 3               # N-step returns

Fleet Configuration

fleet:
  n_urban: 20              # Urban bridges (higher importance)
  n_rural: 80              # Rural bridges
  urban_importance: 1.5    # Reward weight multiplier

---

📈 C51 Theory Overview

Distributional Bellman Equation

Instead of learning $Q(s, a)$, C51 learns the distribution $Z(s, a)$:

$$Z(s, a) \sim \text{Categorical}(z_1, ..., z_{51})$$

Bellman Update: $$\mathcal{T}Z(s, a) \stackrel{D}{=} R(s, a) + \gamma Z(s', a^*)$$

Projection: Project $\mathcal{T}Z$ onto fixed support $[V_{\min}, V_{\max}]$ with 51 atoms.

Cross-Entropy Loss

$$\mathcal{L} = -\sum_{i=1}^{51} m_i \log p_i$$

where:

• $m_i$: Projected target distribution
• $p_i$: Current distribution

Q-Value Estimation

$$Q(s, a) = \mathbb{E}[Z(s, a)] = \sum_{i=1}^{51} z_i \cdot p_i$$

---

🔬 Expected Improvements over v0.7

Theoretical Advantages

1. Better approximation of return distribution (not just mean)
2. Reduced instability from distributional Bellman
3. Improved exploration from uncertainty in distribution
4. Better credit assignment with full distribution

Empirical Benefits (to verify)

• Higher final rewards
• More stable training
• Better worst-case performance
• Faster convergence

---

📊 Evaluation

Training Metrics

• Episode Reward: Total fleet health improvement
• Episode Cost: Total maintenance spending (30 years)
• Cross-Entropy Loss: Distributional learning progress

Visualization Tools

# Training curves
python visualize_markov_v08.py <checkpoint>

# C51 distributions (per-action return distributions)
python visualize_markov_v08.py <checkpoint> --plot-dist

# Detailed distribution analysis
python analyze_c51_distribution.py <checkpoint> --save-dir <output_dir>

Distribution Plot: Shows probability distribution over return values for each action at a sample state.

C51 Distribution Analysis Results (200-Bridge Fleet)

Distribution Statistics:

Risk Profile (VaR/CVaR):

Mean vs Uncertainty:

State-Dependent Distributions:

Uncertainty Analysis:

---

🧪 Testing & Debugging

Run Tests

python test_noisy_net.py

Tests:

1. ✅ NoisyLinear layer functionality
2. ✅ FleetC51 forward pass
3. ✅ Distribution validity (sum to 1)
4. ✅ Q-value calculation from distribution
5. ✅ Training loop with C51 loss

Common Issues

1. Distribution doesn't sum to 1

• Issue: Numerical instability in softmax
• Fix: Added torch.softmax() normalization

2. Projection produces NaN

• Issue: Invalid support range
• Fix: Clamp t_z to [v_min, v_max]

3. Loss explodes

• Issue: Log of zero probability
• Fix: Add epsilon log(p + 1e-8)

---

📚 References

C51 Paper

"A Distributional Perspective on Reinforcement Learning"

• Authors: Bellemare, Dabney, Munos
• Conference: PMLR 2017
• Key Idea: Learn return distribution instead of expected value

Previous Methods (v0.7)

• Noisy Networks: Fortunato et al., ICLR 2018
• Dueling DQN: Wang et al., ICML 2016
• Double DQN: van Hasselt et al., AAAI 2016
• Prioritized Replay: Schaul et al., ICLR 2016

---

🎯 Command Line Arguments

python train_markov_fleet.py \
    --episodes 1000 \           # Number of episodes
    --n-envs 4 \                # Parallel environments
    --n-atoms 51 \              # C51 atoms
    --v-min -100.0 \            # Min return value
    --v-max 100.0 \             # Max return value
    --lr 0.0015 \               # Learning rate
    --batch-size 64 \           # Batch size
    --buffer-size 10000 \       # Replay buffer
    --target-sync 500 \         # Target network sync
    --device cuda               # cuda or cpu

---

🔄 Version History

v0.8 (Current) - C51 Distributional RL

• ✅ C51 categorical distribution
• ✅ Distributional Bellman update with projection
• ✅ Cross-entropy loss
• ✅ Distribution visualization

v0.7 - Noisy Networks

• ✅ Noisy Networks for exploration
• ✅ Removed ε-greedy
• ✅ Dueling DQN + Double DQN
• ✅ Prioritized N-step Replay

v0.6 - Baseline

• ✅ Standard DQN with ε-greedy
• ✅ Basic fleet environment

---

📧 Contact & Support

For questions or issues with C51 implementation:

• Check test_noisy_net.py for debugging
• Review visualize_markov_v08.py --plot-dist for distribution visualization
• Verify config.yaml C51 parameters

---

✨ Next Steps

Potential Extensions

1. QR-DQN: Quantile regression for distribution learning
2. IQN: Implicit quantile networks
3. Rainbow: Combine all DQN improvements
4. PPO/A3C: Actor-Critic methods

Hyperparameter Tuning

• Try different n_atoms (21, 51, 101)
• Adjust v_min and v_max based on actual return range
• Experiment with learning rate schedules

---

Happy training with C51 Distributional RL! 🚀