Training Acceleration Strategies for Bridge Fleet DQN

Current Performance: ~1.9 seconds/episode (2000 episodes → 60 minutes)
Target: 6-10x speedup for large-scale experiments

Current Implementation Status

✅ Already Implemented

GPU Acceleration: CUDA support on NVIDIA RTX 4060 Ti (16GB)
Experience Replay: 200k capacity buffer
Batch Learning: 512 samples per update
Target Network: 1000-step synchronization
Gradient Clipping: norm = 10.0

Performance Baseline

Hardware: NVIDIA GeForce RTX 4060 Ti (16GB)
Python: 3.12.10
PyTorch: 2.6.0+cu124
CUDA: 12.4

Episodes    Time      Speed
100         ~2 min    1.9s/ep
2000        ~60 min   1.8s/ep

Acceleration Strategies

Phase 1: Low-Cost, High-Impact (Immediate Implementation)

1. torch.compile() - Model Compilation

Expected Speedup: 1.3-2x
Implementation Effort: Minimal (1 line per model)
Compatibility: PyTorch 2.0+

# Before
policy_net = UrbanAgentDQN(state_dim, n_bridges, n_actions).to(device)
target_net = UrbanAgentDQN(state_dim, n_bridges, n_actions).to(device)

# After
policy_net = torch.compile(
    UrbanAgentDQN(state_dim, n_bridges, n_actions).to(device),
    mode="max-autotune"  # or "reduce-overhead" for faster compilation
)
target_net = torch.compile(
    UrbanAgentDQN(state_dim, n_bridges, n_actions).to(device),
    mode="max-autotune"
)

Pros:

Graph-level optimizations (kernel fusion, memory layout)
Minimal code changes
Automatic optimization

Cons:

Initial compilation overhead (~30 seconds)
Debugging becomes harder
May not work with dynamic control flow

Impact on Current System:

2000 episodes: 60 min → 30-45 min

2. Mixed Precision Training (AMP)

Expected Speedup: 1.5-2x
Implementation Effort: Low (5-10 lines)
Compatibility: CUDA Compute Capability 7.0+ (RTX 4060 Ti ✓)

from torch.cuda.amp import autocast, GradScaler

# Initialize scaler
scaler = GradScaler()

# In training loop
for batch in replay_buffer.sample(batch_size):
    optimizer.zero_grad()
    
    with autocast():
        q_values = policy_net(states)
        q_selected = q_values.gather(1, actions)
        
        with torch.no_grad():
            next_q = target_net(next_states).max(1)[0]
            q_target = rewards + gamma * next_q * (1 - dones)
        
        loss = F.mse_loss(q_selected.squeeze(), q_target)
    
    # Scaled backpropagation
    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(policy_net.parameters(), gradient_clip)
    scaler.step(optimizer)
    scaler.update()

Pros:

Leverages Tensor Cores (FP16 computation)
Reduces memory usage by ~50%
Allows larger batch sizes

Cons:

Requires careful handling of gradient scaling
Potential numerical instability (rare with modern PyTorch)

Impact on Current System:

2000 episodes: 60 min → 30-40 min
Can increase batch size: 512 → 768 or 1024

3. Double DQN

Expected Speedup: None (but improves learning efficiency)
Implementation Effort: Minimal (1 line change)
Benefits: Reduces overestimation, faster convergence

# Before (Vanilla DQN)
with torch.no_grad():
    next_q_values = target_net(next_states).max(1)[0]
    q_target = rewards + gamma * next_q_values * (1 - dones)

# After (Double DQN)
with torch.no_grad():
    next_actions = policy_net(next_states).argmax(1)
    next_q_values = target_net(next_states).gather(1, next_actions.unsqueeze(1)).squeeze()
    q_target = rewards + gamma * next_q_values * (1 - dones)

Pros:

Prevents Q-value overestimation
More stable training
Zero computational overhead

Cons:

None

Impact on Current System:

May reduce required episodes by 10-20%

Phase 2: Moderate Refactoring (Learning Efficiency)

4. Dueling DQN Architecture

Expected Speedup: None (but 20-30% fewer episodes needed)
Implementation Effort: Moderate (network redesign)
Computation Overhead: +10-15%

class DuelingUrbanDQN(nn.Module):
    def __init__(self, state_dim, n_bridges, n_actions):
        super().__init__()
        
        # Shared feature extraction
        self.feature = nn.Sequential(
            nn.Linear(state_dim, 512),
            nn.ReLU(),
            nn.Linear(512, 1024),
            nn.ReLU(),
            nn.Linear(1024, 512),
            nn.ReLU()
        )
        
        # Value stream V(s)
        self.value_stream = nn.Sequential(
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Linear(256, 1)
        )
        
        # Advantage stream A(s,a)
        self.advantage_stream = nn.Sequential(
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Linear(256, n_bridges * n_actions)
        )
        
        self.n_bridges = n_bridges
        self.n_actions = n_actions
    
    def forward(self, x):
        features = self.feature(x)
        value = self.value_stream(features)
        advantage = self.advantage_stream(features)
        advantage = advantage.view(-1, self.n_bridges, self.n_actions)
        
        # Q(s,a) = V(s) + (A(s,a) - mean(A(s,a)))
        q_values = value.unsqueeze(1).unsqueeze(2) + \
                   (advantage - advantage.mean(2, keepdim=True))
        
        return q_values.view(-1, self.n_bridges * self.n_actions)

Pros:

Separates state value from action advantage
Better learning in high-dimensional action spaces
Especially useful for Urban agent (20×5=100 actions)

Cons:

Slightly slower per iteration (+10-15%)
More complex architecture

Impact on Current System:

Episodes needed: 2000 → 1400-1600
Total time similar (fewer episodes, slightly slower per episode)

5. Multi-Step Learning (n-step Returns)

Expected Speedup: 1.5-2x reduction in episodes
Implementation Effort: Moderate (buffer redesign)
Computation Overhead: +15-20%

class NStepReplayBuffer:
    def __init__(self, capacity, n_steps=3, gamma=0.95):
        self.capacity = capacity
        self.n_steps = n_steps
        self.gamma = gamma
        self.buffer = deque(maxlen=capacity)
        self.n_step_buffer = deque(maxlen=n_steps)
    
    def push(self, state, action, reward, next_state, done):
        self.n_step_buffer.append((state, action, reward, next_state, done))
        
        if len(self.n_step_buffer) < self.n_steps:
            return
        
        # Compute n-step return
        R = 0
        for i in range(self.n_steps):
            s, a, r, ns, d = self.n_step_buffer[i]
            R += (self.gamma ** i) * r
            if d:
                break
        
        # Store n-step transition
        s0, a0 = self.n_step_buffer[0][:2]
        sn, dn = self.n_step_buffer[-1][3:]
        self.buffer.append((s0, a0, R, sn, dn))

Pros:

Faster credit assignment (rewards propagate faster)
Reduces bias (less bootstrapping)
Core component of Rainbow DQN

Cons:

More complex buffer implementation
+15-20% computation per episode

Impact on Current System:

Episodes needed: 2000 → 1200-1400
Time per episode: +15-20%
Total time: 60 min → 35-45 min

6. Prioritized Experience Replay (PER)

Expected Speedup: 2-3x reduction in episodes
Implementation Effort: High (complex data structure)
Computation Overhead: +20-30%

class PrioritizedReplayBuffer:
    def __init__(self, capacity, alpha=0.6, beta=0.4):
        self.capacity = capacity
        self.alpha = alpha  # Prioritization strength
        self.beta = beta    # Importance sampling correction
        self.buffer = []
        self.priorities = np.zeros(capacity, dtype=np.float32)
        self.pos = 0
    
    def push(self, transition):
        max_priority = self.priorities.max() if self.buffer else 1.0
        
        if len(self.buffer) < self.capacity:
            self.buffer.append(transition)
        else:
            self.buffer[self.pos] = transition
        
        self.priorities[self.pos] = max_priority
        self.pos = (self.pos + 1) % self.capacity
    
    def sample(self, batch_size, beta=None):
        if beta is None:
            beta = self.beta
        
        priorities = self.priorities[:len(self.buffer)]
        probs = priorities ** self.alpha
        probs /= probs.sum()
        
        indices = np.random.choice(len(self.buffer), batch_size, p=probs)
        samples = [self.buffer[idx] for idx in indices]
        
        # Importance sampling weights
        total = len(self.buffer)
        weights = (total * probs[indices]) ** (-beta)
        weights /= weights.max()
        
        return samples, indices, weights
    
    def update_priorities(self, indices, td_errors):
        for idx, error in zip(indices, td_errors):
            self.priorities[idx] = abs(error) + 1e-6

Usage in training:

# Sample with priorities
batch, indices, weights = buffer.sample(batch_size, beta=current_beta)

# Compute TD errors
td_errors = (q_values - q_targets).detach().cpu().numpy()

# Update priorities
buffer.update_priorities(indices, td_errors)

# Weight loss by importance sampling
loss = (F.mse_loss(q_values, q_targets, reduction='none') * 
        torch.FloatTensor(weights).to(device)).mean()

Pros:

2-3x faster convergence (learns from important transitions)
More sample-efficient
Core component of Rainbow DQN

Cons:

+20-30% overhead per episode
Complex implementation (sum-tree for efficiency)
Requires careful tuning (α, β)

Impact on Current System:

Episodes needed: 2000 → 800-1000
Time per episode: +20-30%
Total time: 60 min → 25-35 min

Phase 3: Major Refactoring (Parallelization)

7. Vectorized Environments (Async/Sync)

Expected Speedup: 3-4x for data collection
Implementation Effort: High (environment redesign)
Memory: 4-8x increase

import gymnasium as gym
from gymnasium.vector import AsyncVectorEnv, SyncVectorEnv

# Create 4 parallel environments
def make_env():
    return FleetEnvironmentGym(config)

# Asynchronous (faster, uses multiprocessing)
envs = AsyncVectorEnv([make_env for _ in range(4)])

# Synchronous (simpler, uses threading)
envs = SyncVectorEnv([make_env for _ in range(4)])

# Collect 4x experience simultaneously
observations, rewards, dones, infos = envs.step(actions)

Training loop adaptation:

def train_vectorized(envs, agent, n_episodes):
    n_envs = envs.num_envs
    observations = envs.reset()
    episode_rewards = np.zeros(n_envs)
    
    while episodes_completed < n_episodes:
        # Select actions for all environments
        actions = agent.select_actions(observations)
        
        # Step all environments
        next_obs, rewards, dones, infos = envs.step(actions)
        
        # Store transitions
        for i in range(n_envs):
            buffer.push(observations[i], actions[i], rewards[i], 
                       next_obs[i], dones[i])
        
        # Learn when buffer is ready
        if len(buffer) >= batch_size:
            agent.learn(buffer.sample(batch_size))
        
        observations = next_obs

Pros:

3-4x data collection speed (utilize GPU parallelism)
More diverse experience
Better utilization of GPU (100 bridges × 4 envs = 400 bridges)

Cons:

4x memory usage
Complex environment management
Requires thread-safe/process-safe implementation

Impact on Current System:

Data collection: 4x faster
GPU utilization: 30% → 70-80%
2000 episodes: 60 min → 15-20 min

8. Gymnasium API Migration

Expected Speedup: ~5-10% (minor)
Implementation Effort: Moderate
Benefits: Ecosystem integration, maintainability

import gymnasium as gym
from gymnasium import spaces

class FleetEnvironmentGym(gym.Env):
    metadata = {'render_modes': ['human', 'rgb_array']}
    
    def __init__(self, config):
        super().__init__()
        
        # Define action space
        # Urban: 20 bridges × 5 actions
        # Rural: 8 strategies
        self.action_space = spaces.Dict({
            'urban': spaces.MultiDiscrete([5] * 20),
            'rural': spaces.Discrete(8)
        })
        
        # Define observation space
        # Urban: 81D (20 bridges × 4 + 1 budget)
        # Rural: 10D (9 stats + 1 budget)
        self.observation_space = spaces.Dict({
            'urban': spaces.Box(low=0, high=100, shape=(81,), dtype=np.float32),
            'rural': spaces.Box(low=0, high=100, shape=(10,), dtype=np.float32)
        })
    
    def reset(self, seed=None, options=None):
        super().reset(seed=seed)
        # ... reset logic ...
        return observation, info
    
    def step(self, action):
        # ... step logic ...
        return observation, reward, terminated, truncated, info
    
    def render(self):
        # Optional: visualization
        pass

Pros:

Standard API for RL community
Compatible with stable-baselines3, RLlib, etc.
Better documentation and examples

Cons:

Refactoring effort
Dict action/observation spaces add complexity
Minimal speed improvement

Impact on Current System:

Direct speed: +5-10%
Main benefit: compatibility and maintainability

Phase 4: Advanced Techniques (Optional)

9. Noisy Networks for Exploration

Expected Speedup: 10-20% fewer episodes
Implementation Effort: Moderate
Replaces: ε-greedy exploration

class NoisyLinear(nn.Module):
    def __init__(self, in_features, out_features, sigma_init=0.5):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.sigma_init = sigma_init
        
        # Learnable parameters
        self.weight_mu = nn.Parameter(torch.Tensor(out_features, in_features))
        self.weight_sigma = nn.Parameter(torch.Tensor(out_features, in_features))
        self.bias_mu = nn.Parameter(torch.Tensor(out_features))
        self.bias_sigma = nn.Parameter(torch.Tensor(out_features))
        
        # Factorized noise
        self.register_buffer('weight_epsilon', torch.Tensor(out_features, in_features))
        self.register_buffer('bias_epsilon', torch.Tensor(out_features))
        
        self.reset_parameters()
        self.reset_noise()
    
    def forward(self, x):
        if self.training:
            weight = self.weight_mu + self.weight_sigma * self.weight_epsilon
            bias = self.bias_mu + self.bias_sigma * self.bias_epsilon
        else:
            weight = self.weight_mu
            bias = self.bias_mu
        
        return F.linear(x, weight, bias)
    
    def reset_noise(self):
        epsilon_in = self._scale_noise(self.in_features)
        epsilon_out = self._scale_noise(self.out_features)
        self.weight_epsilon.copy_(epsilon_out.outer(epsilon_in))
        self.bias_epsilon.copy_(epsilon_out)
    
    def _scale_noise(self, size):
        x = torch.randn(size)
        return x.sign() * x.abs().sqrt()

Pros:

State-dependent exploration (more efficient)
No ε-decay schedule needed
Rainbow DQN component

Cons:

+10-15% computation overhead
More parameters to learn

Impact on Current System:

Episodes needed: 2000 → 1600-1800
Overhead: +10-15%
Net benefit: marginal

10. PTAN Integration

Expected Speedup: ~5-10%
Implementation Effort: Moderate
Benefits: Code simplification, best practices

from ptan import experience, agent, common

# Experience source
exp_source = experience.ExperienceSourceFirstLast(
    env, agent, gamma=0.95, steps_count=1
)

# Or n-step experience
exp_source = experience.ExperienceSourceFirstLast(
    env, agent, gamma=0.95, steps_count=3
)

# Agent wrapper
dqn_agent = agent.DQNAgent(dqn_model, action_selector, device="cuda")

# Simple training loop
for step, exp in enumerate(exp_source):
    buffer.append(exp)
    
    if len(buffer) < batch_size:
        continue
    
    batch = buffer.sample(batch_size)
    loss = calc_loss(batch, dqn_model, target_model)
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

Pros:

Cleaner, more maintainable code
Tested implementations
Good for rapid prototyping

Cons:

Additional dependency
Less flexibility for custom logic
Multi-agent support is limited

Impact on Current System:

Speed: +5-10%
Main benefit: code quality, not performance

Recommended Implementation Roadmap

Quick Wins (1-2 hours)

# Step 1: Enable model compilation
policy_net = torch.compile(policy_net, mode="max-autotune")
target_net = torch.compile(target_net, mode="max-autotune")

# Step 2: Enable mixed precision
from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()

# Step 3: Implement Double DQN
next_actions = policy_net(next_states).argmax(1)
next_q = target_net(next_states).gather(1, next_actions.unsqueeze(1))

Expected Result:

2000 episodes: 60 min → 20-30 min
Code changes: ~20 lines

Medium-Term (1-2 days)

# Step 4: Implement Dueling DQN architecture
# Step 5: Add 3-step returns
# Step 6: Implement Prioritized Experience Replay

Expected Result:

Episodes needed: 2000 → 800-1000
Total time: 60 min → 20-25 min
Code changes: ~200 lines

Long-Term (1-2 weeks)

# Step 7: Vectorized environments (4-8 parallel)
# Step 8: Gymnasium API migration
# Step 9: Noisy networks

Expected Result:

2000 episodes: 60 min → 6-10 min
Code changes: ~500 lines + refactoring

Expected Final Performance

Conservative Estimate (Phase 1 + Phase 2)

Before: 2000 episodes = 60 minutes
After:  1000 episodes = 20 minutes

Speedup: 6x (3x fewer episodes + 2x faster per episode)

Aggressive Estimate (All Phases)

Before: 2000 episodes = 60 minutes
After:  800 episodes = 8 minutes

Speedup: 10x (2.5x fewer episodes + 4x faster execution)

Benchmark Targets

Configuration	Episodes	Time	vs Baseline
Baseline	2000	60 min	1x
+ torch.compile + AMP	2000	30 min	2x
+ Double DQN	1800	27 min	2.2x
+ Multi-step (3)	1200	24 min	2.5x
+ PER	800	20 min	3x
+ Dueling DQN	700	18 min	3.3x
+ Vectorized (4 envs)	700	8 min	7.5x
+ All optimizations	600	6 min	10x

References

Rainbow DQN: Hessel et al. (2018). "Rainbow: Combining Improvements in Deep Reinforcement Learning"
Double DQN: van Hasselt et al. (2015). "Deep Reinforcement Learning with Double Q-learning"
Dueling DQN: Wang et al. (2016). "Dueling Network Architectures for Deep Reinforcement Learning"
PER: Schaul et al. (2015). "Prioritized Experience Replay"
Noisy Networks: Fortunato et al. (2017). "Noisy Networks for Exploration"
PyTorch AMP: https://pytorch.org/docs/stable/amp.html
torch.compile: https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html
PTAN: https://github.com/Shmuma/ptan
Gymnasium: https://gymnasium.farama.org/

Document Version: 1.0
Last Updated: 2025-12-06
Current Status: Planning Phase
Next Action: Implement Phase 1 (Quick Wins)

FilesExpand file tree

Faster_acceleration.md

Latest commit

History

Faster_acceleration.md

File metadata and controls

Training Acceleration Strategies for Bridge Fleet DQN

Current Implementation Status

✅ Already Implemented

Performance Baseline

Acceleration Strategies

Phase 1: Low-Cost, High-Impact (Immediate Implementation)

1. torch.compile() - Model Compilation

2. Mixed Precision Training (AMP)

3. Double DQN

Phase 2: Moderate Refactoring (Learning Efficiency)

4. Dueling DQN Architecture

5. Multi-Step Learning (n-step Returns)

6. Prioritized Experience Replay (PER)

Phase 3: Major Refactoring (Parallelization)

7. Vectorized Environments (Async/Sync)

8. Gymnasium API Migration

Phase 4: Advanced Techniques (Optional)

9. Noisy Networks for Exploration

10. PTAN Integration

Recommended Implementation Roadmap

Quick Wins (1-2 hours)

Medium-Term (1-2 days)

Long-Term (1-2 weeks)

Expected Final Performance

Conservative Estimate (Phase 1 + Phase 2)

Aggressive Estimate (All Phases)

Benchmark Targets

References